Allow workspace specific custom agents (#15457)

fixed #15456

Author: JonasHelming

packages/ai-core/src/browser/frontend-prompt-customization-service.ts

@@ -519,46 +519,93 @@ export class FrontendPromptCustomizationServiceImpl implements PromptCustomizati
     }
 
     async getCustomAgents(): Promise<CustomAgentDescription[]> {
-        const customAgentYamlUri = (await this.getTemplatesDirectoryURI()).resolve('customAgents.yml');
+        const agentsById = new Map<string, CustomAgentDescription>();
+        // 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<string, CustomAgentDescription>
+    ): Promise<void> {
+        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<string>();
-            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<void> {
-        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<void> {
         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);
     }
 }

packages/ai-core/src/common/prompt-service.ts

@@ -202,9 +202,19 @@ export interface PromptCustomizationService {
     readonly onDidChangeCustomAgents: Event<void>;
 
     /**
-     * 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<void>;
 }
 
 @injectable()

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<void> {
+        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 {