Merge pull request #1691 from QwenLM/fix/subagent-tool-restriction

fix(core): enforce tool restrictions in subagents

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

@@ -38,6 +38,8 @@ import {
   SubAgentEventEmitter,
   SubAgentEventType,
   type SubAgentStreamTextEvent,
+  type SubAgentToolCallEvent,
+  type SubAgentToolResultEvent,
 } from './subagent-events.js';
 import type {
   ModelConfig,
@@ -933,5 +935,165 @@ describe('subagent.ts', () => {
         expect(mockSendMessageStream).toHaveBeenCalledTimes(2);
       });
     });
+
+    describe('runNonInteractive - Tool Restriction Enforcement (Issue #1121)', () => {
+      const promptConfig: PromptConfig = { systemPrompt: 'Execute task.' };
+
+      it('should NOT execute tools that are not in the allowed tools list', async () => {
+        // Define two tools: one allowed (read_file), one not allowed (edit_file)
+        const readFileToolDef: FunctionDeclaration = {
+          name: 'read_file',
+          description: 'Reads a file',
+          parameters: { type: Type.OBJECT, properties: {} },
+        };
+        const editFileToolDef: FunctionDeclaration = {
+          name: 'edit_file',
+          description: 'Edits a file',
+          parameters: { type: Type.OBJECT, properties: {} },
+        };
+
+        // Track which tools were executed
+        const executedTools: string[] = [];
+
+        const readFileInvocation = {
+          params: { path: 'test.txt' },
+          getDescription: vi.fn().mockReturnValue('Read file'),
+          toolLocations: vi.fn().mockReturnValue([]),
+          shouldConfirmExecute: vi.fn().mockResolvedValue(false),
+          execute: vi.fn().mockImplementation(async () => {
+            executedTools.push('read_file');
+            return {
+              llmContent: 'file contents',
+              returnDisplay: 'Read file contents',
+            };
+          }),
+        };
+
+        const editFileInvocation = {
+          params: { path: 'test.txt', content: 'malicious content' },
+          getDescription: vi.fn().mockReturnValue('Edit file'),
+          toolLocations: vi.fn().mockReturnValue([]),
+          shouldConfirmExecute: vi.fn().mockResolvedValue(false),
+          execute: vi.fn().mockImplementation(async () => {
+            executedTools.push('edit_file');
+            return {
+              llmContent: 'file edited',
+              returnDisplay: 'Edited file',
+            };
+          }),
+        };
+
+        const readFileTool = {
+          name: 'read_file',
+          displayName: 'Read File',
+          description: 'Read file contents',
+          kind: 'READ' as const,
+          schema: readFileToolDef,
+          build: vi.fn().mockImplementation(() => readFileInvocation),
+          canUpdateOutput: false,
+          isOutputMarkdown: true,
+        } as unknown as AnyDeclarativeTool;
+
+        const editFileTool = {
+          name: 'edit_file',
+          displayName: 'Edit File',
+          description: 'Edit file contents',
+          kind: 'WRITE' as const,
+          schema: editFileToolDef,
+          build: vi.fn().mockImplementation(() => editFileInvocation),
+          canUpdateOutput: false,
+          isOutputMarkdown: true,
+        } as unknown as AnyDeclarativeTool;
+
+        const { config } = await createMockConfig({
+          // Only return read_file in the filtered list (this is what the subagent should see)
+          getFunctionDeclarationsFiltered: vi
+            .fn()
+            .mockReturnValue([readFileToolDef]),
+          // But the full registry has both tools (simulating the bug)
+          getFunctionDeclarations: vi
+            .fn()
+            .mockReturnValue([readFileToolDef, editFileToolDef]),
+          getTool: vi.fn().mockImplementation((name: string) => {
+            if (name === 'read_file') return readFileTool;
+            if (name === 'edit_file') return editFileTool;
+            return undefined;
+          }),
+        });
+
+        // Only allow read_file in the subagent's tool config
+        const toolConfig: ToolConfig = { tools: ['read_file'] };
+
+        // Model calls BOTH read_file (allowed) AND edit_file (NOT allowed)
+        // This simulates the bug where the model hallucinates an unauthorized tool call
+        mockSendMessageStream.mockImplementation(
+          createMockStream([
+            [
+              {
+                id: 'call_read',
+                name: 'read_file',
+                args: { path: 'test.txt' },
+              },
+              {
+                id: 'call_edit',
+                name: 'edit_file', // This tool is NOT in the allowed list!
+                args: { path: 'test.txt', content: 'malicious content' },
+              },
+            ],
+            'stop',
+          ]),
+        );
+
+        // Track emitted events
+        const toolCallEvents: SubAgentToolCallEvent[] = [];
+        const toolResultEvents: SubAgentToolResultEvent[] = [];
+
+        // Create event emitter BEFORE the scope and subscribe to events
+        const eventEmitter = new SubAgentEventEmitter();
+        eventEmitter.on(SubAgentEventType.TOOL_CALL, (event: unknown) => {
+          toolCallEvents.push(event as SubAgentToolCallEvent);
+        });
+        eventEmitter.on(SubAgentEventType.TOOL_RESULT, (event: unknown) => {
+          toolResultEvents.push(event as SubAgentToolResultEvent);
+        });
+
+        const scope = await SubAgentScope.create(
+          'test-agent',
+          config,
+          promptConfig,
+          defaultModelConfig,
+          defaultRunConfig,
+          toolConfig,
+          eventEmitter,
+        );
+
+        await scope.runNonInteractive(new ContextState());
+
+        // 1. Only allowed tool should be executed
+        expect(executedTools).toContain('read_file');
+        expect(executedTools).not.toContain('edit_file');
+        expect(editFileInvocation.execute).not.toHaveBeenCalled();
+
+        // 2. TOOL_CALL events should be emitted for BOTH tools (for visibility)
+        expect(toolCallEvents).toHaveLength(2);
+        expect(toolCallEvents.map((e) => e.name)).toContain('read_file');
+        expect(toolCallEvents.map((e) => e.name)).toContain('edit_file');
+
+        // 3. TOOL_RESULT events should be emitted for both
+        expect(toolResultEvents).toHaveLength(2);
+
+        // 4. Verify blocked tool result has success=false and error message
+        const editResult = toolResultEvents.find((e) => e.name === 'edit_file');
+        expect(editResult).toBeDefined();
+        expect(editResult!.success).toBe(false);
+        expect(editResult!.error).toContain('not found');
+        expect(editResult!.callId).toBe('call_edit');
+
+        // 5. Verify allowed tool result has success=true
+        const readResult = toolResultEvents.find((e) => e.name === 'read_file');
+        expect(readResult).toBeDefined();
+        expect(readResult!.success).toBe(true);
+      });
+    });
   });
 });

