Fix: TaskTool Dynamic Updates (QwenLM#697)

Author: tanzhenxin

packages/core/src/core/client.test.ts

@@ -239,6 +239,7 @@ describe('Gemini Client (client.ts)', () => {
     };
     const mockSubagentManager = {
       listSubagents: vi.fn().mockResolvedValue([]),
+      addChangeListener: vi.fn().mockReturnValue(() => {}),
     };
     mockConfigObject = {
       getContentGeneratorConfig: vi

packages/core/src/subagents/subagent-manager.ts

@@ -40,11 +40,29 @@ const AGENT_CONFIG_DIR = 'agents';
 export class SubagentManager {
   private readonly validator: SubagentValidator;
   private subagentsCache: Map<SubagentLevel, SubagentConfig[]> | null = null;
+  private readonly changeListeners: Set<() => void> = new Set();
 
   constructor(private readonly config: Config) {
     this.validator = new SubagentValidator();
   }
 
+  addChangeListener(listener: () => void): () => void {
+    this.changeListeners.add(listener);
+    return () => {
+      this.changeListeners.delete(listener);
+    };
+  }
+
+  private notifyChangeListeners(): void {
+    for (const listener of this.changeListeners) {
+      try {
+        listener();
+      } catch (error) {
+        console.warn('Subagent change listener threw an error:', error);
+      }
+    }
+  }
+
   /**
    * Creates a new subagent configuration.
    *
@@ -93,8 +111,8 @@ export class SubagentManager {
 
     try {
       await fs.writeFile(filePath, content, 'utf8');
-      // Clear cache after successful creation
-      this.clearCache();
+      // Refresh cache after successful creation
+      await this.refreshCache();
     } catch (error) {
       throw new SubagentError(
         `Failed to write subagent file: ${error instanceof Error ? error.message : 'Unknown error'}`,
@@ -183,8 +201,8 @@ export class SubagentManager {
 
     try {
       await fs.writeFile(existing.filePath, content, 'utf8');
-      // Clear cache after successful update
-      this.clearCache();
+      // Refresh cache after successful update
+      await this.refreshCache();
     } catch (error) {
       throw new SubagentError(
         `Failed to update subagent file: ${error instanceof Error ? error.message : 'Unknown error'}`,
@@ -242,8 +260,8 @@ export class SubagentManager {
       );
     }
 
-    // Clear cache after successful deletion
-    this.clearCache();
+    // Refresh cache after successful deletion
+    await this.refreshCache();
   }
 
   /**
@@ -327,21 +345,17 @@ export class SubagentManager {
    * @private
    */
   private async refreshCache(): Promise<void> {
-    this.subagentsCache = new Map();
+    const subagentsCache = new Map();
 
     const levels: SubagentLevel[] = ['project', 'user', 'builtin'];
 
     for (const level of levels) {
       const levelSubagents = await this.listSubagentsAtLevel(level);
-      this.subagentsCache.set(level, levelSubagents);
+      subagentsCache.set(level, levelSubagents);
     }
-  }
 
-  /**
-   * Clears the subagents cache, forcing the next listSubagents call to reload from disk.
-   */
-  clearCache(): void {
-    this.subagentsCache = null;
+    this.subagentsCache = subagentsCache;
+    this.notifyChangeListeners();
   }
 
   /**

packages/core/src/subagents/subagent.test.ts

@@ -41,12 +41,14 @@ import type {
   ToolConfig,
 } from './types.js';
 import { SubagentTerminateMode } from './types.js';
+import { GeminiClient } from '../core/client.js';
 
 vi.mock('../core/geminiChat.js');
 vi.mock('../core/contentGenerator.js');
 vi.mock('../utils/environmentContext.js');
 vi.mock('../core/nonInteractiveToolExecutor.js');
 vi.mock('../ide/ide-client.js');
+vi.mock('../core/client.js');
 
 async function createMockConfig(
   toolRegistryMocks = {},
@@ -194,6 +196,28 @@ describe('subagent.ts', () => {
           }) as unknown as GeminiChat,
       );
 
+      // Mock GeminiClient constructor to return a properly mocked client
+      const mockGeminiChat = {
+        setTools: vi.fn(),
+        getHistory: vi.fn().mockReturnValue([]),
+        setHistory: vi.fn(),
+        sendMessageStream: vi.fn(),
+      };
+
+      const mockGeminiClient = {
+        getChat: vi.fn().mockReturnValue(mockGeminiChat),
+        setTools: vi.fn().mockResolvedValue(undefined),
+        isInitialized: vi.fn().mockReturnValue(true),
+        getHistory: vi.fn().mockReturnValue([]),
+        initialize: vi.fn().mockResolvedValue(undefined),
+        setHistory: vi.fn(),
+      };
+
+      // Mock the GeminiClient constructor
+      vi.mocked(GeminiClient).mockImplementation(
+        () => mockGeminiClient as unknown as GeminiClient,
+      );
+
       // Default mock for executeToolCall
       vi.mocked(executeToolCall).mockResolvedValue({
         callId: 'default-call',

packages/core/src/tools/task.test.ts

@@ -43,6 +43,7 @@ describe('TaskTool', () => {
   let config: Config;
   let taskTool: TaskTool;
   let mockSubagentManager: SubagentManager;
+  let changeListeners: Array<() => void>;
 
   const mockSubagents: SubagentConfig[] = [
     {
@@ -70,13 +71,25 @@ describe('TaskTool', () => {
       getProjectRoot: vi.fn().mockReturnValue('/test/project'),
       getSessionId: vi.fn().mockReturnValue('test-session-id'),
       getSubagentManager: vi.fn(),
+      getGeminiClient: vi.fn().mockReturnValue(undefined),
     } as unknown as Config;
 
+    changeListeners = [];
+
     // Setup SubagentManager mock
     mockSubagentManager = {
       listSubagents: vi.fn().mockResolvedValue(mockSubagents),
       loadSubagent: vi.fn(),
       createSubagentScope: vi.fn(),
+      addChangeListener: vi.fn((listener: () => void) => {
+        changeListeners.push(listener);
+        return () => {
+          const index = changeListeners.indexOf(listener);
+          if (index >= 0) {
+            changeListeners.splice(index, 1);
+          }
+        };
+      }),
     } as unknown as SubagentManager;
 
     MockedSubagentManager.mockImplementation(() => mockSubagentManager);
@@ -106,6 +119,10 @@ describe('TaskTool', () => {
       expect(mockSubagentManager.listSubagents).toHaveBeenCalled();
     });
 
+    it('should subscribe to subagent manager changes', () => {
+      expect(mockSubagentManager.addChangeListener).toHaveBeenCalledTimes(1);
+    });
+
     it('should update description with available subagents', () => {
       expect(taskTool.description).toContain('file-search');
       expect(taskTool.description).toContain(
@@ -232,6 +249,31 @@ describe('TaskTool', () => {
   });
 
   describe('refreshSubagents', () => {
+    it('should refresh when change listener fires', async () => {
+      const newSubagents: SubagentConfig[] = [
+        {
+          name: 'new-agent',
+          description: 'A brand new agent',
+          systemPrompt: 'Do new things.',
+          level: 'project',
+          filePath: '/project/.qwen/agents/new-agent.md',
+        },
+      ];
+
+      vi.mocked(mockSubagentManager.listSubagents).mockResolvedValueOnce(
+        newSubagents,
+      );
+
+      const listener = changeListeners[0];
+      expect(listener).toBeDefined();
+
+      listener?.();
+      await vi.runAllTimersAsync();
+
+      expect(taskTool.description).toContain('new-agent');
+      expect(taskTool.description).toContain('A brand new agent');
+    });
+
     it('should refresh available subagents and update description', async () => {
       const newSubagents: SubagentConfig[] = [
         {

packages/core/src/tools/task.ts

@@ -86,23 +86,32 @@ export class TaskTool extends BaseDeclarativeTool<TaskParams, ToolResult> {
     );
 
     this.subagentManager = config.getSubagentManager();
+    this.subagentManager.addChangeListener(() => {
+      void this.refreshSubagents();
+    });
 
     // Initialize the tool asynchronously
-    this.initializeAsync();
+    this.refreshSubagents();
   }
 
   /**
    * Asynchronously initializes the tool by loading available subagents
    * and updating the description and schema.
    */
-  private async initializeAsync(): Promise<void> {
+  async refreshSubagents(): Promise<void> {
     try {
       this.availableSubagents = await this.subagentManager.listSubagents();
       this.updateDescriptionAndSchema();
     } catch (error) {
       console.warn('Failed to load subagents for Task tool:', error);
       this.availableSubagents = [];
       this.updateDescriptionAndSchema();
+    } finally {
+      // Update the client with the new tools
+      const geminiClient = this.config.getGeminiClient();
+      if (geminiClient) {
+        await geminiClient.setTools();
+      }
     }
   }
 
@@ -201,14 +210,6 @@ assistant: "I'm going to use the Task tool to launch the with the greeting-respo
     }
   }
 
-  /**
-   * Refreshes the available subagents and updates the tool description.
-   * This can be called when subagents are added or removed.
-   */
-  async refreshSubagents(): Promise<void> {
-    await this.initializeAsync();
-  }
-
   override validateToolParams(params: TaskParams): string | null {
     // Validate required fields
     if (