Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 27 additions & 0 deletions packages/cli/src/ui/components/messages/ToolMessage.test.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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%)');
});
});
4 changes: 4 additions & 0 deletions packages/cli/src/ui/components/messages/ToolMessage.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,8 @@ export const ToolMessage: React.FC<ToolMessageProps> = ({
embeddedShellFocused,
ptyId,
config,
progressMessage,
progressPercent,
}) => {
const isThisShellFocused = checkIsShellFocused(
name,
Expand Down Expand Up @@ -89,6 +91,8 @@ export const ToolMessage: React.FC<ToolMessageProps> = ({
status={status}
description={description}
emphasis={emphasis}
progressMessage={progressMessage}
progressPercent={progressPercent}
/>
<FocusHint
shouldShowFocusHint={shouldShowFocusHint}
Expand Down
24 changes: 23 additions & 1 deletion packages/cli/src/ui/components/messages/ToolShared.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -183,13 +183,17 @@ type ToolInfoProps = {
description: string;
status: CoreToolCallStatus;
emphasis: TextEmphasis;
progressMessage?: string;
progressPercent?: number;
};

export const ToolInfo: React.FC<ToolInfoProps> = ({
name,
description,
status: coreStatus,
emphasis,
progressMessage,
progressPercent,
}) => {
const status = mapCoreStatusToDisplayStatus(coreStatus);
const nameColor = React.useMemo<string>(() => {
Expand All @@ -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">
Expand All @@ -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>
Expand Down
6 changes: 6 additions & 0 deletions packages/cli/src/ui/hooks/toolMapping.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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:
Expand All @@ -100,6 +104,8 @@ export function mapToDisplay(
outputFile,
ptyId,
correlationId,
progressMessage,
progressPercent,
approvalMode: call.approvalMode,
};
});
Expand Down
2 changes: 2 additions & 0 deletions packages/cli/src/ui/hooks/useToolScheduler.ts
Original file line number Diff line number Diff line change
Expand Up @@ -104,6 +104,8 @@ export function useToolScheduler(
[config, messageBus],
);

useEffect(() => () => scheduler.dispose(), [scheduler]);

const internalAdaptToolCalls = useCallback(
(coreCalls: ToolCall[], prevTracked: TrackedToolCall[]) =>
adaptToolCalls(coreCalls, prevTracked),
Expand Down
2 changes: 2 additions & 0 deletions packages/cli/src/ui/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -108,6 +108,8 @@ export interface IndividualToolCallDisplay {
outputFile?: string;
correlationId?: string;
approvalMode?: ApprovalMode;
progressMessage?: string;
progressPercent?: number;
}

export interface CompressionProps {
Expand Down
27 changes: 27 additions & 0 deletions packages/core/src/scheduler/scheduler.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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),
);
});
});
});
22 changes: 22 additions & 0 deletions packages/core/src/scheduler/scheduler.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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[];
Expand Down Expand Up @@ -115,8 +120,25 @@ export class Scheduler {
this.modifier = new ToolModificationHandler();

this.setupMessageBusListener(this.messageBus);

coreEvents.on(CoreEvent.McpProgress, this.handleMcpProgress);
Comment thread
NTaylorMullen marked this conversation as resolved.
}

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;
Expand Down
32 changes: 32 additions & 0 deletions packages/core/src/scheduler/state-manager.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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', () => {
Expand Down
8 changes: 8 additions & 0 deletions packages/core/src/scheduler/state-manager.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -527,6 +533,8 @@ export class SchedulerStateManager {
invocation: call.invocation,
liveOutput,
pid,
progressMessage,
progressPercent,
schedulerId: call.schedulerId,
approvalMode: call.approvalMode,
};
Expand Down
2 changes: 2 additions & 0 deletions packages/core/src/scheduler/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -109,6 +109,8 @@ export type ExecutingToolCall = {
tool: AnyDeclarativeTool;
invocation: AnyToolInvocation;
liveOutput?: string | AnsiOutput;
progressMessage?: string;
progressPercent?: number;
startTime?: number;
outcome?: ToolConfirmationOutcome;
pid?: number;
Expand Down
26 changes: 17 additions & 9 deletions packages/core/src/tools/mcp-client.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -140,7 +144,7 @@ describe('mcp-client', () => {
await client.discover({} as Config);
expect(mockedClient.listTools).toHaveBeenCalledWith(
{},
{ timeout: 600000 },
expect.objectContaining({ timeout: 600000, progressReporter: client }),
);
});

Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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?.({
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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?.({
Expand Down Expand Up @@ -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 () => {
Expand Down
Loading
Loading