Support backgrounding remote agent executions via Ctrl+B

packages/core/src/agents/remote-invocation.test.ts

@@ -22,6 +22,7 @@ import type { RemoteAgentDefinition } from './types.js';
 import { createMockMessageBus } from '../test-utils/mock-message-bus.js';
 import { A2AAuthProviderFactory } from './auth-provider/factory.js';
 import type { A2AAuthProvider } from './auth-provider/types.js';
+import { ShellExecutionService } from '../services/shellExecutionService.js';
 
 // Mock A2AClientManager
 vi.mock('./a2a-client-manager.js', () => ({
@@ -583,6 +584,88 @@ describe('RemoteAgentInvocation', () => {
         'Generating...\n\nArtifact (Result):\nPart 1 Part 2',
       );
     });
+
+    it('should support Ctrl+B backgrounding through ShellExecutionService', async () => {
+      mockClientManager.getClient.mockReturnValue({});
+
+      let releaseSecondChunk: (() => void) | undefined;
+      const secondChunkGate = new Promise<void>((resolve) => {
+        releaseSecondChunk = resolve;
+      });
+
+      mockClientManager.sendMessageStream.mockImplementation(
+        async function* () {
+          yield {
+            kind: 'message',
+            messageId: 'msg-1',
+            role: 'agent',
+            parts: [{ kind: 'text', text: 'Chunk 1' }],
+          };
+          await secondChunkGate;
+          yield {
+            kind: 'message',
+            messageId: 'msg-2',
+            role: 'agent',
+            parts: [{ kind: 'text', text: 'Chunk 2' }],
+          };
+        },
+      );
+
+      let pid: number | undefined;
+      const onExit = vi.fn();
+      let unsubscribeOnExit: (() => void) | undefined;
+      const streamedOutputChunks: string[] = [];
+      let unsubscribeStream: (() => void) | undefined;
+
+      const updateOutput = vi.fn((output: unknown) => {
+        if (output === 'Chunk 1' && pid) {
+          ShellExecutionService.background(pid);
+          unsubscribeStream = ShellExecutionService.subscribe(pid, (event) => {
+            if (event.type === 'data' && typeof event.chunk === 'string') {
+              streamedOutputChunks.push(event.chunk);
+            }
+          });
+        }
+      });
+
+      const invocation = new RemoteAgentInvocation(
+        mockDefinition,
+        { query: 'long task' },
+        mockMessageBus,
+      );
+
+      const resultPromise = invocation.execute(
+        new AbortController().signal,
+        updateOutput,
+        undefined,
+        (newPid) => {
+          pid = newPid;
+          unsubscribeOnExit = ShellExecutionService.onExit(newPid, onExit);
+        },
+      );
+
+      const result = await resultPromise;
+      expect(pid).toBeDefined();
+      expect(result.returnDisplay).toContain(
+        'Remote agent moved to background',
+      );
+      expect(result.data).toMatchObject({
+        pid,
+        initialOutput: 'Chunk 1',
+      });
+
+      releaseSecondChunk?.();
+
+      await vi.waitFor(() => {
+        expect(onExit).toHaveBeenCalledWith(0, undefined);
+      });
+      await vi.waitFor(() => {
+        expect(streamedOutputChunks.join('')).toContain('Chunk 2');
+      });
+
+      unsubscribeStream?.();
+      unsubscribeOnExit?.();
+    });
   });
 
   describe('Confirmations', () => {

packages/core/src/agents/remote-invocation.ts

@@ -27,6 +27,7 @@ import type { AuthenticationHandler } from '@a2a-js/sdk/client';
 import { debugLogger } from '../utils/debugLogger.js';
 import type { AnsiOutput } from '../utils/terminalSerializer.js';
 import { A2AAuthProviderFactory } from './auth-provider/factory.js';
+import { ExecutionLifecycleService } from '../services/executionLifecycleService.js';
 
 /**
  * Authentication handler implementation using Google Application Default Credentials (ADC).
@@ -145,102 +146,192 @@ export class RemoteAgentInvocation extends BaseToolInvocation<
     };
   }
 
+  private publishBackgroundDelta(
+    pid: number,
+    previousOutput: string,
+    nextOutput: string,
+  ): string {
+    if (nextOutput.length === 0 || nextOutput === previousOutput) {
+      return previousOutput;
+    }
+
+    if (nextOutput.startsWith(previousOutput)) {
+      ExecutionLifecycleService.appendOutput(
+        pid,
+        nextOutput.slice(previousOutput.length),
+      );
+      return nextOutput;
+    }
+
+    // If the reassembled output changes non-monotonically, resync by appending
+    // the full latest snapshot with a clear separator.
+    ExecutionLifecycleService.appendOutput(
+      pid,
+      `\n\n[Output updated]\n${nextOutput}`,
+    );
+    return nextOutput;
+  }
+
   async execute(
     _signal: AbortSignal,
     updateOutput?: (output: string | AnsiOutput) => void,
+    _shellExecutionConfig?: unknown,
+    setPidCallback?: (pid: number) => void,
   ): Promise<ToolResult> {
-    // 1. Ensure the agent is loaded (cached by manager)
-    // We assume the user has provided an access token via some mechanism (TODO),
-    // or we rely on ADC.
     const reassembler = new A2AResultReassembler();
-    try {
+    const executionController = new AbortController();
-      const priorState = RemoteAgentInvocation.sessionState.get(
+    const onAbort = () => executionController.abort();
-        this.definition.name,
+    _signal.addEventListener('abort', onAbort, { once: true });
-      );
-      if (priorState) {
-        this.contextId = priorState.contextId;
-        this.taskId = priorState.taskId;
-      }
 
-      const authHandler = await this.getAuthHandler();
+    const { pid, result } = ExecutionLifecycleService.createExecution(
+      '',
+      () => executionController.abort(),
+    );
+    if (pid === undefined) {
+      _signal.removeEventListener('abort', onAbort);
+      return {
+        llmContent: [
+          { text: 'Error calling remote agent: missing execution pid.' },
+        ],
+        returnDisplay: 'Error calling remote agent: missing execution pid.',
+        error: {
+          message: 'Error calling remote agent: missing execution pid.',
+        },
+      };
+    }
+    const backgroundPid = pid;
+    setPidCallback?.(backgroundPid);
 
-      if (!this.clientManager.getClient(this.definition.name)) {
+    const run = async () => {
-        await this.clientManager.loadAgent(
+      let lastOutput = '';
+      try {
+        const priorState = RemoteAgentInvocation.sessionState.get(
           this.definition.name,
-          this.definition.agentCardUrl,
-          authHandler,
         );
-      }
+        if (priorState) {
+          this.contextId = priorState.contextId;
+          this.taskId = priorState.taskId;
+        }
 
-      const message = this.params.query;
+        const authHandler = await this.getAuthHandler();
 
-      const stream = this.clientManager.sendMessageStream(
+        if (!this.clientManager.getClient(this.definition.name)) {
-        this.definition.name,
+          await this.clientManager.loadAgent(
-        message,
+            this.definition.name,
-        {
+            this.definition.agentCardUrl,
+            authHandler,
+          );
+        }
+
+        const stream = this.clientManager.sendMessageStream(
+          this.definition.name,
+          this.params.query,
+          {
+            contextId: this.contextId,
+            taskId: this.taskId,
+            signal: executionController.signal,
+          },
+        );
+
+        let finalResponse: SendMessageResult | undefined;
+
+        for await (const chunk of stream) {
+          if (executionController.signal.aborted) {
+            throw new Error('Operation aborted');
+          }
+          finalResponse = chunk;
+          reassembler.update(chunk);
+
+          const currentOutput = reassembler.toString();
+          lastOutput = this.publishBackgroundDelta(
+            backgroundPid,
+            lastOutput,
+            currentOutput,
+          );
+          if (updateOutput) {
+            updateOutput(currentOutput);
+          }
+
+          const {
+            contextId: newContextId,
+            taskId: newTaskId,
+            clearTaskId,
+          } = extractIdsFromResponse(chunk);
+
+          if (newContextId) {
+            this.contextId = newContextId;
+          }
+
+          this.taskId = clearTaskId ? undefined : (newTaskId ?? this.taskId);
+        }
+
+        if (!finalResponse) {
+          throw new Error('No response from remote agent.');
+        }
+
+        debugLogger.debug(
+          `[RemoteAgent] Final response from ${this.definition.name}:\n${JSON.stringify(finalResponse, null, 2)}`,
+        );
+
+        ExecutionLifecycleService.completeExecution(backgroundPid, {
+          exitCode: 0,
+        });
+      } catch (error: unknown) {
+        const partialOutput = reassembler.toString();
+        lastOutput = this.publishBackgroundDelta(
+          backgroundPid,
+          lastOutput,
+          partialOutput,
+        );
+        const errorMessage = `Error calling remote agent: ${
+          error instanceof Error ? error.message : String(error)
+        }`;
+        ExecutionLifecycleService.completeExecution(backgroundPid, {
+          error: new Error(errorMessage),
+          aborted: executionController.signal.aborted,
+          exitCode: executionController.signal.aborted ? 130 : 1,
+        });
+      } finally {
+        _signal.removeEventListener('abort', onAbort);
+        // Persist state even on partial failures or aborts to maintain conversational continuity.
+        RemoteAgentInvocation.sessionState.set(this.definition.name, {
           contextId: this.contextId,
           taskId: this.taskId,
-          signal: _signal,
+        });
-        },
-      );
-
-      let finalResponse: SendMessageResult | undefined;
-
-      for await (const chunk of stream) {
-        if (_signal.aborted) {
-          throw new Error('Operation aborted');
-        }
-        finalResponse = chunk;
-        reassembler.update(chunk);
-
-        if (updateOutput) {
-          updateOutput(reassembler.toString());
-        }
-
-        const {
-          contextId: newContextId,
-          taskId: newTaskId,
-          clearTaskId,
-        } = extractIdsFromResponse(chunk);
-
-        if (newContextId) {
-          this.contextId = newContextId;
-        }
-
-        this.taskId = clearTaskId ? undefined : (newTaskId ?? this.taskId);
       }
+    };
 
-      if (!finalResponse) {
+    void run();
-        throw new Error('No response from remote agent.');
+    const executionResult = await result;
-      }
-
-      const finalOutput = reassembler.toString();
-
-      debugLogger.debug(
-        `[RemoteAgent] Final response from ${this.definition.name}:\n${JSON.stringify(finalResponse, null, 2)}`,
-      );
 
+    if (executionResult.backgrounded) {
+      const command = `${this.getDescription()}: ${this.params.query}`;
+      const backgroundMessage = `Remote agent moved to background (PID: ${backgroundPid}). Output hidden. Press Ctrl+B to view.`;
       return {
-        llmContent: [{ text: finalOutput }],
+        llmContent: [{ text: backgroundMessage }],
-        returnDisplay: finalOutput,
+        returnDisplay: backgroundMessage,
+        data: {
+          pid: backgroundPid,
+          command,
+          initialOutput: executionResult.output,
+        },
       };
-    } catch (error: unknown) {
+    }
-      const partialOutput = reassembler.toString();
+
-      const errorMessage = `Error calling remote agent: ${error instanceof Error ? error.message : String(error)}`;
+    if (executionResult.error) {
-      const fullDisplay = partialOutput
+      const fullDisplay = executionResult.output
-        ? `${partialOutput}\n\n${errorMessage}`
+        ? `${executionResult.output}\n\n${executionResult.error.message}`
-        : errorMessage;
+        : executionResult.error.message;
       return {
         llmContent: [{ text: fullDisplay }],
         returnDisplay: fullDisplay,
-        error: { message: errorMessage },
+        error: { message: executionResult.error.message },
       };
-    } finally {
-      // Persist state even on partial failures or aborts to maintain conversational continuity.
-      RemoteAgentInvocation.sessionState.set(this.definition.name, {
-        contextId: this.contextId,
-        taskId: this.taskId,
-      });
     }
+
+    return {
+      llmContent: [{ text: executionResult.output }],
+      returnDisplay: executionResult.output,
+    };
   }
 }