diff --git a/packages/cli/src/ui/hooks/useToolScheduler.test.ts b/packages/cli/src/ui/hooks/useToolScheduler.test.ts index 764e5f641f..0d010f25fa 100644 --- a/packages/cli/src/ui/hooks/useToolScheduler.test.ts +++ b/packages/cli/src/ui/hooks/useToolScheduler.test.ts @@ -20,6 +20,7 @@ import { type AnyToolInvocation, ROOT_SCHEDULER_ID, CoreToolCallStatus, + type WaitingToolCall, } from '@google/gemini-cli-core'; import { createMockMessageBus } from '@google/gemini-cli-core/src/test-utils/mock-message-bus.js'; @@ -32,6 +33,7 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => { Scheduler: vi.fn().mockImplementation(() => ({ schedule: vi.fn().mockResolvedValue([]), cancelAll: vi.fn(), + dispose: vi.fn(), })), }; }); @@ -341,7 +343,9 @@ describe('useToolScheduler', () => { const callSub = { ...callRoot, request: { ...callRoot.request, callId: 'call-sub' }, + status: CoreToolCallStatus.AwaitingApproval as const, // Must be awaiting approval to be tracked schedulerId: 'subagent-1', + confirmationDetails: { type: 'info', title: 'Confirm', prompt: 'Yes?' }, }; // 1. Populate state with multiple schedulers @@ -360,9 +364,13 @@ describe('useToolScheduler', () => { }); const [toolCalls] = result.current; - expect(toolCalls).toHaveLength(1); - expect(toolCalls[0].request.callId).toBe('call-root'); - expect(toolCalls[0].schedulerId).toBe(ROOT_SCHEDULER_ID); + expect(toolCalls).toHaveLength(2); + expect( + toolCalls.find((t) => t.request.callId === 'call-root'), + ).toBeDefined(); + expect( + toolCalls.find((t) => t.request.callId === 'call-sub'), + ).toBeDefined(); // 2. Call setToolCallsForDisplay (e.g., simulate a manual update or clear) act(() => { @@ -374,12 +382,11 @@ describe('useToolScheduler', () => { // 3. Verify that tools are still present and maintain their scheduler IDs const [toolCalls2] = result.current; - expect(toolCalls2).toHaveLength(1); - expect(toolCalls2[0].responseSubmittedToGemini).toBe(true); - expect(toolCalls2[0].schedulerId).toBe(ROOT_SCHEDULER_ID); + expect(toolCalls2).toHaveLength(2); + expect(toolCalls2.every((t) => t.responseSubmittedToGemini)).toBe(true); }); - it('ignores TOOL_CALLS_UPDATE from non-root schedulers', () => { + it('ignores TOOL_CALLS_UPDATE from non-root schedulers when no tools await approval', () => { const { result } = renderHook(() => useToolScheduler( vi.fn().mockResolvedValue(undefined), @@ -410,8 +417,125 @@ describe('useToolScheduler', () => { } as ToolCallsUpdateMessage); }); + expect(result.current[0]).toHaveLength(0); + }); + + it('allows TOOL_CALLS_UPDATE from non-root schedulers when tools are awaiting approval', () => { + const { result } = renderHook(() => + useToolScheduler( + vi.fn().mockResolvedValue(undefined), + mockConfig, + () => undefined, + ), + ); + + const subagentCall = { + status: CoreToolCallStatus.AwaitingApproval as const, + request: { + callId: 'call-sub', + name: 'test', + args: {}, + isClientInitiated: false, + prompt_id: 'p1', + }, + tool: createMockTool(), + invocation: createMockInvocation(), + schedulerId: 'subagent-1', + confirmationDetails: { type: 'info', title: 'Confirm', prompt: 'Yes?' }, + } as WaitingToolCall; + + act(() => { + void mockMessageBus.publish({ + type: MessageBusType.TOOL_CALLS_UPDATE, + toolCalls: [subagentCall], + schedulerId: 'subagent-1', + } as ToolCallsUpdateMessage); + }); + const [toolCalls] = result.current; - expect(toolCalls).toHaveLength(0); + expect(toolCalls).toHaveLength(1); + expect(toolCalls[0].request.callId).toBe('call-sub'); + expect(toolCalls[0].status).toBe(CoreToolCallStatus.AwaitingApproval); + }); + + it('preserves subagent tools in the UI after they have been approved', () => { + const { result } = renderHook(() => + useToolScheduler( + vi.fn().mockResolvedValue(undefined), + mockConfig, + () => undefined, + ), + ); + + const subagentCall = { + status: CoreToolCallStatus.AwaitingApproval as const, + request: { + callId: 'call-sub', + name: 'test', + args: {}, + isClientInitiated: false, + prompt_id: 'p1', + }, + tool: createMockTool(), + invocation: createMockInvocation(), + schedulerId: 'subagent-1', + confirmationDetails: { type: 'info', title: 'Confirm', prompt: 'Yes?' }, + } as WaitingToolCall; + + // 1. Initial approval request + act(() => { + void mockMessageBus.publish({ + type: MessageBusType.TOOL_CALLS_UPDATE, + toolCalls: [subagentCall], + schedulerId: 'subagent-1', + } as ToolCallsUpdateMessage); + }); + + expect(result.current[0]).toHaveLength(1); + + // 2. Approved and executing + const approvedCall = { + ...subagentCall, + status: CoreToolCallStatus.Executing as const, + } as unknown as ExecutingToolCall; + + act(() => { + void mockMessageBus.publish({ + type: MessageBusType.TOOL_CALLS_UPDATE, + toolCalls: [approvedCall], + schedulerId: 'subagent-1', + } as ToolCallsUpdateMessage); + }); + + expect(result.current[0]).toHaveLength(1); + expect(result.current[0][0].status).toBe(CoreToolCallStatus.Executing); + + // 3. New turn with a background tool (should NOT be shown) + const backgroundTool = { + status: CoreToolCallStatus.Executing as const, + request: { + callId: 'call-background', + name: 'read_file', + args: {}, + isClientInitiated: false, + prompt_id: 'p1', + }, + tool: createMockTool(), + invocation: createMockInvocation(), + schedulerId: 'subagent-1', + } as ExecutingToolCall; + + act(() => { + void mockMessageBus.publish({ + type: MessageBusType.TOOL_CALLS_UPDATE, + toolCalls: [backgroundTool], + schedulerId: 'subagent-1', + } as ToolCallsUpdateMessage); + }); + + // The subagent list should now be empty because the previously approved tool + // is gone from the current list, and the new tool doesn't need approval. + expect(result.current[0]).toHaveLength(0); }); it('adapts success/error status to executing when a tail call is present', () => { diff --git a/packages/cli/src/ui/hooks/useToolScheduler.ts b/packages/cli/src/ui/hooks/useToolScheduler.ts index 0fe732b26e..496143e590 100644 --- a/packages/cli/src/ui/hooks/useToolScheduler.ts +++ b/packages/cli/src/ui/hooks/useToolScheduler.ts @@ -115,11 +115,42 @@ export function useToolScheduler( useEffect(() => { const handler = (event: ToolCallsUpdateMessage) => { - // Only process updates for the root scheduler. - // Subagent internal tools should not be displayed in the main tool list. - if (event.schedulerId !== ROOT_SCHEDULER_ID) { - return; - } + const isRoot = event.schedulerId === ROOT_SCHEDULER_ID; + + setToolCallsMap((prev) => { + const prevCalls = prev[event.schedulerId] ?? []; + const prevCallIds = new Set(prevCalls.map((tc) => tc.request.callId)); + + // For non-root schedulers, we only show tool calls that: + // 1. Are currently awaiting approval. + // 2. Were previously shown (e.g., they are now executing or completed). + // This prevents "thinking" tools (reads/searches) from flickering in the UI + // unless they specifically required user interaction. + const filteredToolCalls = isRoot + ? event.toolCalls + : event.toolCalls.filter( + (tc) => + tc.status === CoreToolCallStatus.AwaitingApproval || + prevCallIds.has(tc.request.callId), + ); + + // If this is a subagent and we have no tools to show and weren't showing any, + // we can skip the update entirely to avoid unnecessary re-renders. + if ( + !isRoot && + filteredToolCalls.length === 0 && + prevCalls.length === 0 + ) { + return prev; + } + + const adapted = internalAdaptToolCalls(filteredToolCalls, prevCalls); + + return { + ...prev, + [event.schedulerId]: adapted, + }; + }); // Update output timer for UI spinners (Side Effect) const hasExecuting = event.toolCalls.some( @@ -134,18 +165,6 @@ export function useToolScheduler( if (hasExecuting) { setLastToolOutputTime(Date.now()); } - - setToolCallsMap((prev) => { - const adapted = internalAdaptToolCalls( - event.toolCalls, - prev[event.schedulerId] ?? [], - ); - - return { - ...prev, - [event.schedulerId]: adapted, - }; - }); }; messageBus.subscribe(MessageBusType.TOOL_CALLS_UPDATE, handler);