From 14415316c0719d7a943aed23eb7216e72f4d3a7a Mon Sep 17 00:00:00 2001 From: "N. Taylor Mullen" Date: Wed, 18 Feb 2026 12:46:12 -0800 Subject: [PATCH] feat(core): add support for MCP progress updates (#19046) --- .../components/messages/ToolMessage.test.tsx | 27 +++++++ .../ui/components/messages/ToolMessage.tsx | 4 + .../src/ui/components/messages/ToolShared.tsx | 24 +++++- packages/cli/src/ui/hooks/toolMapping.ts | 6 ++ packages/cli/src/ui/hooks/useToolScheduler.ts | 2 + packages/cli/src/ui/types.ts | 2 + packages/core/src/scheduler/scheduler.test.ts | 27 +++++++ packages/core/src/scheduler/scheduler.ts | 22 +++++ .../core/src/scheduler/state-manager.test.ts | 32 ++++++++ packages/core/src/scheduler/state-manager.ts | 8 ++ packages/core/src/scheduler/types.ts | 2 + packages/core/src/tools/mcp-client.test.ts | 26 +++--- packages/core/src/tools/mcp-client.ts | 81 ++++++++++++++++++- packages/core/src/utils/events.ts | 21 +++++ 14 files changed, 270 insertions(+), 14 deletions(-) diff --git a/packages/cli/src/ui/components/messages/ToolMessage.test.tsx b/packages/cli/src/ui/components/messages/ToolMessage.test.tsx index 29012bbd26..34d6b5e52b 100644 --- a/packages/cli/src/ui/components/messages/ToolMessage.test.tsx +++ b/packages/cli/src/ui/components/messages/ToolMessage.test.tsx @@ -320,4 +320,31 @@ describe('', () => { ); expect(lastFrame()).toMatchSnapshot(); }); + + it('renders progress information appended to description for executing tools', () => { + const { lastFrame } = renderWithContext( + , + StreamingState.Responding, + ); + expect(lastFrame()).toContain( + 'A tool for testing (Working on it... - 42%)', + ); + }); + + it('renders only percentage when progressMessage is missing', () => { + const { lastFrame } = renderWithContext( + , + StreamingState.Responding, + ); + expect(lastFrame()).toContain('A tool for testing (75%)'); + }); }); diff --git a/packages/cli/src/ui/components/messages/ToolMessage.tsx b/packages/cli/src/ui/components/messages/ToolMessage.tsx index 06ad6b3f7b..557e0bd857 100644 --- a/packages/cli/src/ui/components/messages/ToolMessage.tsx +++ b/packages/cli/src/ui/components/messages/ToolMessage.tsx @@ -55,6 +55,8 @@ export const ToolMessage: React.FC = ({ embeddedShellFocused, ptyId, config, + progressMessage, + progressPercent, }) => { const isThisShellFocused = checkIsShellFocused( name, @@ -89,6 +91,8 @@ export const ToolMessage: React.FC = ({ status={status} description={description} emphasis={emphasis} + progressMessage={progressMessage} + progressPercent={progressPercent} /> = ({ @@ -190,6 +192,8 @@ export const ToolInfo: React.FC = ({ description, status: coreStatus, emphasis, + progressMessage, + progressPercent, }) => { const status = mapCoreStatusToDisplayStatus(coreStatus); const nameColor = React.useMemo(() => { @@ -210,6 +214,24 @@ export const ToolInfo: React.FC = ({ // Hide description for completed Ask User tools (the result display speaks for itself) const isCompletedAskUser = isCompletedAskUserTool(name, status); + let displayDescription = description; + if (status === ToolCallStatus.Executing) { + const parts: string[] = []; + if (progressMessage) { + parts.push(progressMessage); + } + if (progressPercent !== undefined) { + parts.push(`${Math.round(progressPercent)}%`); + } + + if (parts.length > 0) { + const progressInfo = parts.join(' - '); + displayDescription = description + ? `${description} (${progressInfo})` + : progressInfo; + } + } + return ( @@ -219,7 +241,7 @@ export const ToolInfo: React.FC = ({ {!isCompletedAskUser && ( <> {' '} - {description} + {displayDescription} )} diff --git a/packages/cli/src/ui/hooks/toolMapping.ts b/packages/cli/src/ui/hooks/toolMapping.ts index bd8718b9bd..ded17f29a9 100644 --- a/packages/cli/src/ui/hooks/toolMapping.ts +++ b/packages/cli/src/ui/hooks/toolMapping.ts @@ -59,6 +59,8 @@ export function mapToDisplay( let outputFile: string | undefined = undefined; let ptyId: number | undefined = undefined; let correlationId: string | undefined = undefined; + let progressMessage: string | undefined = undefined; + let progressPercent: number | undefined = undefined; switch (call.status) { case CoreToolCallStatus.Success: @@ -77,6 +79,8 @@ export function mapToDisplay( case CoreToolCallStatus.Executing: resultDisplay = call.liveOutput; ptyId = call.pid; + progressMessage = call.progressMessage; + progressPercent = call.progressPercent; break; case CoreToolCallStatus.Scheduled: case CoreToolCallStatus.Validating: @@ -100,6 +104,8 @@ export function mapToDisplay( outputFile, ptyId, correlationId, + progressMessage, + progressPercent, approvalMode: call.approvalMode, }; }); diff --git a/packages/cli/src/ui/hooks/useToolScheduler.ts b/packages/cli/src/ui/hooks/useToolScheduler.ts index 89bee14342..56b1622468 100644 --- a/packages/cli/src/ui/hooks/useToolScheduler.ts +++ b/packages/cli/src/ui/hooks/useToolScheduler.ts @@ -104,6 +104,8 @@ export function useToolScheduler( [config, messageBus], ); + useEffect(() => () => scheduler.dispose(), [scheduler]); + const internalAdaptToolCalls = useCallback( (coreCalls: ToolCall[], prevTracked: TrackedToolCall[]) => adaptToolCalls(coreCalls, prevTracked), diff --git a/packages/cli/src/ui/types.ts b/packages/cli/src/ui/types.ts index 290ab63417..7cb06fbd15 100644 --- a/packages/cli/src/ui/types.ts +++ b/packages/cli/src/ui/types.ts @@ -108,6 +108,8 @@ export interface IndividualToolCallDisplay { outputFile?: string; correlationId?: string; approvalMode?: ApprovalMode; + progressMessage?: string; + progressPercent?: number; } export interface CompressionProps { diff --git a/packages/core/src/scheduler/scheduler.test.ts b/packages/core/src/scheduler/scheduler.test.ts index ad2d094b4e..f909a11d23 100644 --- a/packages/core/src/scheduler/scheduler.test.ts +++ b/packages/core/src/scheduler/scheduler.test.ts @@ -85,6 +85,7 @@ import { getToolCallContext, type ToolCallContext, } from '../utils/toolCallContext.js'; +import { coreEvents, CoreEvent } from '../utils/events.js'; describe('Scheduler (Orchestrator)', () => { let scheduler: Scheduler; @@ -1242,4 +1243,30 @@ describe('Scheduler (Orchestrator)', () => { expect(capturedContext!.parentCallId).toBe(parentCallId); }); }); + + describe('Cleanup', () => { + it('should unregister McpProgress listener on dispose()', () => { + const onSpy = vi.spyOn(coreEvents, 'on'); + const offSpy = vi.spyOn(coreEvents, 'off'); + + const s = new Scheduler({ + config: mockConfig, + messageBus: mockMessageBus, + getPreferredEditor, + schedulerId: 'cleanup-test', + }); + + expect(onSpy).toHaveBeenCalledWith( + CoreEvent.McpProgress, + expect.any(Function), + ); + + s.dispose(); + + expect(offSpy).toHaveBeenCalledWith( + CoreEvent.McpProgress, + expect.any(Function), + ); + }); + }); }); diff --git a/packages/core/src/scheduler/scheduler.ts b/packages/core/src/scheduler/scheduler.ts index b177fe0318..d71381eb33 100644 --- a/packages/core/src/scheduler/scheduler.ts +++ b/packages/core/src/scheduler/scheduler.ts @@ -39,6 +39,11 @@ import { type ToolConfirmationRequest, } from '../confirmation-bus/types.js'; import { runWithToolCallContext } from '../utils/toolCallContext.js'; +import { + coreEvents, + CoreEvent, + type McpProgressPayload, +} from '../utils/events.js'; interface SchedulerQueueItem { requests: ToolCallRequestInfo[]; @@ -115,8 +120,25 @@ export class Scheduler { this.modifier = new ToolModificationHandler(); this.setupMessageBusListener(this.messageBus); + + coreEvents.on(CoreEvent.McpProgress, this.handleMcpProgress); } + dispose(): void { + coreEvents.off(CoreEvent.McpProgress, this.handleMcpProgress); + } + + private readonly handleMcpProgress = (payload: McpProgressPayload) => { + const callId = payload.callId; + this.state.updateStatus(callId, CoreToolCallStatus.Executing, { + progressMessage: payload.message, + progressPercent: + payload.total && payload.total > 0 + ? (payload.progress / payload.total) * 100 + : undefined, + }); + }; + private setupMessageBusListener(messageBus: MessageBus): void { if (Scheduler.subscribedMessageBuses.has(messageBus)) { return; diff --git a/packages/core/src/scheduler/state-manager.test.ts b/packages/core/src/scheduler/state-manager.test.ts index 758ff354c0..6d25841b2e 100644 --- a/packages/core/src/scheduler/state-manager.test.ts +++ b/packages/core/src/scheduler/state-manager.test.ts @@ -495,6 +495,38 @@ describe('SchedulerStateManager', () => { expect(active.liveOutput).toBe('chunk 2'); expect(active.pid).toBe(1234); }); + + it('should update progressMessage and progressPercent during executing updates', () => { + const call = createValidatingCall(); + stateManager.enqueue([call]); + stateManager.dequeue(); + + // Update with progress + stateManager.updateStatus( + call.request.callId, + CoreToolCallStatus.Executing, + { + progressMessage: 'Starting...', + progressPercent: 10, + }, + ); + let active = stateManager.firstActiveCall as ExecutingToolCall; + expect(active.progressMessage).toBe('Starting...'); + expect(active.progressPercent).toBe(10); + + // Update progress further + stateManager.updateStatus( + call.request.callId, + CoreToolCallStatus.Executing, + { + progressMessage: 'Halfway!', + progressPercent: 50, + }, + ); + active = stateManager.firstActiveCall as ExecutingToolCall; + expect(active.progressMessage).toBe('Halfway!'); + expect(active.progressPercent).toBe(50); + }); }); describe('Argument Updates', () => { diff --git a/packages/core/src/scheduler/state-manager.ts b/packages/core/src/scheduler/state-manager.ts index 6a473ad47c..b282c3eb78 100644 --- a/packages/core/src/scheduler/state-manager.ts +++ b/packages/core/src/scheduler/state-manager.ts @@ -517,6 +517,12 @@ export class SchedulerStateManager { execData?.liveOutput ?? ('liveOutput' in call ? call.liveOutput : undefined); const pid = execData?.pid ?? ('pid' in call ? call.pid : undefined); + const progressMessage = + execData?.progressMessage ?? + ('progressMessage' in call ? call.progressMessage : undefined); + const progressPercent = + execData?.progressPercent ?? + ('progressPercent' in call ? call.progressPercent : undefined); return { request: call.request, @@ -527,6 +533,8 @@ export class SchedulerStateManager { invocation: call.invocation, liveOutput, pid, + progressMessage, + progressPercent, schedulerId: call.schedulerId, approvalMode: call.approvalMode, }; diff --git a/packages/core/src/scheduler/types.ts b/packages/core/src/scheduler/types.ts index b09c42fe51..7da611f23a 100644 --- a/packages/core/src/scheduler/types.ts +++ b/packages/core/src/scheduler/types.ts @@ -109,6 +109,8 @@ export type ExecutingToolCall = { tool: AnyDeclarativeTool; invocation: AnyToolInvocation; liveOutput?: string | AnsiOutput; + progressMessage?: string; + progressPercent?: number; startTime?: number; outcome?: ToolConfirmationOutcome; pid?: number; diff --git a/packages/core/src/tools/mcp-client.test.ts b/packages/core/src/tools/mcp-client.test.ts index 3f289f1732..19430c2f9a 100644 --- a/packages/core/src/tools/mcp-client.test.ts +++ b/packages/core/src/tools/mcp-client.test.ts @@ -18,7 +18,11 @@ import { MCPOAuthProvider } from '../mcp/oauth-provider.js'; import { MCPOAuthTokenStorage } from '../mcp/oauth-token-storage.js'; import { OAuthUtils } from '../mcp/oauth-utils.js'; import type { PromptRegistry } from '../prompts/prompt-registry.js'; -import { ToolListChangedNotificationSchema } from '@modelcontextprotocol/sdk/types.js'; +import { + PromptListChangedNotificationSchema, + ResourceListChangedNotificationSchema, + ToolListChangedNotificationSchema, +} from '@modelcontextprotocol/sdk/types.js'; import { ApprovalMode, PolicyDecision } from '../policy/types.js'; import { WorkspaceContext } from '../utils/workspaceContext.js'; @@ -140,7 +144,7 @@ describe('mcp-client', () => { await client.discover({} as Config); expect(mockedClient.listTools).toHaveBeenCalledWith( {}, - { timeout: 600000 }, + expect.objectContaining({ timeout: 600000, progressReporter: client }), ); }); @@ -710,8 +714,10 @@ describe('mcp-client', () => { getStatus: vi.fn(), registerCapabilities: vi.fn(), setRequestHandler: vi.fn(), - setNotificationHandler: vi.fn((_, handler) => { - resourceListHandler = handler; + setNotificationHandler: vi.fn((schema, handler) => { + if (schema === ResourceListChangedNotificationSchema) { + resourceListHandler = handler; + } }), getServerCapabilities: vi .fn() @@ -772,7 +778,7 @@ describe('mcp-client', () => { await client.connect(); await client.discover({} as Config); - expect(mockedClient.setNotificationHandler).toHaveBeenCalledOnce(); + expect(mockedClient.setNotificationHandler).toHaveBeenCalledTimes(2); expect(resourceListHandler).toBeDefined(); await resourceListHandler?.({ @@ -802,8 +808,10 @@ describe('mcp-client', () => { getStatus: vi.fn(), registerCapabilities: vi.fn(), setRequestHandler: vi.fn(), - setNotificationHandler: vi.fn((_, handler) => { - promptListHandler = handler; + setNotificationHandler: vi.fn((schema, handler) => { + if (schema === PromptListChangedNotificationSchema) { + promptListHandler = handler; + } }), getServerCapabilities: vi .fn() @@ -854,7 +862,7 @@ describe('mcp-client', () => { await client.connect(); await client.discover({ sanitizationConfig: EMPTY_CONFIG } as Config); - expect(mockedClient.setNotificationHandler).toHaveBeenCalledOnce(); + expect(mockedClient.setNotificationHandler).toHaveBeenCalledTimes(2); expect(promptListHandler).toBeDefined(); await promptListHandler?.({ @@ -1023,7 +1031,7 @@ describe('mcp-client', () => { await client.connect(); - expect(mockedClient.setNotificationHandler).not.toHaveBeenCalled(); + expect(mockedClient.setNotificationHandler).toHaveBeenCalledOnce(); }); it('should refresh tools and notify manager when notification is received', async () => { diff --git a/packages/core/src/tools/mcp-client.ts b/packages/core/src/tools/mcp-client.ts index 7902d8953a..19ea7d5054 100644 --- a/packages/core/src/tools/mcp-client.ts +++ b/packages/core/src/tools/mcp-client.ts @@ -30,6 +30,7 @@ import { ResourceListChangedNotificationSchema, ToolListChangedNotificationSchema, PromptListChangedNotificationSchema, + ProgressNotificationSchema, type Tool as McpTool, } from '@modelcontextprotocol/sdk/types.js'; import { ApprovalMode, PolicyDecision } from '../policy/types.js'; @@ -44,6 +45,7 @@ import { XcodeMcpBridgeFixTransport } from './xcode-mcp-fix-transport.js'; import type { CallableTool, FunctionCall, Part, Tool } from '@google/genai'; import { basename } from 'node:path'; import { pathToFileURL } from 'node:url'; +import { randomUUID } from 'node:crypto'; import type { McpAuthProvider } from '../mcp/auth-provider.js'; import { MCPOAuthProvider } from '../mcp/oauth-provider.js'; import { MCPOAuthTokenStorage } from '../mcp/oauth-token-storage.js'; @@ -58,6 +60,7 @@ import type { Unsubscribe, WorkspaceContext, } from '../utils/workspaceContext.js'; +import { getToolCallContext } from '../utils/toolCallContext.js'; import type { ToolRegistry } from './tool-registry.js'; import { debugLogger } from '../utils/debugLogger.js'; import { type MessageBus } from '../confirmation-bus/message-bus.js'; @@ -105,13 +108,21 @@ export enum MCPDiscoveryState { COMPLETED = 'completed', } +/** + * Interface for reporting progress from MCP tool calls. + */ +export interface McpProgressReporter { + registerProgressToken(token: string | number, callId: string): void; + unregisterProgressToken(token: string | number): void; +} + /** * A client for a single MCP server. * * This class is responsible for connecting to, discovering tools from, and * managing the state of a single MCP server. */ -export class McpClient { +export class McpClient implements McpProgressReporter { private client: Client | undefined; private transport: Transport | undefined; private status: MCPServerStatus = MCPServerStatus.DISCONNECTED; @@ -122,6 +133,12 @@ export class McpClient { private isRefreshingPrompts: boolean = false; private pendingPromptRefresh: boolean = false; + /** + * Map of progress tokens to tool call IDs. + * This allows us to route progress notifications to the correct tool call. + */ + private readonly progressTokenToCallId = new Map(); + constructor( private readonly serverName: string, private readonly serverConfig: MCPServerConfig, @@ -254,8 +271,11 @@ export class McpClient { this.client!, cliConfig, this.toolRegistry.getMessageBus(), - options ?? { - timeout: this.serverConfig.timeout ?? MCP_DEFAULT_TIMEOUT_MSEC, + { + ...(options ?? { + timeout: this.serverConfig.timeout ?? MCP_DEFAULT_TIMEOUT_MSEC, + }), + progressReporter: this, }, ); } @@ -349,6 +369,25 @@ export class McpClient { }, ); } + + this.client.setNotificationHandler( + ProgressNotificationSchema, + (notification) => { + const { progressToken, progress, total, message } = notification.params; + const callId = this.progressTokenToCallId.get(progressToken); + + if (callId) { + coreEvents.emitMcpProgress({ + serverName: this.serverName, + callId, + progressToken, + progress, + total, + message, + }); + } + }, + ); } /** @@ -409,6 +448,20 @@ export class McpClient { } } + /** + * Registers a progress token for a tool call. + */ + registerProgressToken(token: string | number, callId: string): void { + this.progressTokenToCallId.set(token, callId); + } + + /** + * Unregisters a progress token. + */ + unregisterProgressToken(token: string | number): void { + this.progressTokenToCallId.delete(token); + } + /** * Refreshes prompts for this server by re-querying the MCP `prompts/list` endpoint. */ @@ -994,7 +1047,11 @@ export async function discoverTools( mcpClient: Client, cliConfig: Config, messageBus: MessageBus, - options?: { timeout?: number; signal?: AbortSignal }, + options?: { + timeout?: number; + signal?: AbortSignal; + progressReporter?: McpProgressReporter; + }, ): Promise { try { // Only request tools if the server supports them. @@ -1012,6 +1069,7 @@ export async function discoverTools( mcpClient, toolDef, mcpServerConfig.timeout ?? MCP_DEFAULT_TIMEOUT_MSEC, + options?.progressReporter, ); // Extract readOnlyHint from annotations @@ -1078,6 +1136,7 @@ class McpCallableTool implements CallableTool { private readonly client: Client, private readonly toolDef: McpTool, private readonly timeout: number, + private readonly progressReporter?: McpProgressReporter, ) {} async tool(): Promise { @@ -1099,12 +1158,22 @@ class McpCallableTool implements CallableTool { } const call = functionCalls[0]; + const progressToken = randomUUID(); + const context = getToolCallContext(); + if (context && this.progressReporter) { + this.progressReporter.registerProgressToken( + progressToken, + context.callId, + ); + } + try { const result = await this.client.callTool( { name: call.name!, // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion arguments: call.args as Record, + _meta: { progressToken }, }, undefined, { timeout: this.timeout }, @@ -1133,6 +1202,10 @@ class McpCallableTool implements CallableTool { }, }, ]; + } finally { + if (this.progressReporter) { + this.progressReporter.unregisterProgressToken(progressToken); + } } } } diff --git a/packages/core/src/utils/events.ts b/packages/core/src/utils/events.ts index 014c2eec7a..1495ba63b5 100644 --- a/packages/core/src/utils/events.ts +++ b/packages/core/src/utils/events.ts @@ -124,6 +124,18 @@ export interface ConsentRequestPayload { onConfirm: (confirmed: boolean) => void; } +/** + * Payload for the 'mcp-progress' event. + */ +export interface McpProgressPayload { + serverName: string; + callId: string; + progressToken: string | number; + progress: number; + total?: number; + message?: string; +} + /** * Payload for the 'agents-discovered' event. */ @@ -167,6 +179,7 @@ export enum CoreEvent { AdminSettingsChanged = 'admin-settings-changed', RetryAttempt = 'retry-attempt', ConsentRequest = 'consent-request', + McpProgress = 'mcp-progress', AgentsDiscovered = 'agents-discovered', RequestEditorSelection = 'request-editor-selection', EditorSelected = 'editor-selected', @@ -200,6 +213,7 @@ export interface CoreEvents extends ExtensionEvents { [CoreEvent.AdminSettingsChanged]: never[]; [CoreEvent.RetryAttempt]: [RetryAttemptPayload]; [CoreEvent.ConsentRequest]: [ConsentRequestPayload]; + [CoreEvent.McpProgress]: [McpProgressPayload]; [CoreEvent.AgentsDiscovered]: [AgentsDiscoveredPayload]; [CoreEvent.RequestEditorSelection]: never[]; [CoreEvent.EditorSelected]: [EditorSelectedPayload]; @@ -335,6 +349,13 @@ export class CoreEventEmitter extends EventEmitter { this._emitOrQueue(CoreEvent.ConsentRequest, payload); } + /** + * Notifies subscribers that progress has been made on an MCP tool call. + */ + emitMcpProgress(payload: McpProgressPayload): void { + this.emit(CoreEvent.McpProgress, payload); + } + /** * Notifies subscribers that new unacknowledged agents have been discovered. */