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
140 changes: 132 additions & 8 deletions packages/cli/src/ui/hooks/useToolScheduler.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';

Expand All @@ -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(),
})),
};
});
Expand Down Expand Up @@ -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
Expand All @@ -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(() => {
Expand All @@ -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),
Expand Down Expand Up @@ -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', () => {
Expand Down
53 changes: 36 additions & 17 deletions packages/cli/src/ui/hooks/useToolScheduler.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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);
Expand Down
Loading