diff --git a/packages/ai-workspace-agent/src/browser/frontend-module.ts b/packages/ai-workspace-agent/src/browser/frontend-module.ts index b4adcd4da9d79..ef3f02174dc7a 100644 --- a/packages/ai-workspace-agent/src/browser/frontend-module.ts +++ b/packages/ai-workspace-agent/src/browser/frontend-module.ts @@ -17,7 +17,7 @@ import { ContainerModule } from '@theia/core/shared/inversify'; import { ChatAgent } from '@theia/ai-chat/lib/common'; import { Agent, ToolProvider } from '@theia/ai-core/lib/common'; import { WorkspaceAgent } from './workspace-agent'; -import { FileContentFunction, GetWorkspaceDirectoryStructure, GetWorkspaceFileList } from './functions'; +import { FileContentFunction, GetWorkspaceDirectoryStructure, GetWorkspaceFileList, WorkspaceUtils } from './functions'; export default new ContainerModule(bind => { bind(WorkspaceAgent).toSelf().inSingletonScope(); @@ -26,4 +26,5 @@ export default new ContainerModule(bind => { bind(ToolProvider).to(GetWorkspaceFileList); bind(ToolProvider).to(FileContentFunction); bind(ToolProvider).to(GetWorkspaceDirectoryStructure); + bind(WorkspaceUtils).toSelf().inSingletonScope(); }); diff --git a/packages/ai-workspace-agent/src/browser/functions.ts b/packages/ai-workspace-agent/src/browser/functions.ts index 4e90830d34f41..0b3847cea23fe 100644 --- a/packages/ai-workspace-agent/src/browser/functions.ts +++ b/packages/ai-workspace-agent/src/browser/functions.ts @@ -21,9 +21,34 @@ import { FileStat } from '@theia/filesystem/lib/common/files'; import { WorkspaceService } from '@theia/workspace/lib/browser'; import { FILE_CONTENT_FUNCTION_ID, GET_WORKSPACE_DIRECTORY_STRUCTURE_FUNCTION_ID, GET_WORKSPACE_FILE_LIST_FUNCTION_ID } from '../common/functions'; -function shouldExclude(stat: FileStat): boolean { - const excludedFolders = ['node_modules', 'lib']; - return stat.resource.path.base.startsWith('.') || excludedFolders.includes(stat.resource.path.base); +@injectable() +export class WorkspaceUtils { + @inject(WorkspaceService) + protected workspaceService: WorkspaceService; + + async getWorkspaceRoot(): Promise { + const wsRoots = await this.workspaceService.roots; + if (wsRoots.length === 0) { + throw new Error('No workspace has been opened yet'); + } + return wsRoots[0].resource; + } + + ensureWithinWorkspace(targetUri: URI, workspaceRootUri: URI): void { + if (!targetUri.toString().startsWith(workspaceRootUri.toString())) { + throw new Error('Access outside of the workspace is not allowed'); + } + } + /** + * Determines whether a given file or directory should be excluded from workspace operations. + * + * @param stat - The `FileStat` object representing the file or directory to check. + * @returns `true` if the file or directory should be excluded, `false` otherwise. + */ + shouldExclude(stat: FileStat): boolean { + const excludedFolders = ['node_modules', 'lib']; + return stat.resource.path.base.startsWith('.') || excludedFolders.includes(stat.resource.path.base); + } } @injectable() @@ -40,22 +65,21 @@ export class GetWorkspaceDirectoryStructure implements ToolProvider { }; } - @inject(WorkspaceService) - protected workspaceService: WorkspaceService; - @inject(FileService) protected readonly fileService: FileService; - private async getDirectoryStructure(): Promise { - const wsRoots = await this.workspaceService.roots; + @inject(WorkspaceUtils) + protected workspaceUtils: WorkspaceUtils; - if (wsRoots.length === 0) { - throw new Error('Workspace root not found'); + private async getDirectoryStructure(): Promise { + let workspaceRoot; + try { + workspaceRoot = await this.workspaceUtils.getWorkspaceRoot(); + } catch (error) { + return [`Error: ${error.message}`]; } - const workspaceRootUri = wsRoots[0].resource; - - return this.buildDirectoryStructure(workspaceRootUri); + return this.buildDirectoryStructure(workspaceRoot); } private async buildDirectoryStructure(uri: URI, prefix: string = ''): Promise { @@ -64,7 +88,7 @@ export class GetWorkspaceDirectoryStructure implements ToolProvider { if (stat && stat.isDirectory && stat.children) { for (const child of stat.children) { - if (!child.isDirectory || shouldExclude(child)) { continue; }; + if (!child.isDirectory || this.workspaceUtils.shouldExclude(child)) { continue; }; const path = `${prefix}${child.resource.path.base}/`; result.push(path); result.push(...await this.buildDirectoryStructure(child.resource, `${path}`)); @@ -102,31 +126,27 @@ export class FileContentFunction implements ToolProvider { }; } - @inject(WorkspaceService) - protected workspaceService: WorkspaceService; - @inject(FileService) protected readonly fileService: FileService; + @inject(WorkspaceUtils) + protected readonly workspaceUtils: WorkspaceUtils; + private parseArg(arg_string: string): string { const result = JSON.parse(arg_string); return result.file; } private async getFileContent(file: string): Promise { - const wsRoots = await this.workspaceService.roots; - - if (wsRoots.length === 0) { - throw new Error('Workspace root not found'); + let workspaceRoot; + try { + workspaceRoot = await this.workspaceUtils.getWorkspaceRoot(); + } catch (error) { + return JSON.stringify({ error: error.message }); } - const workspaceRootUri = wsRoots[0].resource; - - const targetUri = workspaceRootUri.resolve(file); - - if (!targetUri.toString().startsWith(workspaceRootUri.toString())) { - throw new Error('Access outside of the workspace is not allowed'); - } + const targetUri = workspaceRoot.resolve(file); + this.workspaceUtils.ensureWithinWorkspace(targetUri, workspaceRoot); try { const fileStat = await this.fileService.resolve(targetUri); @@ -171,32 +191,30 @@ export class GetWorkspaceFileList implements ToolProvider { }; } - @inject(WorkspaceService) - protected workspaceService: WorkspaceService; - @inject(FileService) protected readonly fileService: FileService; - async getProjectFileList(path?: string): Promise { - const wsRoots = await this.workspaceService.roots; + @inject(WorkspaceUtils) + protected workspaceUtils: WorkspaceUtils; - if (wsRoots.length === 0) { - throw new Error('Workspace root not found'); + async getProjectFileList(path?: string): Promise { + let workspaceRoot; + try { + workspaceRoot = await this.workspaceUtils.getWorkspaceRoot(); + } catch (error) { + return [`Error: ${error.message}`]; } - const workspaceRootUri = wsRoots[0].resource; - const targetUri = path ? workspaceRootUri.resolve(path) : workspaceRootUri; + const targetUri = path ? workspaceRoot.resolve(path) : workspaceRoot; + this.workspaceUtils.ensureWithinWorkspace(targetUri, workspaceRoot); - if (!targetUri.toString().startsWith(workspaceRootUri.toString())) { - throw new Error('Access outside of the workspace is not allowed'); - } try { const stat = await this.fileService.resolve(targetUri); if (!stat || !stat.isDirectory) { return ['Error: Directory not found']; } - return await this.listFilesDirectly(targetUri, workspaceRootUri); + return await this.listFilesDirectly(targetUri, workspaceRoot); } catch (error) { return ['Error: Directory not found']; @@ -208,7 +226,7 @@ export class GetWorkspaceFileList implements ToolProvider { const result: string[] = []; if (stat && stat.isDirectory) { - if (shouldExclude(stat)) { + if (this.workspaceUtils.shouldExclude(stat)) { return result; } const children = await this.fileService.resolve(uri); diff --git a/packages/ai-workspace-agent/src/common/template.ts b/packages/ai-workspace-agent/src/common/template.ts index 51897af1eca93..eb03aeea7eb90 100644 --- a/packages/ai-workspace-agent/src/common/template.ts +++ b/packages/ai-workspace-agent/src/common/template.ts @@ -30,8 +30,9 @@ Use the following functions to interact with the workspace files as needed: ### Workspace Navigation Guidelines -1. **Confirm Paths**: Always verify paths by listing directories or files as you navigate. Avoid assumptions based on user input alone. -2. **Start from Root**: Begin at the root and navigate subdirectories step-by-step. +1. **Start at the Root**: For general questions (e.g., "How to build the project"), check root-level documentation files or setup files before browsing subdirectories. +2. **Confirm Paths**: Always verify paths by listing directories or files as you navigate. Avoid assumptions based on user input alone. +3. **Navigate Step-by-Step**: Move into subdirectories only as needed, confirming each directory level. ### Response Guidelines