refactor: generalize CLI background execution handling

This commit is contained in:
Adam Weidman
2026-03-08 20:41:08 -04:00
parent 6860284344
commit f2e2014894
2 changed files with 57 additions and 24 deletions
@@ -599,6 +599,35 @@ describe('useGeminiStream', () => {
expect(mockSendMessageStream).not.toHaveBeenCalled(); // submitQuery uses this expect(mockSendMessageStream).not.toHaveBeenCalled(); // submitQuery uses this
}); });
it('should expose activePtyId for non-shell executing tools that report an execution ID', () => {
const remoteExecutingTool: TrackedExecutingToolCall = {
request: {
callId: 'remote-call-1',
name: 'remote_agent_call',
args: {},
isClientInitiated: false,
prompt_id: 'prompt-id-remote',
},
status: CoreToolCallStatus.Executing,
responseSubmittedToGemini: false,
tool: {
name: 'remote_agent_call',
displayName: 'Remote Agent',
description: 'Remote agent execution',
build: vi.fn(),
} as any,
invocation: {
getDescription: () => 'Calling remote agent',
} as unknown as AnyToolInvocation,
startTime: Date.now(),
liveOutput: 'working...',
pid: 4242,
};
const { result } = renderTestHook([remoteExecutingTool]);
expect(result.current.activePtyId).toBe(4242);
});
it('should submit tool responses when all tool calls are completed and ready', async () => { it('should submit tool responses when all tool calls are completed and ready', async () => {
const toolCall1ResponseParts: Part[] = [{ text: 'tool 1 final response' }]; const toolCall1ResponseParts: Part[] = [{ text: 'tool 1 final response' }];
const toolCall2ResponseParts: Part[] = [{ text: 'tool 2 final response' }]; const toolCall2ResponseParts: Part[] = [{ text: 'tool 2 final response' }];
+28 -24
View File
@@ -94,8 +94,8 @@ type ToolResponseWithParts = ToolCallResponseInfo & {
llmContent?: PartListUnion; llmContent?: PartListUnion;
}; };
interface BackgroundedShellInfo { interface BackgroundedToolInfo {
pid: number; executionId: number;
command: string; command: string;
initialOutput: string; initialOutput: string;
} }
@@ -151,29 +151,34 @@ function getBackgroundExecutionId(
return undefined; return undefined;
} }
function getBackgroundedShellInfo( function getBackgroundedToolInfo(
toolCall: TrackedCompletedToolCall | TrackedCancelledToolCall, toolCall: TrackedCompletedToolCall | TrackedCancelledToolCall,
): BackgroundedShellInfo | undefined { ): BackgroundedToolInfo | undefined {
if (toolCall.request.name !== SHELL_COMMAND_NAME) {
return undefined;
}
const response = toolCall.response as ToolResponseWithParts; const response = toolCall.response as ToolResponseWithParts;
const rawData = response?.data; const rawData = response?.data;
const data = isBackgroundExecutionData(rawData) ? rawData : undefined; const data = isBackgroundExecutionData(rawData) ? rawData : undefined;
const executionId = data ? getBackgroundExecutionId(data) : undefined; const executionId = data ? getBackgroundExecutionId(data) : undefined;
if (!executionId) { if (executionId === undefined) {
return undefined; return undefined;
} }
return { return {
pid: executionId, executionId,
command: data.command ?? 'shell', command: data.command ?? toolCall.request.name,
initialOutput: data.initialOutput ?? '', initialOutput: data.initialOutput ?? '',
}; };
} }
function isBackgroundableExecutingToolCall(
toolCall: TrackedToolCall,
): toolCall is TrackedExecutingToolCall {
return (
toolCall.status === CoreToolCallStatus.Executing &&
typeof toolCall.pid === 'number'
);
}
function showCitations(settings: LoadedSettings): boolean { function showCitations(settings: LoadedSettings): boolean {
const enabled = settings.merged.ui.showCitations; const enabled = settings.merged.ui.showCitations;
if (enabled !== undefined) { if (enabled !== undefined) {
@@ -362,13 +367,11 @@ export const useGeminiStream = (
getPreferredEditor, getPreferredEditor,
); );
const activeToolExecutionId = useMemo(() => { const activeBackgroundExecutionId = useMemo(() => {
const executingShellTool = toolCalls.find( const executingBackgroundableTool = toolCalls.find(
(tc): tc is TrackedExecutingToolCall => isBackgroundableExecutingToolCall,
tc.status === CoreToolCallStatus.Executing &&
tc.request.name === SHELL_COMMAND_NAME,
); );
return executingShellTool?.pid; return executingBackgroundableTool?.pid;
}, [toolCalls]); }, [toolCalls]);
const onExec = useCallback(async (done: Promise<void>) => { const onExec = useCallback(async (done: Promise<void>) => {
@@ -398,7 +401,7 @@ export const useGeminiStream = (
setShellInputFocused, setShellInputFocused,
terminalWidth, terminalWidth,
terminalHeight, terminalHeight,
activeToolExecutionId, activeBackgroundExecutionId,
); );
const streamingState = useMemo( const streamingState = useMemo(
@@ -576,7 +579,8 @@ export const useGeminiStream = (
onComplete: (result: { userSelection: 'disable' | 'keep' }) => void; onComplete: (result: { userSelection: 'disable' | 'keep' }) => void;
} | null>(null); } | null>(null);
const activePtyId = activeShellPtyId || activeToolExecutionId; const activePtyId =
activeShellPtyId ?? activeBackgroundExecutionId ?? undefined;
const prevActiveShellPtyIdRef = useRef<number | null>(null); const prevActiveShellPtyIdRef = useRef<number | null>(null);
useEffect(() => { useEffect(() => {
@@ -1703,12 +1707,12 @@ export const useGeminiStream = (
); );
for (const toolCall of completedAndReadyToSubmitTools) { for (const toolCall of completedAndReadyToSubmitTools) {
const backgroundedShell = getBackgroundedShellInfo(toolCall); const backgroundedTool = getBackgroundedToolInfo(toolCall);
if (backgroundedShell) { if (backgroundedTool) {
registerBackgroundShell( registerBackgroundShell(
backgroundedShell.pid, backgroundedTool.executionId,
backgroundedShell.command, backgroundedTool.command,
backgroundedShell.initialOutput, backgroundedTool.initialOutput,
); );
} }
} }