fix(cli): allow sub-agent confirmation requests in UI while preventing background flicker (#20722)

This commit is contained in:
Abhi
2026-02-28 21:39:25 -05:00
committed by GitHub
parent 0063581e47
commit 703759cfae
2 changed files with 168 additions and 25 deletions

View File

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

View File

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