diff --git a/packages/cli/src/nonInteractive/io/BaseJsonOutputAdapter.test.ts b/packages/cli/src/nonInteractive/io/BaseJsonOutputAdapter.test.ts index be04b7f2b..dbb746a10 100644 --- a/packages/cli/src/nonInteractive/io/BaseJsonOutputAdapter.test.ts +++ b/packages/cli/src/nonInteractive/io/BaseJsonOutputAdapter.test.ts @@ -1086,6 +1086,26 @@ describe('BaseJsonOutputAdapter', () => { }); }); + describe('emitToolProgress', () => { + it('should be a no-op in base class (does not emit any message)', () => { + const request: ToolCallRequestInfo = { + callId: 'tool-call-1', + name: 'mcp__echo-test__echo', + args: {}, + isClientInitiated: false, + prompt_id: '', + }; + adapter.emitToolProgress(request, { + type: 'mcp_tool_progress', + progress: 1, + total: 10, + message: 'Echo: 1', + }); + + expect(adapter.emittedMessages).toHaveLength(0); + }); + }); + describe('buildResultMessage', () => { beforeEach(() => { adapter.startAssistantMessage(); diff --git a/packages/cli/src/nonInteractive/io/BaseJsonOutputAdapter.ts b/packages/cli/src/nonInteractive/io/BaseJsonOutputAdapter.ts index 072497000..b0d6736a5 100644 --- a/packages/cli/src/nonInteractive/io/BaseJsonOutputAdapter.ts +++ b/packages/cli/src/nonInteractive/io/BaseJsonOutputAdapter.ts @@ -12,6 +12,7 @@ import type { SessionMetrics, ServerGeminiStreamEvent, TaskResultDisplay, + McpToolProgressData, } from '@qwen-code/qwen-code-core'; import { GeminiEventType, @@ -82,6 +83,18 @@ export interface MessageEmitter { parentToolUseId?: string | null, ): void; emitSystemMessage(subtype: string, data?: unknown): void; + /** + * Emits a tool progress stream event. + * Only emits when the adapter supports partial messages (stream mode). + * In non-streaming mode, this is a no-op. + * + * @param request - Tool call request info + * @param progress - Structured MCP progress data + */ + emitToolProgress( + request: ToolCallRequestInfo, + progress: McpToolProgressData, + ): void; } /** @@ -1051,6 +1064,22 @@ export abstract class BaseJsonOutputAdapter { this.emitMessageImpl(systemMessage); } + /** + * Emits a tool progress stream event. + * Default implementation is a no-op. StreamJsonOutputAdapter overrides this + * to emit stream events when includePartialMessages is enabled. + * + * @param _request - Tool call request info + * @param _progress - Structured MCP progress data + */ + emitToolProgress( + _request: ToolCallRequestInfo, + _progress: McpToolProgressData, + ): void { + // No-op in base class. Only StreamJsonOutputAdapter emits tool progress + // as stream events when includePartialMessages is enabled. + } + /** * Builds a result message from options. * Helper method used by both emitResult implementations. diff --git a/packages/cli/src/nonInteractive/io/StreamJsonOutputAdapter.test.ts b/packages/cli/src/nonInteractive/io/StreamJsonOutputAdapter.test.ts index ff3aa1f5d..96977d5b0 100644 --- a/packages/cli/src/nonInteractive/io/StreamJsonOutputAdapter.test.ts +++ b/packages/cli/src/nonInteractive/io/StreamJsonOutputAdapter.test.ts @@ -882,6 +882,115 @@ describe('StreamJsonOutputAdapter', () => { }); }); + describe('emitToolProgress', () => { + const mockRequest = { + callId: 'tool-call-1', + name: 'mcp__echo-test__echo', + args: {}, + isClientInitiated: false, + prompt_id: '', + }; + + it('should emit tool_progress stream event when includePartialMessages is true', () => { + adapter = new StreamJsonOutputAdapter(mockConfig, true); + stdoutWriteSpy.mockClear(); + + adapter.emitToolProgress(mockRequest, { + type: 'mcp_tool_progress', + progress: 1, + total: 10, + message: 'Echo: 1', + }); + + expect(stdoutWriteSpy).toHaveBeenCalledTimes(1); + const output = stdoutWriteSpy.mock.calls[0][0] as string; + const parsed = JSON.parse(output); + + expect(parsed.type).toBe('stream_event'); + expect(parsed.parent_tool_use_id).toBeNull(); + expect(parsed.session_id).toBe('test-session-id'); + expect(parsed.uuid).toBeDefined(); + expect(parsed.event).toEqual({ + type: 'tool_progress', + tool_use_id: 'tool-call-1', + content: { + type: 'mcp_tool_progress', + progress: 1, + total: 10, + message: 'Echo: 1', + }, + }); + }); + + it('should not emit tool_progress when includePartialMessages is false', () => { + adapter = new StreamJsonOutputAdapter(mockConfig, false); + stdoutWriteSpy.mockClear(); + + adapter.emitToolProgress(mockRequest, { + type: 'mcp_tool_progress', + progress: 1, + total: 10, + message: 'Echo: 1', + }); + + expect(stdoutWriteSpy).not.toHaveBeenCalled(); + }); + + it('should emit multiple tool_progress events for sequential progress updates', () => { + adapter = new StreamJsonOutputAdapter(mockConfig, true); + stdoutWriteSpy.mockClear(); + + adapter.emitToolProgress(mockRequest, { + type: 'mcp_tool_progress', + progress: 1, + total: 3, + message: 'Echo: 1', + }); + adapter.emitToolProgress(mockRequest, { + type: 'mcp_tool_progress', + progress: 2, + total: 3, + message: 'Echo: 1, 2', + }); + adapter.emitToolProgress(mockRequest, { + type: 'mcp_tool_progress', + progress: 3, + total: 3, + message: 'Echo: 1, 2, 3', + }); + + expect(stdoutWriteSpy).toHaveBeenCalledTimes(3); + + const events = stdoutWriteSpy.mock.calls.map( + (call: unknown[]) => JSON.parse(call[0] as string).event, + ); + expect(events[0].content).toEqual({ + type: 'mcp_tool_progress', + progress: 1, + total: 3, + message: 'Echo: 1', + }); + expect(events[1].content).toEqual({ + type: 'mcp_tool_progress', + progress: 2, + total: 3, + message: 'Echo: 1, 2', + }); + expect(events[2].content).toEqual({ + type: 'mcp_tool_progress', + progress: 3, + total: 3, + message: 'Echo: 1, 2, 3', + }); + + // All events should share the same tool_use_id + for (const event of events) { + expect(event.type).toBe('tool_progress'); + expect(event.tool_use_id).toBe('tool-call-1'); + } + }); + }); + describe('getSessionId and getModel', () => { beforeEach(() => { adapter = new StreamJsonOutputAdapter(mockConfig, false); diff --git a/packages/cli/src/nonInteractive/io/StreamJsonOutputAdapter.ts b/packages/cli/src/nonInteractive/io/StreamJsonOutputAdapter.ts index af2f0bb6c..bf76d025c 100644 --- a/packages/cli/src/nonInteractive/io/StreamJsonOutputAdapter.ts +++ b/packages/cli/src/nonInteractive/io/StreamJsonOutputAdapter.ts @@ -5,7 +5,11 @@ */ import { randomUUID } from 'node:crypto'; -import type { Config } from '@qwen-code/qwen-code-core'; +import type { + Config, + ToolCallRequestInfo, + McpToolProgressData, +} from '@qwen-code/qwen-code-core'; import type { CLIAssistantMessage, CLIMessage, @@ -267,6 +271,32 @@ export class StreamJsonOutputAdapter } } + /** + * Emits a tool progress stream event when partial messages are enabled. + * This overrides the no-op in BaseJsonOutputAdapter. + */ + override emitToolProgress( + request: ToolCallRequestInfo, + progress: McpToolProgressData, + ): void { + if (!this.includePartialMessages) { + return; + } + + const partial: CLIPartialAssistantMessage = { + type: 'stream_event', + uuid: randomUUID(), + session_id: this.getSessionId(), + parent_tool_use_id: null, + event: { + type: 'tool_progress', + tool_use_id: request.callId, + content: progress, + }, + }; + this.emitMessageImpl(partial); + } + /** * Emits stream events when partial messages are enabled. * This is a private method specific to StreamJsonOutputAdapter. diff --git a/packages/cli/src/nonInteractive/types.ts b/packages/cli/src/nonInteractive/types.ts index 1d5e800dd..84c2d0ff7 100644 --- a/packages/cli/src/nonInteractive/types.ts +++ b/packages/cli/src/nonInteractive/types.ts @@ -1,5 +1,8 @@ /* eslint-disable @typescript-eslint/no-explicit-any */ -import type { SubagentConfig } from '@qwen-code/qwen-code-core'; +import type { + SubagentConfig, + McpToolProgressData, +} from '@qwen-code/qwen-code-core'; /** * Annotation for attaching metadata to content blocks @@ -236,12 +239,19 @@ export interface MessageStopStreamEvent { type: 'message_stop'; } +export interface ToolProgressStreamEvent { + type: 'tool_progress'; + tool_use_id: string; + content: McpToolProgressData; +} + export type StreamEvent = | MessageStartStreamEvent | ContentBlockStartEvent | ContentBlockDeltaEvent | ContentBlockStopEvent - | MessageStopStreamEvent; + | MessageStopStreamEvent + | ToolProgressStreamEvent; export interface CLIPartialAssistantMessage { type: 'stream_event'; diff --git a/packages/cli/src/nonInteractiveCli.test.ts b/packages/cli/src/nonInteractiveCli.test.ts index 2931118fc..6a6b33b87 100644 --- a/packages/cli/src/nonInteractiveCli.test.ts +++ b/packages/cli/src/nonInteractiveCli.test.ts @@ -296,7 +296,9 @@ describe('runNonInteractive', () => { mockConfig, expect.objectContaining({ name: 'testTool' }), expect.any(AbortSignal), - undefined, + expect.objectContaining({ + outputUpdateHandler: expect.any(Function), + }), ); // Verify first call has isContinuation: false expect(mockGeminiClient.sendMessageStream).toHaveBeenNthCalledWith( @@ -641,7 +643,9 @@ describe('runNonInteractive', () => { mockConfig, expect.objectContaining({ name: 'testTool' }), expect.any(AbortSignal), - undefined, + expect.objectContaining({ + outputUpdateHandler: expect.any(Function), + }), ); // JSON adapter emits array of messages, last one is result with stats diff --git a/packages/cli/src/nonInteractiveCli.ts b/packages/cli/src/nonInteractiveCli.ts index ef96d8e63..129bec380 100644 --- a/packages/cli/src/nonInteractiveCli.ts +++ b/packages/cli/src/nonInteractiveCli.ts @@ -41,6 +41,7 @@ import { normalizePartList, extractPartsFromUserMessage, buildSystemMessage, + createToolProgressHandler, createTaskToolProgressHandler, computeUsageFromMetrics, } from './utils/nonInteractiveHelpers.js'; @@ -313,31 +314,29 @@ export async function runNonInteractive( ? options.controlService.permission.getToolCallUpdateCallback() : undefined; - // Create output handler for Task tool (for subagent execution) + // Build outputUpdateHandler for this tool call. + // Task tool has its own complex handler (subagent messages). + // All other tools with canUpdateOutput=true (e.g., MCP tools) + // get a generic handler that emits progress via the adapter. const isTaskTool = finalRequestInfo.name === 'task'; - const taskToolProgress = isTaskTool + const { handler: outputUpdateHandler } = isTaskTool ? createTaskToolProgressHandler( config, finalRequestInfo.callId, adapter, ) - : undefined; - const taskToolProgressHandler = taskToolProgress?.handler; + : createToolProgressHandler(finalRequestInfo, adapter); const toolResponse = await executeToolCall( config, finalRequestInfo, abortController.signal, - taskToolProgressHandler || toolCallUpdateCallback - ? { - ...(taskToolProgressHandler && { - outputUpdateHandler: taskToolProgressHandler, - }), - ...(toolCallUpdateCallback && { - onToolCallsUpdate: toolCallUpdateCallback, - }), - } - : undefined, + { + outputUpdateHandler, + ...(toolCallUpdateCallback && { + onToolCallsUpdate: toolCallUpdateCallback, + }), + }, ); // Note: In JSON mode, subagent messages are automatically added to the main diff --git a/packages/cli/src/ui/components/messages/ToolMessage.tsx b/packages/cli/src/ui/components/messages/ToolMessage.tsx index afc16317c..c0c981197 100644 --- a/packages/cli/src/ui/components/messages/ToolMessage.tsx +++ b/packages/cli/src/ui/components/messages/ToolMessage.tsx @@ -20,6 +20,7 @@ import type { PlanResultDisplay, AnsiOutput, Config, + McpToolProgressData, } from '@qwen-code/qwen-code-core'; import { AgentExecutionDisplay } from '../subagents/index.js'; import { PlanSummaryDisplay } from '../PlanSummaryDisplay.js'; @@ -113,6 +114,22 @@ const useResultDisplayRenderer = ( }; } + // Check for McpToolProgressData + if ( + typeof resultDisplay === 'object' && + resultDisplay !== null && + 'type' in resultDisplay && + resultDisplay.type === 'mcp_tool_progress' + ) { + const progress = resultDisplay as McpToolProgressData; + const msg = progress.message ?? `Progress: ${progress.progress}`; + const totalStr = progress.total != null ? `/${progress.total}` : ''; + return { + type: 'string', + data: `⏳ [${progress.progress}${totalStr}] ${msg}`, + }; + } + // Check for AnsiOutput if ( typeof resultDisplay === 'object' && diff --git a/packages/cli/src/utils/nonInteractiveHelpers.test.ts b/packages/cli/src/utils/nonInteractiveHelpers.test.ts index b5565bb52..1892e6e41 100644 --- a/packages/cli/src/utils/nonInteractiveHelpers.test.ts +++ b/packages/cli/src/utils/nonInteractiveHelpers.test.ts @@ -29,6 +29,7 @@ import { extractUsageFromGeminiClient, computeUsageFromMetrics, buildSystemMessage, + createToolProgressHandler, createTaskToolProgressHandler, functionResponsePartsToString, toolResultContent, @@ -621,6 +622,115 @@ describe('buildSystemMessage', () => { }); }); +describe('createToolProgressHandler', () => { + const mockRequest = { + callId: 'tool-call-1', + name: 'mcp__echo-test__echo', + args: {}, + isClientInitiated: false, + prompt_id: '', + }; + + it('should call emitToolProgress with request and McpToolProgressData', () => { + const mockAdapter = { + emitToolProgress: vi.fn(), + } as unknown as JsonOutputAdapterInterface; + + const { handler } = createToolProgressHandler(mockRequest, mockAdapter); + + const progressData = { + type: 'mcp_tool_progress' as const, + progress: 1, + total: 10, + message: 'Echo: 1', + }; + handler('tool-call-1', progressData); + + expect(mockAdapter.emitToolProgress).toHaveBeenCalledWith( + mockRequest, + progressData, + ); + }); + + it('should not call emitToolProgress for non-McpToolProgressData output', () => { + const mockAdapter = { + emitToolProgress: vi.fn(), + } as unknown as JsonOutputAdapterInterface; + + const { handler } = createToolProgressHandler( + { ...mockRequest, name: 'test_tool' }, + mockAdapter, + ); + + // Pass a non-McpToolProgressData ToolResultDisplay (e.g., FileDiff) + handler('tool-call-1', { + fileDiff: 'diff', + fileName: 'test.ts', + originalContent: null, + newContent: 'new', + }); + + expect(mockAdapter.emitToolProgress).not.toHaveBeenCalled(); + + // Also test with a plain string — should not emit + handler('tool-call-1', 'plain string progress'); + + expect(mockAdapter.emitToolProgress).not.toHaveBeenCalled(); + }); + + it('should forward multiple progress updates', () => { + const mockAdapter = { + emitToolProgress: vi.fn(), + } as unknown as JsonOutputAdapterInterface; + + const browserRequest = { + ...mockRequest, + name: 'mcp__browser__navigate', + }; + const { handler } = createToolProgressHandler(browserRequest, mockAdapter); + + const progress1 = { + type: 'mcp_tool_progress' as const, + progress: 1, + total: 3, + message: 'Navigating...', + }; + const progress2 = { + type: 'mcp_tool_progress' as const, + progress: 2, + total: 3, + message: 'Loading page...', + }; + const progress3 = { + type: 'mcp_tool_progress' as const, + progress: 3, + total: 3, + message: 'Complete', + }; + + handler('tool-call-1', progress1); + handler('tool-call-1', progress2); + handler('tool-call-1', progress3); + + expect(mockAdapter.emitToolProgress).toHaveBeenCalledTimes(3); + expect(mockAdapter.emitToolProgress).toHaveBeenNthCalledWith( + 1, + browserRequest, + progress1, + ); + expect(mockAdapter.emitToolProgress).toHaveBeenNthCalledWith( + 2, + browserRequest, + progress2, + ); + expect(mockAdapter.emitToolProgress).toHaveBeenNthCalledWith( + 3, + browserRequest, + progress3, + ); + }); +}); + describe('createTaskToolProgressHandler', () => { let mockAdapter: JsonOutputAdapterInterface; let mockConfig: Config; diff --git a/packages/cli/src/utils/nonInteractiveHelpers.ts b/packages/cli/src/utils/nonInteractiveHelpers.ts index 4e2317b2e..a5b23e6c9 100644 --- a/packages/cli/src/utils/nonInteractiveHelpers.ts +++ b/packages/cli/src/utils/nonInteractiveHelpers.ts @@ -12,6 +12,7 @@ import type { ToolCallRequestInfo, ToolCallResponseInfo, SessionMetrics, + McpToolProgressData, } from '@qwen-code/qwen-code-core'; import { OutputFormat, @@ -26,7 +27,10 @@ import type { PermissionMode, CLISystemMessage, } from '../nonInteractive/types.js'; -import type { JsonOutputAdapterInterface } from '../nonInteractive/io/BaseJsonOutputAdapter.js'; +import type { + JsonOutputAdapterInterface, + MessageEmitter, +} from '../nonInteractive/io/BaseJsonOutputAdapter.js'; import { computeSessionStats } from '../ui/utils/computeStats.js'; import { getAvailableCommands } from '../nonInteractiveCliCommands.js'; @@ -291,6 +295,45 @@ export async function buildSystemMessage( return systemMessage; } +function isMcpToolProgressData( + output: ToolResultDisplay, +): output is McpToolProgressData { + return ( + typeof output === 'object' && + output !== null && + 'type' in output && + (output as McpToolProgressData).type === 'mcp_tool_progress' + ); +} + +/** + * Creates a generic output update handler for tools with canUpdateOutput=true. + * This handler forwards MCP progress data (McpToolProgressData) as tool_progress + * stream events via the adapter. Progress events are only emitted when the adapter + * supports partial messages (i.e., includePartialMessages is true). + * + * @param request - Tool call request info + * @param adapter - The adapter instance for emitting messages + * @returns An object containing the output update handler + */ +export function createToolProgressHandler( + request: ToolCallRequestInfo, + adapter: MessageEmitter, +): { + handler: OutputUpdateHandler; +} { + const handler: OutputUpdateHandler = ( + _callId: string, + output: ToolResultDisplay, + ) => { + if (isMcpToolProgressData(output)) { + adapter.emitToolProgress(request, output); + } + }; + + return { handler }; +} + /** * Creates an output update handler specifically for Task tool subagent execution. * This handler monitors TaskResultDisplay updates and converts them to protocol messages diff --git a/packages/core/src/tools/mcp-client.ts b/packages/core/src/tools/mcp-client.ts index e61318edc..0e8fc9cce 100644 --- a/packages/core/src/tools/mcp-client.ts +++ b/packages/core/src/tools/mcp-client.ts @@ -638,6 +638,7 @@ export async function discoverTools( return []; } + const mcpTimeout = mcpServerConfig.timeout ?? MCP_DEFAULT_TIMEOUT_MSEC; const discoveredTools: DiscoveredMCPTool[] = []; for (const funcDecl of tool.functionDeclarations) { try { @@ -655,6 +656,8 @@ export async function discoverTools( mcpServerConfig.trust, undefined, cliConfig, + mcpClient, // raw MCP Client for direct callTool with progress + mcpTimeout, ), ); } catch (error) { diff --git a/packages/core/src/tools/mcp-tool.test.ts b/packages/core/src/tools/mcp-tool.test.ts index 680fa9299..997449eba 100644 --- a/packages/core/src/tools/mcp-tool.test.ts +++ b/packages/core/src/tools/mcp-tool.test.ts @@ -8,9 +8,13 @@ import type { Mocked } from 'vitest'; import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; import { safeJsonStringify } from '../utils/safeJsonStringify.js'; -import { DiscoveredMCPTool, generateValidName } from './mcp-tool.js'; // Added getStringifiedResultForDisplay +import { + DiscoveredMCPTool, + generateValidName, + type McpDirectClient, +} from './mcp-tool.js'; import type { ToolResult } from './tools.js'; -import { ToolConfirmationOutcome } from './tools.js'; // Added ToolConfirmationOutcome +import { ToolConfirmationOutcome } from './tools.js'; import type { CallableTool, Part } from '@google/genai'; import { ToolErrorType } from './tool-error.js'; @@ -958,4 +962,157 @@ describe('DiscoveredMCPTool', () => { expect(description).toBe('{"param":"testValue","param2":"anotherOne"}'); }); }); + + describe('streaming progress for long-running MCP tools', () => { + it('should have canUpdateOutput set to true so the scheduler creates liveOutputCallback', () => { + // For long-running MCP tools (e.g., browseruse), the scheduler needs + // canUpdateOutput=true to create a liveOutputCallback. Without this, + // users see no progress during potentially minutes-long operations. + expect(tool.canUpdateOutput).toBe(true); + }); + + it('should forward MCP progress notifications to updateOutput callback during execution', async () => { + const params = { param: 'https://example.com' }; + + // Create a mock MCP direct client that simulates progress notifications. + // When callTool is called with an onprogress callback, it invokes + // the callback to simulate the MCP server sending progress updates. + const mockMcpClient: McpDirectClient = { + callTool: vi.fn(async (_params, _schema, options) => { + // Simulate 3 progress notifications from the MCP server + for (let i = 1; i <= 3; i++) { + await new Promise((resolve) => setTimeout(resolve, 10)); + options?.onprogress?.({ + progress: i, + total: 3, + message: `Step ${i} of 3`, + }); + } + return { + content: [ + { + type: 'text', + text: 'Browser automation completed successfully.', + }, + ], + }; + }), + }; + + // Create a tool with the direct MCP client + const streamingTool = new DiscoveredMCPTool( + mockCallableToolInstance, + serverName, + serverToolName, + baseDescription, + inputSchema, + undefined, // trust + undefined, // nameOverride + undefined, // cliConfig + mockMcpClient, + ); + + const invocation = streamingTool.build(params); + const updateOutputSpy = vi.fn(); + + const result = await invocation.execute( + new AbortController().signal, + updateOutputSpy, + ); + + // The final result should still be correct + expect(result.llmContent).toEqual([ + { text: 'Browser automation completed successfully.' }, + ]); + + // The updateOutput callback SHOULD have been called at least once + // with intermediate progress, so users can see what's happening + // during the long wait. + expect(updateOutputSpy).toHaveBeenCalled(); + expect(updateOutputSpy).toHaveBeenCalledTimes(3); + // Verify progress data contains structured MCP progress info + expect(updateOutputSpy).toHaveBeenCalledWith( + expect.objectContaining({ + type: 'mcp_tool_progress', + progress: 1, + total: 3, + message: 'Step 1 of 3', + }), + ); + expect(updateOutputSpy).toHaveBeenCalledWith( + expect.objectContaining({ + type: 'mcp_tool_progress', + progress: 3, + total: 3, + message: 'Step 3 of 3', + }), + ); + }); + + it('should show incremental progress for multi-step browser automation', async () => { + const params = { param: 'fill-form' }; + const steps = [ + 'Navigating to page...', + 'Filling username field...', + 'Filling password field...', + 'Clicking submit...', + ]; + + const mockMcpClient: McpDirectClient = { + callTool: vi.fn(async (_params, _schema, options) => { + for (let i = 0; i < steps.length; i++) { + await new Promise((resolve) => setTimeout(resolve, 10)); + options?.onprogress?.({ + progress: i + 1, + total: steps.length, + message: steps[i], + }); + } + return { + content: [{ type: 'text', text: steps.join('\n') }], + }; + }), + }; + + const streamingTool = new DiscoveredMCPTool( + mockCallableToolInstance, + serverName, + serverToolName, + baseDescription, + inputSchema, + undefined, + undefined, + undefined, + mockMcpClient, + ); + + const invocation = streamingTool.build(params); + const receivedUpdates: unknown[] = []; + const updateOutputCallback = (output: unknown) => { + receivedUpdates.push(output); + }; + + await invocation.execute( + new AbortController().signal, + updateOutputCallback, + ); + + // User should have received one update per step + expect(receivedUpdates.length).toBeGreaterThan(0); + expect(receivedUpdates).toHaveLength(steps.length); + // Each update should be structured McpToolProgressData + expect(receivedUpdates[0]).toEqual({ + type: 'mcp_tool_progress', + progress: 1, + total: steps.length, + message: 'Navigating to page...', + }); + expect(receivedUpdates[3]).toEqual({ + type: 'mcp_tool_progress', + progress: 4, + total: steps.length, + message: 'Clicking submit...', + }); + }); + }); }); diff --git a/packages/core/src/tools/mcp-tool.ts b/packages/core/src/tools/mcp-tool.ts index 15f461e90..4917fa2c3 100644 --- a/packages/core/src/tools/mcp-tool.ts +++ b/packages/core/src/tools/mcp-tool.ts @@ -10,7 +10,9 @@ import type { ToolInvocation, ToolMcpConfirmationDetails, ToolResult, + ToolResultDisplay, ToolConfirmationPayload, + McpToolProgressData, } from './tools.js'; import { BaseDeclarativeTool, @@ -24,6 +26,40 @@ import type { Config } from '../config/config.js'; type ToolParams = Record; +/** + * Minimal interface for the raw MCP Client's callTool method. + * This avoids a direct import of @modelcontextprotocol/sdk in this file, + * keeping the dependency contained in mcp-client.ts. + */ +export interface McpDirectClient { + callTool( + params: { name: string; arguments?: Record }, + resultSchema?: unknown, + options?: { + onprogress?: (progress: { + progress: number; + total?: number; + message?: string; + }) => void; + timeout?: number; + signal?: AbortSignal; + }, + ): Promise; +} + +/** The result shape returned by MCP SDK Client.callTool(). */ +interface McpCallToolResult { + content?: Array<{ + type: string; + text?: string; + data?: string; + mimeType?: string; + [key: string]: unknown; + }>; + isError?: boolean; + [key: string]: unknown; +} + // Discriminated union for MCP Content Blocks to ensure type safety. type McpTextBlock = { type: 'text'; @@ -72,6 +108,8 @@ class DiscoveredMCPToolInvocation extends BaseToolInvocation< readonly trust?: boolean, params: ToolParams = {}, private readonly cliConfig?: Config, + private readonly mcpClient?: McpDirectClient, + private readonly mcpTimeout?: number, ) { super(params); } @@ -135,7 +173,91 @@ class DiscoveredMCPToolInvocation extends BaseToolInvocation< return false; } - async execute(signal: AbortSignal): Promise { + async execute( + signal: AbortSignal, + updateOutput?: (output: ToolResultDisplay) => void, + ): Promise { + // Use direct MCP client if available (supports progress notifications), + // otherwise fall back to the @google/genai mcpToTool wrapper. + if (this.mcpClient) { + return this.executeWithDirectClient(signal, updateOutput); + } + return this.executeWithCallableTool(signal); + } + + /** + * Execute using the raw MCP SDK Client, which supports progress + * notifications via the onprogress callback. This enables real-time + * streaming of progress updates to the user during long-running + * MCP tool calls (e.g., browser automation). + */ + private async executeWithDirectClient( + signal: AbortSignal, + updateOutput?: (output: ToolResultDisplay) => void, + ): Promise { + const callToolResult = await this.mcpClient!.callTool( + { + name: this.serverToolName, + arguments: this.params as Record, + }, + undefined, + { + onprogress: (progress) => { + if (updateOutput) { + const progressData: McpToolProgressData = { + type: 'mcp_tool_progress', + progress: progress.progress, + ...(progress.total != null && { total: progress.total }), + ...(progress.message != null && { message: progress.message }), + }; + updateOutput(progressData); + } + }, + timeout: this.mcpTimeout, + signal, + }, + ); + + // Wrap the raw CallToolResult into the Part[] format that the + // existing transform/display functions expect. + const rawResponseParts = wrapMcpCallToolResultAsParts( + this.serverToolName, + callToolResult, + ); + + // Ensure the response is not an error + if (this.isMCPToolError(rawResponseParts)) { + const errorMessage = `MCP tool '${ + this.serverToolName + }' reported tool error for function call: ${safeJsonStringify({ + name: this.serverToolName, + args: this.params, + })} with response: ${safeJsonStringify(rawResponseParts)}`; + return { + llmContent: errorMessage, + returnDisplay: `Error: MCP tool '${this.serverToolName}' reported an error.`, + error: { + message: errorMessage, + type: ToolErrorType.MCP_TOOL_ERROR, + }, + }; + } + + const transformedParts = transformMcpContentToParts(rawResponseParts); + + return { + llmContent: transformedParts, + returnDisplay: getStringifiedResultForDisplay(rawResponseParts), + }; + } + + /** + * Fallback: execute using the @google/genai CallableTool wrapper. + * This path does NOT support progress notifications. + */ + private async executeWithCallableTool( + signal: AbortSignal, + ): Promise { const functionCalls: FunctionCall[] = [ { name: this.serverToolName, @@ -217,6 +339,8 @@ export class DiscoveredMCPTool extends BaseDeclarativeTool< readonly trust?: boolean, nameOverride?: string, private readonly cliConfig?: Config, + private readonly mcpClient?: McpDirectClient, + private readonly mcpTimeout?: number, ) { super( nameOverride ?? generateValidName(serverToolName), @@ -225,7 +349,7 @@ export class DiscoveredMCPTool extends BaseDeclarativeTool< Kind.Other, parameterSchema, true, // isOutputMarkdown - false, // canUpdateOutput + true, // canUpdateOutput — enables streaming progress for MCP tools ); } @@ -239,6 +363,8 @@ export class DiscoveredMCPTool extends BaseDeclarativeTool< this.trust, `${this.serverName}__${this.serverToolName}`, this.cliConfig, + this.mcpClient, + this.mcpTimeout, ); } @@ -253,10 +379,37 @@ export class DiscoveredMCPTool extends BaseDeclarativeTool< this.trust, params, this.cliConfig, + this.mcpClient, + this.mcpTimeout, ); } } +/** + * Wraps a raw MCP CallToolResult into the Part[] format that the + * existing transform/display functions expect. This bridges the gap + * between the raw MCP SDK response and the @google/genai Part format. + */ +function wrapMcpCallToolResultAsParts( + toolName: string, + result: { + content?: Array<{ [key: string]: unknown }>; + isError?: boolean; + }, +): Part[] { + const response = result.isError + ? { error: result, content: result.content } + : result; + return [ + { + functionResponse: { + name: toolName, + response, + }, + }, + ]; +} + function transformTextBlock(block: McpTextBlock): Part { return { text: block.text }; } diff --git a/packages/core/src/tools/tools.ts b/packages/core/src/tools/tools.ts index 7b3c893e6..96ae53402 100644 --- a/packages/core/src/tools/tools.ts +++ b/packages/core/src/tools/tools.ts @@ -470,13 +470,28 @@ export interface AnsiOutputDisplay { ansiOutput: AnsiOutput; } +/** + * Structured progress data following the MCP notifications/progress spec. + * @see https://modelcontextprotocol.io/specification/2025-06-18/basic/utilities/progress + */ +export interface McpToolProgressData { + type: 'mcp_tool_progress'; + /** Current progress value (must increase with each notification) */ + progress: number; + /** Optional total value indicating the operation's target */ + total?: number; + /** Optional human-readable progress message */ + message?: string; +} + export type ToolResultDisplay = | string | FileDiff | TodoResultDisplay | PlanResultDisplay | TaskResultDisplay - | AnsiOutputDisplay; + | AnsiOutputDisplay + | McpToolProgressData; export interface FileDiff { fileDiff: string;