fix: properly initialize custom agents

packages/ai-core/src/browser/ai-configuration/template-settings-renderer.tsx

@@ -25,7 +25,7 @@ export interface TemplateSettingProps {
 
 export const TemplateRenderer: React.FC<TemplateSettingProps> = ({ agentId, template, promptCustomizationService }) => {
     const openTemplate = React.useCallback(async () => {
-        promptCustomizationService.editTemplate(template.id);
+        promptCustomizationService.editTemplate(template.id, template.template);
     }, [template, promptCustomizationService]);
     const resetTemplate = React.useCallback(async () => {
         promptCustomizationService.resetTemplate(template.id);

packages/ai-core/src/browser/ai-core-frontend-application-contribution.ts

@@ -15,23 +15,22 @@
 // *****************************************************************************
 
 import { FrontendApplicationContribution } from '@theia/core/lib/browser';
-import { inject, injectable } from '@theia/core/shared/inversify';
-import { PromptService } from '../common';
+import { inject, injectable, named } from '@theia/core/shared/inversify';
+import { Agent } from '../common';
 import { AgentService } from '../common/agent-service';
+import { ContributionProvider } from '@theia/core/lib/common/contribution-provider';
 
 @injectable()
 export class AICoreFrontendApplicationContribution implements FrontendApplicationContribution {
     @inject(AgentService)
     private readonly agentService: AgentService;
 
-    @inject(PromptService)
-    private readonly promptService: PromptService;
+    @inject(ContributionProvider) @named(Agent)
+    protected readonly agentsProvider: ContributionProvider<Agent>;
 
     onStart(): void {
-        this.agentService.getAllAgents().forEach(a => {
-            a.promptTemplates.forEach(t => {
-                this.promptService.storePrompt(t.id, t.template);
-            });
+        this.agentsProvider.getContributions().forEach(agent => {
+            this.agentService.registerAgent(agent);
         });
     }
 

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

@@ -17,12 +17,11 @@
 import { DisposableCollection, URI, Event, Emitter } from '@theia/core';
 import { OpenerService } from '@theia/core/lib/browser';
 import { inject, injectable, postConstruct } from '@theia/core/shared/inversify';
-import { PromptCustomizationService, PromptTemplate, CustomAgentDescription } from '../common';
+import { PromptCustomizationService, CustomAgentDescription } from '../common';
 import { BinaryBuffer } from '@theia/core/lib/common/buffer';
 import { FileService } from '@theia/filesystem/lib/browser/file-service';
 import { FileChangesEvent } from '@theia/filesystem/lib/common/files';
 import { AICorePreferences, PREFERENCE_NAME_PROMPT_TEMPLATES } from './ai-core-preferences';
-import { AgentService } from '../common/agent-service';
 import { EnvVariablesServer } from '@theia/core/lib/common/env-variables';
 import { load, dump } from 'js-yaml';
 
@@ -49,9 +48,6 @@ export class FrontendPromptCustomizationServiceImpl implements PromptCustomizati
     @inject(OpenerService)
     protected readonly openerService: OpenerService;
 
-    @inject(AgentService)
-    protected readonly agentService: AgentService;
-
     protected readonly trackedTemplateURIs = new Set<string>();
     protected templates = new Map<string, string>();
 
@@ -168,17 +164,10 @@ export class FrontendPromptCustomizationServiceImpl implements PromptCustomizati
         return this.templates.get(id);
     }
 
-    async editTemplate(id: string, content?: string): Promise<void> {
-        const template = this.getOriginalTemplate(id);
-        if (template === undefined) {
-            throw new Error(`Unable to edit template ${id}: template not found.`);
-        }
+    async editTemplate(id: string, defaultContent?: string): Promise<void> {
         const editorUri = await this.getTemplateURI(id);
         if (! await this.fileService.exists(editorUri)) {
-            await this.fileService.createFile(editorUri, BinaryBuffer.fromString(content ?? template.template));
-        } else if (content) {
-            // Write content to the file before opening it
-            await this.fileService.writeFile(editorUri, BinaryBuffer.fromString(content));
+            await this.fileService.createFile(editorUri, BinaryBuffer.fromString(defaultContent ?? ''));
         }
         const openHandler = await this.openerService.getOpener(editorUri);
         openHandler.open(editorUri);
@@ -191,17 +180,6 @@ export class FrontendPromptCustomizationServiceImpl implements PromptCustomizati
         }
     }
 
-    getOriginalTemplate(id: string): PromptTemplate | undefined {
-        for (const agent of this.agentService.getAllAgents()) {
-            for (const template of agent.promptTemplates) {
-                if (template.id === id) {
-                    return template;
-                }
-            }
-        }
-        return undefined;
-    }
-
     getTemplateIDFromURI(uri: URI): string | undefined {
         const id = this.removePromptTemplateSuffix(uri.path.name);
         if (this.templates.has(id)) {

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

@@ -13,10 +13,11 @@
 //
 // SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-only WITH Classpath-exception-2.0
 // *****************************************************************************
-import { inject, injectable, named, optional, postConstruct } from '@theia/core/shared/inversify';
-import { ContributionProvider, Emitter, Event } from '@theia/core';
+import { inject, injectable, optional, postConstruct } from '@theia/core/shared/inversify';
+import { Emitter, Event } from '@theia/core';
 import { Agent } from './agent';
 import { AISettingsService } from './settings-service';
+import { PromptService } from './prompt-service';
 
 export const AgentService = Symbol('AgentService');
 
@@ -71,12 +72,12 @@ export interface AgentService {
 @injectable()
 export class AgentServiceImpl implements AgentService {
 
-    @inject(ContributionProvider) @named(Agent)
-    protected readonly agentsProvider: ContributionProvider<Agent>;
-
     @inject(AISettingsService) @optional()
     protected readonly aiSettingsService: AISettingsService | undefined;
 
+    @inject(PromptService)
+    protected readonly promptService: PromptService;
+
     protected disabledAgents = new Set<string>();
 
     protected _agents: Agent[] = [];
@@ -95,28 +96,29 @@ export class AgentServiceImpl implements AgentService {
         });
     }
 
-    private get agents(): Agent[] {
-        // We can't collect the contributions at @postConstruct because this will lead to a circular dependency
-        // with agents reusing the chat agent service (e.g. orchestrator) which in turn injects the agent service
-        return [...this.agentsProvider.getContributions(), ...this._agents];
-    }
-
     registerAgent(agent: Agent): void {
         this._agents.push(agent);
+        agent.promptTemplates.forEach(
+            template => this.promptService.storePrompt(template.id, template.template)
+        );
         this.onDidChangeAgentsEmitter.fire();
     }
 
     unregisterAgent(agentId: string): void {
+        const agent = this._agents.find(a => a.id === agentId);
         this._agents = this._agents.filter(a => a.id !== agentId);
         this.onDidChangeAgentsEmitter.fire();
+        agent?.promptTemplates.forEach(
+            template => this.promptService.removePrompt(template.id)
+        );
     }
 
     getAgents(): Agent[] {
-        return this.agents.filter(agent => this.isEnabled(agent.id));
+        return this._agents.filter(agent => this.isEnabled(agent.id));
     }
 
     getAllAgents(): Agent[] {
-        return this.agents;
+        return this._agents;
     }
 
     enableAgent(agentId: string): void {

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

@@ -58,11 +58,16 @@ export interface PromptService {
      */
     getPrompt(id: string, args?: { [key: string]: unknown }): Promise<ResolvedPromptTemplate | undefined>;
     /**
-     * Manually add a prompt to the list of prompts.
+     * Adds a prompt to the list of prompts.
      * @param id the id of the prompt
      * @param prompt the prompt template to store
      */
     storePrompt(id: string, prompt: string): void;
+    /**
+     * Removes a prompt from the list of prompts.
+     * @param id the id of the prompt
+     */
+    removePrompt(id: string): void;
     /**
      * Return all known prompts as a {@link PromptMap map}.
      */
@@ -113,9 +118,9 @@ export interface PromptCustomizationService {
      * on the implementation. Implementation may for example decide to
      * open an editor, or request more information from the user, ...
      * @param id the template id.
-     * @param content optional content to customize the template.
+     * @param content optional default content to initialize the template
      */
-    editTemplate(id: string, content?: string): void;
+    editTemplate(id: string, defaultContent?: string): void;
 
     /**
      * Reset the template to its default value.
@@ -250,4 +255,7 @@ export class PromptServiceImpl implements PromptService {
     storePrompt(id: string, prompt: string): void {
         this._prompts[id] = { id, template: prompt };
     }
+    removePrompt(id: string): void {
+        delete this._prompts[id];
+    }
 }