feat(ux) Expandable (ctrl-O) and scrollable approvals in alternate buffer mode. (#17640)

This commit is contained in:
Jacob Richman
2026-01-27 16:06:24 -08:00
committed by GitHub
parent ff6547857e
commit d165b6d4e7
34 changed files with 1177 additions and 496 deletions

View File

@@ -50,8 +50,10 @@ export function mapCoreStatusToDisplayStatus(
*/
export function mapToDisplay(
toolOrTools: ToolCall[] | ToolCall,
options: { borderTop?: boolean; borderBottom?: boolean } = {},
): HistoryItemToolGroup {
const toolCalls = Array.isArray(toolOrTools) ? toolOrTools : [toolOrTools];
const { borderTop, borderBottom } = options;
const toolDisplays = toolCalls.map((call): IndividualToolCallDisplay => {
let description: string;
@@ -128,5 +130,7 @@ export function mapToDisplay(
return {
type: 'tool_group',
tools: toolDisplays,
borderTop,
borderBottom,
};
}

View File

@@ -8,7 +8,7 @@
import type { Mock, MockInstance } from 'vitest';
import { describe, it, expect, vi, beforeEach } from 'vitest';
import { act } from 'react';
import { renderHook } from '../../test-utils/render.js';
import { renderHookWithProviders } from '../../test-utils/render.js';
import { waitFor } from '../../test-utils/async.js';
import { useGeminiStream } from './useGeminiStream.js';
import { useKeypress } from './useKeypress.js';
@@ -56,7 +56,7 @@ const MockedGeminiClientClass = vi.hoisted(() =>
this.startChat = mockStartChat;
this.sendMessageStream = mockSendMessageStream;
this.addHistory = vi.fn();
this.getCurrentSequenceModel = vi.fn();
this.getCurrentSequenceModel = vi.fn().mockReturnValue('test-model');
this.getChat = vi.fn().mockReturnValue({
recordCompletedToolCalls: vi.fn(),
});
@@ -93,6 +93,8 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => {
ValidationRequiredError: MockValidationRequiredError,
parseAndFormatApiError: mockParseAndFormatApiError,
tokenLimit: vi.fn().mockReturnValue(100), // Mock tokenLimit
recordToolCallInteractions: vi.fn().mockResolvedValue(undefined),
getCodeAssistServer: vi.fn().mockReturnValue(undefined),
};
});
@@ -167,84 +169,92 @@ vi.mock('./useAlternateBuffer.js', () => ({
// --- Tests for useGeminiStream Hook ---
describe('useGeminiStream', () => {
let mockAddItem: Mock;
let mockConfig: Config;
let mockOnDebugMessage: Mock;
let mockHandleSlashCommand: Mock;
let mockAddItem = vi.fn();
let mockOnDebugMessage = vi.fn();
let mockHandleSlashCommand = vi.fn().mockResolvedValue(false);
let mockScheduleToolCalls: Mock;
let mockCancelAllToolCalls: Mock;
let mockMarkToolsAsSubmitted: Mock;
let handleAtCommandSpy: MockInstance;
beforeEach(() => {
vi.clearAllMocks(); // Clear mocks before each test
const emptyHistory: any[] = [];
let capturedOnComplete: any = null;
const mockGetPreferredEditor = vi.fn(() => 'vscode' as EditorType);
const mockOnAuthError = vi.fn();
const mockPerformMemoryRefresh = vi.fn(() => Promise.resolve());
const mockSetModelSwitchedFromQuotaError = vi.fn();
const mockOnCancelSubmit = vi.fn();
const mockSetShellInputFocused = vi.fn();
mockAddItem = vi.fn();
// Define the mock for getGeminiClient
const mockGetGeminiClient = vi.fn().mockImplementation(() => {
// MockedGeminiClientClass is defined in the module scope by the previous change.
// It will use the mockStartChat and mockSendMessageStream that are managed within beforeEach.
const clientInstance = new MockedGeminiClientClass(mockConfig);
return clientInstance;
});
const mockGetGeminiClient = vi.fn().mockImplementation(() => {
const clientInstance = new MockedGeminiClientClass(mockConfig);
return clientInstance;
});
const mockMcpClientManager = {
getDiscoveryState: vi.fn().mockReturnValue(MCPDiscoveryState.COMPLETED),
getMcpServerCount: vi.fn().mockReturnValue(0),
};
const mockMcpClientManager = {
getDiscoveryState: vi.fn().mockReturnValue(MCPDiscoveryState.COMPLETED),
getMcpServerCount: vi.fn().mockReturnValue(0),
};
const contentGeneratorConfig = {
const mockConfig: Config = {
apiKey: 'test-api-key',
model: 'gemini-pro',
sandbox: false,
targetDir: '/test/dir',
debugMode: false,
question: undefined,
coreTools: [],
toolDiscoveryCommand: undefined,
toolCallCommand: undefined,
mcpServerCommand: undefined,
mcpServers: undefined,
userAgent: 'test-agent',
userMemory: '',
geminiMdFileCount: 0,
alwaysSkipModificationConfirmation: false,
vertexai: false,
showMemoryUsage: false,
contextFileName: undefined,
storage: {
getProjectTempDir: vi.fn(() => '/test/temp'),
getProjectTempCheckpointsDir: vi.fn(() => '/test/temp/checkpoints'),
} as any,
getToolRegistry: vi.fn(
() => ({ getToolSchemaList: vi.fn(() => []) }) as any,
),
getProjectRoot: vi.fn(() => '/test/dir'),
getCheckpointingEnabled: vi.fn(() => false),
getGeminiClient: mockGetGeminiClient,
getMcpClientManager: () => mockMcpClientManager as any,
getApprovalMode: vi.fn(() => ApprovalMode.DEFAULT),
getUsageStatisticsEnabled: () => true,
getDebugMode: () => false,
addHistory: vi.fn(),
getSessionId: vi.fn(() => 'test-session-id'),
setQuotaErrorOccurred: vi.fn(),
getQuotaErrorOccurred: vi.fn(() => false),
getModel: vi.fn(() => 'gemini-2.5-pro'),
getContentGeneratorConfig: vi.fn(() => ({
model: 'test-model',
apiKey: 'test-key',
vertexai: false,
authType: AuthType.USE_GEMINI,
};
})),
getContentGenerator: vi.fn(),
isInteractive: () => false,
getExperiments: () => {},
isEventDrivenSchedulerEnabled: vi.fn(() => false),
getMaxSessionTurns: vi.fn(() => 100),
isJitContextEnabled: vi.fn(() => false),
getGlobalMemory: vi.fn(() => ''),
getUserMemory: vi.fn(() => ''),
getIdeMode: vi.fn(() => false),
getEnableHooks: vi.fn(() => false),
} as unknown as Config;
mockConfig = {
apiKey: 'test-api-key',
model: 'gemini-pro',
sandbox: false,
targetDir: '/test/dir',
debugMode: false,
question: undefined,
coreTools: [],
toolDiscoveryCommand: undefined,
toolCallCommand: undefined,
mcpServerCommand: undefined,
mcpServers: undefined,
userAgent: 'test-agent',
userMemory: '',
geminiMdFileCount: 0,
alwaysSkipModificationConfirmation: false,
vertexai: false,
showMemoryUsage: false,
contextFileName: undefined,
getToolRegistry: vi.fn(
() => ({ getToolSchemaList: vi.fn(() => []) }) as any,
),
getProjectRoot: vi.fn(() => '/test/dir'),
getCheckpointingEnabled: vi.fn(() => false),
getGeminiClient: mockGetGeminiClient,
getMcpClientManager: () => mockMcpClientManager as any,
getApprovalMode: () => ApprovalMode.DEFAULT,
getUsageStatisticsEnabled: () => true,
getDebugMode: () => false,
getWorkingDir: () => '/working/dir',
addHistory: vi.fn(),
getSessionId() {
return 'test-session-id';
},
setQuotaErrorOccurred: vi.fn(),
getQuotaErrorOccurred: vi.fn(() => false),
getModel: vi.fn(() => 'gemini-2.5-pro'),
getContentGenerator: vi.fn(),
getContentGeneratorConfig: vi
.fn()
.mockReturnValue(contentGeneratorConfig),
isInteractive: () => false,
getExperiments: () => {},
} as unknown as Config;
beforeEach(() => {
vi.clearAllMocks(); // Clear mocks before each test
mockAddItem = vi.fn();
mockOnDebugMessage = vi.fn();
mockHandleSlashCommand = vi.fn().mockResolvedValue(false);
@@ -253,6 +263,10 @@ describe('useGeminiStream', () => {
mockCancelAllToolCalls = vi.fn();
mockMarkToolsAsSubmitted = vi.fn();
// Reset properties of mockConfig if needed
(mockConfig.getCheckpointingEnabled as Mock).mockReturnValue(false);
(mockConfig.getApprovalMode as Mock).mockReturnValue(ApprovalMode.DEFAULT);
// Default mock for useReactToolScheduler to prevent toolCalls being undefined initially
mockUseToolScheduler.mockReturnValue([
[], // Default to empty array for toolCalls
@@ -289,10 +303,11 @@ describe('useGeminiStream', () => {
geminiClient?: any,
) => {
const client = geminiClient || mockConfig.getGeminiClient();
let lastToolCalls = initialToolCalls;
const initialProps = {
client,
history: [],
history: emptyHistory,
addItem: mockAddItem as unknown as UseHistoryManagerReturn['addItem'],
config: mockConfig,
onDebugMessage: mockOnDebugMessage,
@@ -304,31 +319,26 @@ describe('useGeminiStream', () => {
toolCalls: initialToolCalls,
};
const { result, rerender } = renderHook(
(props: typeof initialProps) => {
// This mock needs to be stateful. When setToolCallsForDisplay is called,
// it should trigger a rerender with the new state.
const mockSetToolCallsForDisplay = vi.fn((updater) => {
const newToolCalls =
typeof updater === 'function' ? updater(props.toolCalls) : updater;
rerender({ ...props, toolCalls: newToolCalls });
});
// Create a stateful mock for cancellation that updates the toolCalls state.
const statefulCancelAllToolCalls = vi.fn((...args) => {
// Call the original spy so `toHaveBeenCalled` checks still work.
mockUseToolScheduler.mockImplementation((onComplete) => {
capturedOnComplete = onComplete;
return [
lastToolCalls,
mockScheduleToolCalls,
mockMarkToolsAsSubmitted,
(updater: any) => {
lastToolCalls =
typeof updater === 'function' ? updater(lastToolCalls) : updater;
rerender({ ...initialProps, toolCalls: lastToolCalls });
},
(...args: any[]) => {
mockCancelAllToolCalls(...args);
const newToolCalls = props.toolCalls.map((tc) => {
// Only cancel tools that are in a cancellable state.
lastToolCalls = lastToolCalls.map((tc) => {
if (
tc.status === 'awaiting_approval' ||
tc.status === 'executing' ||
tc.status === 'scheduled' ||
tc.status === 'validating'
) {
// A real cancelled tool call has a response object.
// We need to simulate this to avoid type errors downstream.
return {
...tc,
status: 'cancelled',
@@ -337,23 +347,20 @@ describe('useGeminiStream', () => {
responseParts: [],
resultDisplay: 'Request cancelled.',
},
responseSubmittedToGemini: true, // Mark as "processed"
responseSubmittedToGemini: true,
} as any as TrackedCancelledToolCall;
}
return tc;
});
rerender({ ...props, toolCalls: newToolCalls });
});
rerender({ ...initialProps, toolCalls: lastToolCalls });
},
0,
];
});
mockUseToolScheduler.mockImplementation(() => [
props.toolCalls,
mockScheduleToolCalls,
mockMarkToolsAsSubmitted,
mockSetToolCallsForDisplay,
statefulCancelAllToolCalls, // Use the stateful mock
]);
return useGeminiStream(
const { result, rerender } = renderHookWithProviders(
(props: typeof initialProps) =>
useGeminiStream(
props.client,
props.history,
props.addItem,
@@ -362,17 +369,16 @@ describe('useGeminiStream', () => {
props.onDebugMessage,
props.handleSlashCommand,
props.shellModeActive,
() => 'vscode' as EditorType,
() => {},
() => Promise.resolve(),
mockGetPreferredEditor,
mockOnAuthError,
mockPerformMemoryRefresh,
false,
() => {},
() => {},
() => {},
mockSetModelSwitchedFromQuotaError,
mockOnCancelSubmit,
mockSetShellInputFocused,
80,
24,
);
},
),
{
initialProps,
},
@@ -454,7 +460,7 @@ describe('useGeminiStream', () => {
modelSwitched = false,
} = options;
return renderHook(() =>
return renderHookWithProviders(() =>
useGeminiStream(
new MockedGeminiClientClass(mockConfig),
[],
@@ -592,10 +598,17 @@ describe('useGeminiStream', () => {
mockUseToolScheduler.mockImplementation((onComplete) => {
capturedOnComplete = onComplete;
return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted, vi.fn()];
return [
[],
mockScheduleToolCalls,
mockMarkToolsAsSubmitted,
vi.fn(),
mockCancelAllToolCalls,
0,
];
});
renderHook(() =>
renderHookWithProviders(() =>
useGeminiStream(
new MockedGeminiClientClass(mockConfig),
[],
@@ -620,6 +633,8 @@ describe('useGeminiStream', () => {
// Trigger the onComplete callback with completed tools
await act(async () => {
if (capturedOnComplete) {
// Wait a tick for refs to be set up
await new Promise((resolve) => setTimeout(resolve, 0));
await capturedOnComplete(completedToolCalls);
}
});
@@ -674,10 +689,17 @@ describe('useGeminiStream', () => {
mockUseToolScheduler.mockImplementation((onComplete) => {
capturedOnComplete = onComplete;
return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted, vi.fn()];
return [
[],
mockScheduleToolCalls,
mockMarkToolsAsSubmitted,
vi.fn(),
mockCancelAllToolCalls,
0,
];
});
renderHook(() =>
renderHookWithProviders(() =>
useGeminiStream(
client,
[],
@@ -702,6 +724,8 @@ describe('useGeminiStream', () => {
// Trigger the onComplete callback with cancelled tools
await act(async () => {
if (capturedOnComplete) {
// Wait a tick for refs to be set up
await new Promise((resolve) => setTimeout(resolve, 0));
await capturedOnComplete(cancelledToolCalls);
}
});
@@ -746,48 +770,12 @@ describe('useGeminiStream', () => {
];
const client = new MockedGeminiClientClass(mockConfig);
// Capture the onComplete callback
let capturedOnComplete:
| ((completedTools: TrackedToolCall[]) => Promise<void>)
| null = null;
mockUseToolScheduler.mockImplementation((onComplete) => {
capturedOnComplete = onComplete;
return [
[],
mockScheduleToolCalls,
mockMarkToolsAsSubmitted,
vi.fn(),
mockCancelAllToolCalls,
];
});
const { result } = renderHook(() =>
useGeminiStream(
client,
[],
mockAddItem,
mockConfig,
mockLoadedSettings,
mockOnDebugMessage,
mockHandleSlashCommand,
false,
() => 'vscode' as EditorType,
() => {},
() => Promise.resolve(),
false,
() => {},
() => {},
() => {},
80,
24,
),
);
const { result } = renderTestHook([], client);
// Trigger the onComplete callback with STOP_EXECUTION tool
await act(async () => {
if (capturedOnComplete) {
await (capturedOnComplete as any)(stopExecutionToolCalls);
await capturedOnComplete(stopExecutionToolCalls);
}
});
@@ -877,10 +865,17 @@ describe('useGeminiStream', () => {
mockUseToolScheduler.mockImplementation((onComplete) => {
capturedOnComplete = onComplete;
return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted, vi.fn()];
return [
[],
mockScheduleToolCalls,
mockMarkToolsAsSubmitted,
vi.fn(),
mockCancelAllToolCalls,
0,
];
});
renderHook(() =>
renderHookWithProviders(() =>
useGeminiStream(
client,
[],
@@ -905,6 +900,8 @@ describe('useGeminiStream', () => {
// Trigger the onComplete callback with multiple cancelled tools
await act(async () => {
if (capturedOnComplete) {
// Wait a tick for refs to be set up
await new Promise((resolve) => setTimeout(resolve, 0));
await capturedOnComplete(allCancelledTools);
}
});
@@ -990,10 +987,12 @@ describe('useGeminiStream', () => {
mockScheduleToolCalls,
mockMarkToolsAsSubmitted,
vi.fn(), // setToolCallsForDisplay
mockCancelAllToolCalls,
0,
];
});
const { result, rerender } = renderHook(() =>
const { result, rerender } = renderHookWithProviders(() =>
useGeminiStream(
new MockedGeminiClientClass(mockConfig),
[],
@@ -1027,6 +1026,8 @@ describe('useGeminiStream', () => {
mockScheduleToolCalls,
mockMarkToolsAsSubmitted,
vi.fn(), // setToolCallsForDisplay
mockCancelAllToolCalls,
0,
];
});
@@ -1041,6 +1042,8 @@ describe('useGeminiStream', () => {
// 4. Trigger the onComplete callback to simulate tool completion
await act(async () => {
if (capturedOnComplete) {
// Wait a tick for refs to be set up
await new Promise((resolve) => setTimeout(resolve, 0));
await capturedOnComplete(completedToolCalls);
}
});
@@ -1124,7 +1127,7 @@ describe('useGeminiStream', () => {
})();
mockSendMessageStream.mockReturnValue(mockStream);
const { result } = renderHook(() =>
const { result } = renderHookWithProviders(() =>
useGeminiStream(
mockConfig.getGeminiClient(),
[],
@@ -1165,7 +1168,7 @@ describe('useGeminiStream', () => {
})();
mockSendMessageStream.mockReturnValue(mockStream);
const { result } = renderHook(() =>
const { result } = renderHookWithProviders(() =>
useGeminiStream(
mockConfig.getGeminiClient(),
[],
@@ -1559,7 +1562,7 @@ describe('useGeminiStream', () => {
});
it('should not call handleSlashCommand is shell mode is active', async () => {
const { result } = renderHook(() =>
const { result } = renderHookWithProviders(() =>
useGeminiStream(
new MockedGeminiClientClass(mockConfig),
[],
@@ -1629,10 +1632,17 @@ describe('useGeminiStream', () => {
mockUseToolScheduler.mockImplementation((onComplete) => {
capturedOnComplete = onComplete;
return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted, vi.fn()];
return [
[],
mockScheduleToolCalls,
mockMarkToolsAsSubmitted,
vi.fn(),
mockCancelAllToolCalls,
0,
];
});
renderHook(() =>
renderHookWithProviders(() =>
useGeminiStream(
new MockedGeminiClientClass(mockConfig),
[],
@@ -1657,6 +1667,8 @@ describe('useGeminiStream', () => {
// Trigger the onComplete callback with the completed save_memory tool
await act(async () => {
if (capturedOnComplete) {
// Wait a tick for refs to be set up
await new Promise((resolve) => setTimeout(resolve, 0));
await capturedOnComplete([completedToolCall]);
}
});
@@ -1689,7 +1701,7 @@ describe('useGeminiStream', () => {
getModel: vi.fn(() => 'gemini-2.5-pro'),
} as unknown as Config;
const { result } = renderHook(() =>
const { result } = renderHookWithProviders(() =>
useGeminiStream(
new MockedGeminiClientClass(testConfig),
[],
@@ -1990,7 +2002,7 @@ describe('useGeminiStream', () => {
})(),
);
const { result } = renderHook(() =>
const { result } = renderHookWithProviders(() =>
useGeminiStream(
new MockedGeminiClientClass(mockConfig),
[],
@@ -2095,7 +2107,7 @@ describe('useGeminiStream', () => {
})(),
);
const { result } = renderHook(() =>
const { result } = renderHookWithProviders(() =>
useGeminiStream(
new MockedGeminiClientClass(mockConfig),
[],
@@ -2243,6 +2255,8 @@ describe('useGeminiStream', () => {
startTime: Date.now(),
endTime: Date.now(),
}));
// Wait a tick for refs to be set up
await new Promise((resolve) => setTimeout(resolve, 0));
await capturedOnComplete(tools);
addItemOrder.push('scheduleToolCalls_END');
});
@@ -2264,7 +2278,7 @@ describe('useGeminiStream', () => {
];
});
const { result } = renderHook(() =>
const { result } = renderHookWithProviders(() =>
useGeminiStream(
new MockedGeminiClientClass(mockConfig),
[],
@@ -2330,7 +2344,7 @@ describe('useGeminiStream', () => {
shouldProceed: true,
});
const { result } = renderHook(() =>
const { result } = renderHookWithProviders(() =>
useGeminiStream(
mockConfig.getGeminiClient(),
[],
@@ -2489,7 +2503,7 @@ describe('useGeminiStream', () => {
})(),
);
const { result } = renderHook(() =>
const { result } = renderHookWithProviders(() =>
useGeminiStream(
new MockedGeminiClientClass(mockConfig),
[],
@@ -2565,11 +2579,13 @@ describe('useGeminiStream', () => {
mockUseToolScheduler.mockReturnValue([
[],
mockScheduleToolCalls,
mockCancelAllToolCalls,
mockMarkToolsAsSubmitted,
vi.fn(),
mockCancelAllToolCalls,
0,
]);
const { result, rerender } = renderHook(() =>
const { result, rerender } = renderHookWithProviders(() =>
useGeminiStream(
mockConfig.getGeminiClient(),
[],
@@ -2616,8 +2632,10 @@ describe('useGeminiStream', () => {
mockUseToolScheduler.mockReturnValue([
newToolCalls,
mockScheduleToolCalls,
mockCancelAllToolCalls,
mockMarkToolsAsSubmitted,
vi.fn(),
mockCancelAllToolCalls,
0,
]);
rerender();
@@ -2638,7 +2656,7 @@ describe('useGeminiStream', () => {
})(),
);
const { result } = renderHook(() =>
const { result } = renderHookWithProviders(() =>
useGeminiStream(
new MockedGeminiClientClass(mockConfig),
[],
@@ -2695,7 +2713,7 @@ describe('useGeminiStream', () => {
})(),
);
const { result } = renderHook(() =>
const { result } = renderHookWithProviders(() =>
useGeminiStream(
new MockedGeminiClientClass(mockConfig),
[],
@@ -2763,7 +2781,7 @@ describe('useGeminiStream', () => {
})(),
);
const { result } = renderHook(() =>
const { result } = renderHookWithProviders(() =>
useGeminiStream(
new MockedGeminiClientClass(mockConfig),
[],

View File

@@ -51,6 +51,7 @@ import type {
HistoryItem,
HistoryItemWithoutId,
HistoryItemToolGroup,
IndividualToolCallDisplay,
SlashCommandProcessorResult,
HistoryItemModel,
} from '../types.js';
@@ -91,6 +92,48 @@ function showCitations(settings: LoadedSettings): boolean {
return true;
}
/**
* Calculates the current streaming state based on tool call status and responding flag.
*/
function calculateStreamingState(
isResponding: boolean,
toolCalls: TrackedToolCall[],
): StreamingState {
if (toolCalls.some((tc) => tc.status === 'awaiting_approval')) {
return StreamingState.WaitingForConfirmation;
}
const isAnyToolActive = toolCalls.some((tc) => {
// These statuses indicate active processing
if (
tc.status === 'executing' ||
tc.status === 'scheduled' ||
tc.status === 'validating'
) {
return true;
}
// Terminal statuses (success, error, cancelled) still count as "Responding"
// if the result hasn't been submitted back to Gemini yet.
if (
tc.status === 'success' ||
tc.status === 'error' ||
tc.status === 'cancelled'
) {
return !(tc as TrackedCompletedToolCall | TrackedCancelledToolCall)
.responseSubmittedToGemini;
}
return false;
});
if (isResponding || isAnyToolActive) {
return StreamingState.Responding;
}
return StreamingState.Idle;
}
/**
* Manages the Gemini stream, including user input, command processing,
* API interaction, and tool call lifecycle.
@@ -130,6 +173,10 @@ export const useGeminiStream = (
useStateAndRef<HistoryItemWithoutId | null>(null);
const [lastGeminiActivityTime, setLastGeminiActivityTime] =
useState<number>(0);
const [pushedToolCallIds, pushedToolCallIdsRef, setPushedToolCallIds] =
useStateAndRef<Set<string>>(new Set());
const [_isFirstToolInGroup, isFirstToolInGroupRef, setIsFirstToolInGroup] =
useStateAndRef<boolean>(true);
const processedMemoryToolsRef = useRef<Set<string>>(new Set());
const { startNewPrompt, getPromptCount } = useSessionStats();
const storage = config.storage;
@@ -162,12 +209,18 @@ export const useGeminiStream = (
async (completedToolCallsFromScheduler) => {
// This onComplete is called when ALL scheduled tools for a given batch are done.
if (completedToolCallsFromScheduler.length > 0) {
// Add the final state of these tools to the history for display.
addItem(
mapTrackedToolCallsToDisplay(
completedToolCallsFromScheduler as TrackedToolCall[],
),
// Add only the tools that haven't been pushed to history yet.
const toolsToPush = completedToolCallsFromScheduler.filter(
(tc) => !pushedToolCallIdsRef.current.has(tc.request.callId),
);
if (toolsToPush.length > 0) {
addItem(
mapTrackedToolCallsToDisplay(toolsToPush as TrackedToolCall[], {
borderTop: isFirstToolInGroupRef.current,
borderBottom: true,
}),
);
}
// Clear the live-updating display now that the final state is in history.
setToolCallsForDisplay([]);
@@ -205,12 +258,139 @@ export const useGeminiStream = (
getPreferredEditor,
);
const pendingToolCallGroupDisplay = useMemo(
() =>
toolCalls.length ? mapTrackedToolCallsToDisplay(toolCalls) : undefined,
[toolCalls],
const streamingState = useMemo(
() => calculateStreamingState(isResponding, toolCalls),
[isResponding, toolCalls],
);
// Reset tracking when a new batch of tools starts
useEffect(() => {
if (toolCalls.length > 0) {
const isNewBatch = !toolCalls.some((tc) =>
pushedToolCallIdsRef.current.has(tc.request.callId),
);
if (isNewBatch) {
setPushedToolCallIds(new Set());
setIsFirstToolInGroup(true);
}
} else if (streamingState === StreamingState.Idle) {
// Clear when idle to be ready for next turn
setPushedToolCallIds(new Set());
setIsFirstToolInGroup(true);
}
}, [
toolCalls,
pushedToolCallIdsRef,
setPushedToolCallIds,
setIsFirstToolInGroup,
streamingState,
]);
// Push completed tools to history as they finish
useEffect(() => {
const toolsToPush: TrackedToolCall[] = [];
for (const tc of toolCalls) {
if (pushedToolCallIdsRef.current.has(tc.request.callId)) continue;
if (
tc.status === 'success' ||
tc.status === 'error' ||
tc.status === 'cancelled'
) {
toolsToPush.push(tc);
} else {
// Stop at first non-terminal tool to preserve order
break;
}
}
if (toolsToPush.length > 0) {
const newPushed = new Set(pushedToolCallIdsRef.current);
let isFirst = isFirstToolInGroupRef.current;
for (const tc of toolsToPush) {
newPushed.add(tc.request.callId);
const isLastInBatch = tc === toolCalls[toolCalls.length - 1];
const historyItem = mapTrackedToolCallsToDisplay(tc, {
borderTop: isFirst,
borderBottom: isLastInBatch,
});
addItem(historyItem);
isFirst = false;
}
setPushedToolCallIds(newPushed);
setIsFirstToolInGroup(false);
}
}, [
toolCalls,
pushedToolCallIdsRef,
isFirstToolInGroupRef,
setPushedToolCallIds,
setIsFirstToolInGroup,
addItem,
]);
const pendingToolGroupItems = useMemo((): HistoryItemWithoutId[] => {
const remainingTools = toolCalls.filter(
(tc) => !pushedToolCallIds.has(tc.request.callId),
);
const items: HistoryItemWithoutId[] = [];
if (remainingTools.length > 0) {
items.push(
mapTrackedToolCallsToDisplay(remainingTools, {
borderTop: pushedToolCallIds.size === 0,
borderBottom: false, // Stay open to connect with the slice below
}),
);
}
// Always show a bottom border slice if we have ANY tools in the batch
// and we haven't finished pushing the whole batch to history yet.
// Once all tools are terminal and pushed, the last history item handles the closing border.
const allTerminal =
toolCalls.length > 0 &&
toolCalls.every(
(tc) =>
tc.status === 'success' ||
tc.status === 'error' ||
tc.status === 'cancelled',
);
const allPushed =
toolCalls.length > 0 &&
toolCalls.every((tc) => pushedToolCallIds.has(tc.request.callId));
const isEventDriven = config.isEventDrivenSchedulerEnabled();
const anyVisibleInHistory = pushedToolCallIds.size > 0;
const anyVisibleInPending = remainingTools.some((tc) => {
if (!isEventDriven) return true;
return (
tc.status !== 'scheduled' &&
tc.status !== 'validating' &&
tc.status !== 'awaiting_approval'
);
});
if (
toolCalls.length > 0 &&
!(allTerminal && allPushed) &&
(anyVisibleInHistory || anyVisibleInPending)
) {
items.push({
type: 'tool_group' as const,
tools: [] as IndividualToolCallDisplay[],
borderTop: false,
borderBottom: true,
});
}
return items;
}, [toolCalls, pushedToolCallIds, config]);
const activeToolPtyId = useMemo(() => {
const executingShellTool = toolCalls?.find(
(tc) =>
@@ -271,29 +451,6 @@ export const useGeminiStream = (
prevActiveShellPtyIdRef.current = activeShellPtyId;
}, [activeShellPtyId, addItem]);
const streamingState = useMemo(() => {
if (toolCalls.some((tc) => tc.status === 'awaiting_approval')) {
return StreamingState.WaitingForConfirmation;
}
if (
isResponding ||
toolCalls.some(
(tc) =>
tc.status === 'executing' ||
tc.status === 'scheduled' ||
tc.status === 'validating' ||
((tc.status === 'success' ||
tc.status === 'error' ||
tc.status === 'cancelled') &&
!(tc as TrackedCompletedToolCall | TrackedCancelledToolCall)
.responseSubmittedToGemini),
)
) {
return StreamingState.Responding;
}
return StreamingState.Idle;
}, [isResponding, toolCalls]);
useEffect(() => {
if (
config.getApprovalMode() === ApprovalMode.YOLO &&
@@ -1349,10 +1506,10 @@ export const useGeminiStream = (
const pendingHistoryItems = useMemo(
() =>
[pendingHistoryItem, pendingToolCallGroupDisplay].filter(
(i) => i !== undefined && i !== null,
[pendingHistoryItem, ...pendingToolGroupItems].filter(
(i): i is HistoryItemWithoutId => i !== undefined && i !== null,
),
[pendingHistoryItem, pendingToolCallGroupDisplay],
[pendingHistoryItem, pendingToolGroupItems],
);
useEffect(() => {

View File

@@ -11,7 +11,7 @@ import React from 'react';
// times in the same function.
export const useStateAndRef = <
// Everything but function.
T extends object | null | undefined | number | string,
T extends object | null | undefined | number | string | boolean,
>(
initialValue: T,
) => {

View File

@@ -36,6 +36,8 @@ export type {
TrackedExecutingToolCall,
TrackedCompletedToolCall,
TrackedCancelledToolCall,
MarkToolsAsSubmittedFn,
CancelAllFn,
};
// Unified type that covers both implementations