feat(core): add support for MCP progress updates (#19046)

This commit is contained in:
N. Taylor Mullen
2026-02-18 12:46:12 -08:00
committed by GitHub
parent 1cf05b0375
commit 14415316c0
14 changed files with 270 additions and 14 deletions

View File

@@ -320,4 +320,31 @@ describe('<ToolMessage />', () => {
);
expect(lastFrame()).toMatchSnapshot();
});
it('renders progress information appended to description for executing tools', () => {
const { lastFrame } = renderWithContext(
<ToolMessage
{...baseProps}
status={CoreToolCallStatus.Executing}
progressMessage="Working on it..."
progressPercent={42}
/>,
StreamingState.Responding,
);
expect(lastFrame()).toContain(
'A tool for testing (Working on it... - 42%)',
);
});
it('renders only percentage when progressMessage is missing', () => {
const { lastFrame } = renderWithContext(
<ToolMessage
{...baseProps}
status={CoreToolCallStatus.Executing}
progressPercent={75}
/>,
StreamingState.Responding,
);
expect(lastFrame()).toContain('A tool for testing (75%)');
});
});

View File

@@ -55,6 +55,8 @@ export const ToolMessage: React.FC<ToolMessageProps> = ({
embeddedShellFocused,
ptyId,
config,
progressMessage,
progressPercent,
}) => {
const isThisShellFocused = checkIsShellFocused(
name,
@@ -89,6 +91,8 @@ export const ToolMessage: React.FC<ToolMessageProps> = ({
status={status}
description={description}
emphasis={emphasis}
progressMessage={progressMessage}
progressPercent={progressPercent}
/>
<FocusHint
shouldShowFocusHint={shouldShowFocusHint}

View File

@@ -183,6 +183,8 @@ type ToolInfoProps = {
description: string;
status: CoreToolCallStatus;
emphasis: TextEmphasis;
progressMessage?: string;
progressPercent?: number;
};
export const ToolInfo: React.FC<ToolInfoProps> = ({
@@ -190,6 +192,8 @@ export const ToolInfo: React.FC<ToolInfoProps> = ({
description,
status: coreStatus,
emphasis,
progressMessage,
progressPercent,
}) => {
const status = mapCoreStatusToDisplayStatus(coreStatus);
const nameColor = React.useMemo<string>(() => {
@@ -210,6 +214,24 @@ export const ToolInfo: React.FC<ToolInfoProps> = ({
// 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 (
<Box overflow="hidden" height={1} flexGrow={1} flexShrink={1}>
<Text strikethrough={status === ToolCallStatus.Canceled} wrap="truncate">
@@ -219,7 +241,7 @@ export const ToolInfo: React.FC<ToolInfoProps> = ({
{!isCompletedAskUser && (
<>
{' '}
<Text color={theme.text.secondary}>{description}</Text>
<Text color={theme.text.secondary}>{displayDescription}</Text>
</>
)}
</Text>

View File

@@ -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,
};
});

View File

@@ -104,6 +104,8 @@ export function useToolScheduler(
[config, messageBus],
);
useEffect(() => () => scheduler.dispose(), [scheduler]);
const internalAdaptToolCalls = useCallback(
(coreCalls: ToolCall[], prevTracked: TrackedToolCall[]) =>
adaptToolCalls(coreCalls, prevTracked),

View File

@@ -108,6 +108,8 @@ export interface IndividualToolCallDisplay {
outputFile?: string;
correlationId?: string;
approvalMode?: ApprovalMode;
progressMessage?: string;
progressPercent?: number;
}
export interface CompressionProps {

View File

@@ -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),
);
});
});
});

View File

@@ -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;

View File

@@ -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', () => {

View File

@@ -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,
};

View File

@@ -109,6 +109,8 @@ export type ExecutingToolCall = {
tool: AnyDeclarativeTool;
invocation: AnyToolInvocation;
liveOutput?: string | AnsiOutput;
progressMessage?: string;
progressPercent?: number;
startTime?: number;
outcome?: ToolConfirmationOutcome;
pid?: number;

View File

@@ -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 () => {

View File

@@ -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<string | number, string>();
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<DiscoveredMCPTool[]> {
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<Tool> {
@@ -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<string, unknown>,
_meta: { progressToken },
},
undefined,
{ timeout: this.timeout },
@@ -1133,6 +1202,10 @@ class McpCallableTool implements CallableTool {
},
},
];
} finally {
if (this.progressReporter) {
this.progressReporter.unregisterProgressToken(progressToken);
}
}
}
}

View File

@@ -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<CoreEvents> {
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.
*/