From f2e2014894fae693d8cb69134ae9863cc60f5a50 Mon Sep 17 00:00:00 2001
From: Adam Weidman <adamweidman@Adams-MacBook-Pro.local>
Date: Sun, 8 Mar 2026 20:41:08 -0400
Subject: [PATCH] refactor: generalize CLI background execution handling

---
 .../cli/src/ui/hooks/useGeminiStream.test.tsx | 29 +++++++++++
 packages/cli/src/ui/hooks/useGeminiStream.ts  | 52 ++++++++++---------
 2 files changed, 57 insertions(+), 24 deletions(-)

diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx
index 1f2ef5f90c..389b757e99 100644
--- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx
+++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx
@@ -599,6 +599,35 @@ describe('useGeminiStream', () => {
     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 () => {
     const toolCall1ResponseParts: Part[] = [{ text: 'tool 1 final response' }];
     const toolCall2ResponseParts: Part[] = [{ text: 'tool 2 final response' }];
diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts
index d30bc78e21..4dd27f1f0b 100644
--- a/packages/cli/src/ui/hooks/useGeminiStream.ts
+++ b/packages/cli/src/ui/hooks/useGeminiStream.ts
@@ -94,8 +94,8 @@ type ToolResponseWithParts = ToolCallResponseInfo & {
   llmContent?: PartListUnion;
 };
 
-interface BackgroundedShellInfo {
-  pid: number;
+interface BackgroundedToolInfo {
+  executionId: number;
   command: string;
   initialOutput: string;
 }
@@ -151,29 +151,34 @@ function getBackgroundExecutionId(
   return undefined;
 }
 
-function getBackgroundedShellInfo(
+function getBackgroundedToolInfo(
   toolCall: TrackedCompletedToolCall | TrackedCancelledToolCall,
-): BackgroundedShellInfo | undefined {
-  if (toolCall.request.name !== SHELL_COMMAND_NAME) {
-    return undefined;
-  }
-
+): BackgroundedToolInfo | undefined {
   const response = toolCall.response as ToolResponseWithParts;
   const rawData = response?.data;
   const data = isBackgroundExecutionData(rawData) ? rawData : undefined;
   const executionId = data ? getBackgroundExecutionId(data) : undefined;
 
-  if (!executionId) {
+  if (executionId === undefined) {
     return undefined;
   }
 
   return {
-    pid: executionId,
-    command: data.command ?? 'shell',
+    executionId,
+    command: data.command ?? toolCall.request.name,
     initialOutput: data.initialOutput ?? '',
   };
 }
 
+function isBackgroundableExecutingToolCall(
+  toolCall: TrackedToolCall,
+): toolCall is TrackedExecutingToolCall {
+  return (
+    toolCall.status === CoreToolCallStatus.Executing &&
+    typeof toolCall.pid === 'number'
+  );
+}
+
 function showCitations(settings: LoadedSettings): boolean {
   const enabled = settings.merged.ui.showCitations;
   if (enabled !== undefined) {
@@ -362,13 +367,11 @@ export const useGeminiStream = (
     getPreferredEditor,
   );
 
-  const activeToolExecutionId = useMemo(() => {
-    const executingShellTool = toolCalls.find(
-      (tc): tc is TrackedExecutingToolCall =>
-        tc.status === CoreToolCallStatus.Executing &&
-        tc.request.name === SHELL_COMMAND_NAME,
+  const activeBackgroundExecutionId = useMemo(() => {
+    const executingBackgroundableTool = toolCalls.find(
+      isBackgroundableExecutingToolCall,
     );
-    return executingShellTool?.pid;
+    return executingBackgroundableTool?.pid;
   }, [toolCalls]);
 
   const onExec = useCallback(async (done: Promise<void>) => {
@@ -398,7 +401,7 @@ export const useGeminiStream = (
     setShellInputFocused,
     terminalWidth,
     terminalHeight,
-    activeToolExecutionId,
+    activeBackgroundExecutionId,
   );
 
   const streamingState = useMemo(
@@ -576,7 +579,8 @@ export const useGeminiStream = (
     onComplete: (result: { userSelection: 'disable' | 'keep' }) => void;
   } | null>(null);
 
-  const activePtyId = activeShellPtyId || activeToolExecutionId;
+  const activePtyId =
+    activeShellPtyId ?? activeBackgroundExecutionId ?? undefined;
 
   const prevActiveShellPtyIdRef = useRef<number | null>(null);
   useEffect(() => {
@@ -1703,12 +1707,12 @@ export const useGeminiStream = (
       );
 
       for (const toolCall of completedAndReadyToSubmitTools) {
-        const backgroundedShell = getBackgroundedShellInfo(toolCall);
-        if (backgroundedShell) {
+        const backgroundedTool = getBackgroundedToolInfo(toolCall);
+        if (backgroundedTool) {
           registerBackgroundShell(
-            backgroundedShell.pid,
-            backgroundedShell.command,
-            backgroundedShell.initialOutput,
+            backgroundedTool.executionId,
+            backgroundedTool.command,
+            backgroundedTool.initialOutput,
           );
         }
       }