diff --git a/packages/ai-core/src/browser/frontend-prompt-customization-service.ts b/packages/ai-core/src/browser/frontend-prompt-customization-service.ts index 5731dee61062b..802dfd4718dcd 100644 --- a/packages/ai-core/src/browser/frontend-prompt-customization-service.ts +++ b/packages/ai-core/src/browser/frontend-prompt-customization-service.ts @@ -519,46 +519,93 @@ export class FrontendPromptCustomizationServiceImpl implements PromptCustomizati } async getCustomAgents(): Promise { - const customAgentYamlUri = (await this.getTemplatesDirectoryURI()).resolve('customAgents.yml'); + const agentsById = new Map(); + // First, process additional (workspace) template directories to give them precedence + for (const dirPath of this.additionalTemplateDirs) { + const dirURI = URI.fromFilePath(dirPath); + await this.loadCustomAgentsFromDirectory(dirURI, agentsById); + } + // Then process global template directory (only adding agents that don't conflict) + const globalTemplateDir = await this.getTemplatesDirectoryURI(); + await this.loadCustomAgentsFromDirectory(globalTemplateDir, agentsById); + // Return the merged list of agents + return Array.from(agentsById.values()); + } + + /** + * Load custom agents from a specific directory + * @param directoryURI The URI of the directory to load from + * @param agentsById Map to store the loaded agents by ID + */ + protected async loadCustomAgentsFromDirectory( + directoryURI: URI, + agentsById: Map + ): Promise { + const customAgentYamlUri = directoryURI.resolve('customAgents.yml'); const yamlExists = await this.fileService.exists(customAgentYamlUri); if (!yamlExists) { - return []; + return; } - const fileContent = await this.fileService.read(customAgentYamlUri, { encoding: 'utf-8' }); + try { + const fileContent = await this.fileService.read(customAgentYamlUri, { encoding: 'utf-8' }); const doc = load(fileContent.value); + if (!Array.isArray(doc) || !doc.every(entry => CustomAgentDescription.is(entry))) { - console.debug('Invalid customAgents.yml file content'); - return []; + console.debug(`Invalid customAgents.yml file content in ${directoryURI.toString()}`); + return; } + const readAgents = doc as CustomAgentDescription[]; - // make sure all agents are unique (id and name) - const uniqueAgentIds = new Set(); - const uniqueAgents: CustomAgentDescription[] = []; - readAgents.forEach(agent => { - if (uniqueAgentIds.has(agent.id)) { - return; + + // Add agents to the map if they don't already exist + for (const agent of readAgents) { + if (!agentsById.has(agent.id)) { + agentsById.set(agent.id, agent); } - uniqueAgentIds.add(agent.id); - uniqueAgents.push(agent); - }); - return uniqueAgents; + } } catch (e) { - console.debug(e.message, e); - return []; + console.debug(`Error loading customAgents.yml from ${directoryURI.toString()}: ${e.message}`, e); } } - async openCustomAgentYaml(): Promise { - const customAgentYamlUri = (await this.getTemplatesDirectoryURI()).resolve('customAgents.yml'); + /** + * Returns all locations of existing customAgents.yml files and potential locations where + * new customAgents.yml files could be created. + * + * @returns An array of objects containing the URI and whether the file exists + */ + async getCustomAgentsLocations(): Promise<{ uri: URI, exists: boolean }[]> { + const locations: { uri: URI, exists: boolean }[] = []; + // Check global template directory + const globalTemplateDir = await this.getTemplatesDirectoryURI(); + const globalAgentsUri = globalTemplateDir.resolve('customAgents.yml'); + const globalExists = await this.fileService.exists(globalAgentsUri); + locations.push({ uri: globalAgentsUri, exists: globalExists }); + // Check additional (workspace) template directories + for (const dirPath of this.additionalTemplateDirs) { + const dirURI = URI.fromFilePath(dirPath); + const agentsUri = dirURI.resolve('customAgents.yml'); + const exists = await this.fileService.exists(agentsUri); + locations.push({ uri: agentsUri, exists: exists }); + } + return locations; + } + + /** + * Opens an existing customAgents.yml file at the given URI, or creates a new one if it doesn't exist. + * + * @param uri The URI of the customAgents.yml file to open or create + */ + async openCustomAgentYaml(uri: URI): Promise { const content = dump([templateEntry]); - if (! await this.fileService.exists(customAgentYamlUri)) { - await this.fileService.createFile(customAgentYamlUri, BinaryBuffer.fromString(content)); + if (! await this.fileService.exists(uri)) { + await this.fileService.createFile(uri, BinaryBuffer.fromString(content)); } else { - const fileContent = (await this.fileService.readFile(customAgentYamlUri)).value; - await this.fileService.writeFile(customAgentYamlUri, BinaryBuffer.concat([fileContent, BinaryBuffer.fromString(content)])); + const fileContent = (await this.fileService.readFile(uri)).value; + await this.fileService.writeFile(uri, BinaryBuffer.concat([fileContent, BinaryBuffer.fromString(content)])); } - const openHandler = await this.openerService.getOpener(customAgentYamlUri); - openHandler.open(customAgentYamlUri); + const openHandler = await this.openerService.getOpener(uri); + openHandler.open(uri); } } diff --git a/packages/ai-core/src/common/prompt-service.ts b/packages/ai-core/src/common/prompt-service.ts index 9d9fc7f00cdb7..306bddf0d3330 100644 --- a/packages/ai-core/src/common/prompt-service.ts +++ b/packages/ai-core/src/common/prompt-service.ts @@ -202,9 +202,19 @@ export interface PromptCustomizationService { readonly onDidChangeCustomAgents: Event; /** - * Open the custom agent yaml file. + * Returns all locations of existing customAgents.yml files and potential locations where + * new customAgents.yml files could be created. + * + * @returns An array of objects containing the URI and whether the file exists + */ + getCustomAgentsLocations(): Promise<{ uri: URI, exists: boolean }[]>; + + /** + * Opens an existing customAgents.yml file at the given URI, or creates a new one if it doesn't exist. + * + * @param uri The URI of the customAgents.yml file to open or create */ - openCustomAgentYaml(): void; + openCustomAgentYaml(uri: URI): Promise; } @injectable() diff --git a/packages/ai-ide/src/browser/ai-configuration/agent-configuration-widget.tsx b/packages/ai-ide/src/browser/ai-configuration/agent-configuration-widget.tsx index cb636dc4fdca2..402ad6ebcc2d1 100644 --- a/packages/ai-ide/src/browser/ai-configuration/agent-configuration-widget.tsx +++ b/packages/ai-ide/src/browser/ai-configuration/agent-configuration-widget.tsx @@ -26,7 +26,8 @@ import { PromptCustomizationService, PromptService, } from '@theia/ai-core/lib/common'; -import { codicon, ReactWidget } from '@theia/core/lib/browser'; +import { codicon, QuickInputService, ReactWidget } from '@theia/core/lib/browser'; +import { URI } from '@theia/core/lib/common'; import { inject, injectable, postConstruct } from '@theia/core/shared/inversify'; import * as React from '@theia/core/shared/react'; import { AIConfigurationSelectionService } from './ai-configuration-service'; @@ -68,6 +69,9 @@ export class AIAgentConfigurationWidget extends ReactWidget { @inject(PromptService) protected promptService: PromptService; + @inject(QuickInputService) + protected readonly quickInputService: QuickInputService; + protected languageModels: LanguageModel[] | undefined; @postConstruct() @@ -237,8 +241,35 @@ export class AIAgentConfigurationWidget extends ReactWidget { this.aiConfigurationSelectionService.selectConfigurationTab(AIVariableConfigurationWidget.ID); } - protected addCustomAgent(): void { - this.promptCustomizationService.openCustomAgentYaml(); + protected async addCustomAgent(): Promise { + const locations = await this.promptCustomizationService.getCustomAgentsLocations(); + + // If only one location is available, use the direct approach + if (locations.length === 1) { + this.promptCustomizationService.openCustomAgentYaml(locations[0].uri); + return; + } + + // Multiple locations - show quick picker + const quickPick = this.quickInputService.createQuickPick(); + quickPick.title = 'Select Location for Custom Agents File'; + quickPick.placeholder = 'Choose where to create or open a custom agents file'; + + quickPick.items = locations.map(location => ({ + label: location.uri.path.toString(), + description: location.exists ? 'Open existing file' : 'Create new file', + location + })); + + quickPick.onDidAccept(async () => { + const selectedItem = quickPick.selectedItems[0] as unknown as { location: { uri: URI, exists: boolean } }; + if (selectedItem && selectedItem.location) { + quickPick.dispose(); + this.promptCustomizationService.openCustomAgentYaml(selectedItem.location.uri); + } + }); + + quickPick.show(); } protected setActiveAgent(agent: Agent): void {