packages/core/src/subagents/subagent.ts

@@ -487,6 +487,7 @@ export class SubAgentScope {
             abortController,
             promptId,
             turnCounter,
+            toolsList,
             currentResponseId,
           );
         } else {
@@ -585,10 +586,67 @@ export class SubAgentScope {
     abortController: AbortController,
     promptId: string,
     currentRound: number,
+    toolsList: FunctionDeclaration[],
     responseId?: string,
   ): Promise<Content[]> {
     const toolResponseParts: Part[] = [];
 
+    // Build allowed tool names set for filtering
+    const allowedToolNames = new Set(toolsList.map((t) => t.name));
+
+    // Filter unauthorized tool calls before scheduling
+    const authorizedCalls: FunctionCall[] = [];
+    for (const fc of functionCalls) {
+      const callId = fc.id ?? `${fc.name}-${Date.now()}`;
+
+      if (!allowedToolNames.has(fc.name)) {
+        const toolName = String(fc.name);
+        const errorMessage = `Tool "${toolName}" not found. Tools must use the exact names provided.`;
+
+        // Emit TOOL_CALL event for visibility
+        this.eventEmitter?.emit(SubAgentEventType.TOOL_CALL, {
+          subagentId: this.subagentId,
+          round: currentRound,
+          callId,
+          name: toolName,
+          args: fc.args ?? {},
+          description: `Tool "${toolName}" not found`,
+          timestamp: Date.now(),
+        } as SubAgentToolCallEvent);
+
+        // Build function response part (used for both event and LLM)
+        const functionResponsePart = {
+          functionResponse: {
+            id: callId,
+            name: toolName,
+            response: { error: errorMessage },
+          },
+        };
+
+        // Emit TOOL_RESULT event with error (include responseParts for UI rendering)
+        this.eventEmitter?.emit(SubAgentEventType.TOOL_RESULT, {
+          subagentId: this.subagentId,
+          round: currentRound,
+          callId,
+          name: toolName,
+          success: false,
+          error: errorMessage,
+          responseParts: [functionResponsePart],
+          resultDisplay: errorMessage,
+          durationMs: 0,
+          timestamp: Date.now(),
+        } as SubAgentToolResultEvent);
+
+        // Record blocked tool call in stats
+        this.recordToolCallStats(toolName, false, 0, errorMessage);
+
+        // Add function response for LLM
+        toolResponseParts.push(functionResponsePart);
+        continue;
+      }
+      authorizedCalls.push(fc);
+    }
+
     // Build scheduler
     const responded = new Set<string>();
     let resolveBatch: (() => void) | null = null;
@@ -605,33 +663,8 @@ export class SubAgentScope {
               ? call.response.error?.message
               : undefined;
 
-          // Update aggregate stats
+          // Record stats
-          this.executionStats.totalToolCalls += 1;
+          this.recordToolCallStats(toolName, success, duration, errorMessage);
-          if (success) {
-            this.executionStats.successfulToolCalls += 1;
-          } else {
-            this.executionStats.failedToolCalls += 1;
-          }
-
-          // Per-tool usage
-          const tu = this.toolUsage.get(toolName) || {
-            count: 0,
-            success: 0,
-            failure: 0,
-            totalDurationMs: 0,
-            averageDurationMs: 0,
-          };
-          tu.count += 1;
-          if (success) {
-            tu.success += 1;
-          } else {
-            tu.failure += 1;
-            tu.lastError = errorMessage || 'Unknown error';
-          }
-          tu.totalDurationMs = (tu.totalDurationMs || 0) + duration;
-          tu.averageDurationMs =
-            tu.count > 0 ? tu.totalDurationMs / tu.count : 0;
-          this.toolUsage.set(toolName, tu);
 
           // Emit tool result event
           this.eventEmitter?.emit(SubAgentEventType.TOOL_RESULT, {
@@ -642,12 +675,6 @@ export class SubAgentScope {
             success,
             error: errorMessage,
             responseParts: call.response.responseParts,
-            /**
-             * Tools like todoWrite will add some extra contents to the result,
-             * making it unable to deserialize the `responseParts` to a JSON object.
-             * While `resultDisplay` is normally a string, if not we stringify it,
-             * so that we can deserialize it to a JSON object when needed.
-             */
             resultDisplay: call.response.resultDisplay
               ? typeof call.response.resultDisplay === 'string'
                 ? call.response.resultDisplay
@@ -657,14 +684,6 @@ export class SubAgentScope {
             timestamp: Date.now(),
           } as SubAgentToolResultEvent);
 
-          // Update statistics service
-          this.stats.recordToolCall(
-            toolName,
-            success,
-            duration,
-            this.toolUsage.get(toolName)?.lastError,
-          );
-
           // post-tool hook
           await this.hooks?.postToolUse?.({
             subagentId: this.subagentId,
@@ -736,7 +755,7 @@ export class SubAgentScope {
     });
 
     // Prepare requests and emit TOOL_CALL events
-    const requests: ToolCallRequestInfo[] = functionCalls.map((fc) => {
+    const requests: ToolCallRequestInfo[] = authorizedCalls.map((fc) => {
       const toolName = String(fc.name || 'unknown');
       const callId = fc.id ?? `${fc.name}-${Date.now()}`;
       const args = (fc.args ?? {}) as Record<string, unknown>;
@@ -902,6 +921,52 @@ export class SubAgentScope {
     }
   }
 
+  /**
+   * Records tool call statistics for both successful and failed tool calls.
+   * This includes updating aggregate stats, per-tool usage, and the statistics service.
+   */
+  private recordToolCallStats(
+    toolName: string,
+    success: boolean,
+    durationMs: number,
+    errorMessage?: string,
+  ): void {
+    // Update aggregate stats
+    this.executionStats.totalToolCalls += 1;
+    if (success) {
+      this.executionStats.successfulToolCalls += 1;
+    } else {
+      this.executionStats.failedToolCalls += 1;
+    }
+
+    // Per-tool usage
+    const tu = this.toolUsage.get(toolName) || {
+      count: 0,
+      success: 0,
+      failure: 0,
+      totalDurationMs: 0,
+      averageDurationMs: 0,
+    };
+    tu.count += 1;
+    if (success) {
+      tu.success += 1;
+    } else {
+      tu.failure += 1;
+      tu.lastError = errorMessage || 'Unknown error';
+    }
+    tu.totalDurationMs = (tu.totalDurationMs || 0) + durationMs;
+    tu.averageDurationMs = tu.count > 0 ? tu.totalDurationMs / tu.count : 0;
+    this.toolUsage.set(toolName, tu);
+
+    // Update statistics service
+    this.stats.recordToolCall(
+      toolName,
+      success,
+      durationMs,
+      this.toolUsage.get(toolName)?.lastError,
+    );
+  }
+
   private buildChatSystemPrompt(context: ContextState): string {
     if (!this.promptConfig.systemPrompt) {
       // This should ideally be caught in createChatObject, but serves as a safeguard.