gemini-cli/packages/core/src/services/executionLifecycleService.ts

/**
 * @license
 * Copyright 2025 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */

import type { AnsiOutput } from '../utils/terminalSerializer.js';

export type ExecutionMethod =
  | 'lydell-node-pty'
  | 'node-pty'
  | 'child_process'
  | 'none';

export interface ExecutionResult {
  rawOutput: Buffer;
  output: string;
  exitCode: number | null;
  signal: number | null;
  error: Error | null;
  aborted: boolean;
  pid: number | undefined;
  executionMethod: ExecutionMethod;
  backgrounded?: boolean;
}

export interface ExecutionHandle {
  pid: number | undefined;
  result: Promise<ExecutionResult>;
}

export type ExecutionOutputEvent =
  | {
      type: 'data';
      chunk: string | AnsiOutput;
    }
  | {
      type: 'binary_detected';
    }
  | {
      type: 'binary_progress';
      bytesReceived: number;
    }
  | {
      type: 'exit';
      exitCode: number | null;
      signal: number | null;
    };

export interface ExecutionCompletionOptions {
  exitCode?: number | null;
  signal?: number | null;
  error?: Error | null;
  aborted?: boolean;
}

export interface ExternalExecutionRegistration {
  executionMethod: ExecutionMethod;
  initialOutput?: string;
  getBackgroundOutput?: () => string;
  getSubscriptionSnapshot?: () => string | AnsiOutput | undefined;
  writeInput?: (input: string) => void;
  kill?: () => void;
  isActive?: () => boolean;
}

interface ManagedExecutionBase {
  executionMethod: ExecutionMethod;
  output: string;
  getBackgroundOutput?: () => string;
  getSubscriptionSnapshot?: () => string | AnsiOutput | undefined;
}

interface VirtualExecutionState extends ManagedExecutionBase {
  kind: 'virtual';
  onKill?: () => void;
}

interface ExternalExecutionState extends ManagedExecutionBase {
  kind: 'external';
  writeInput?: (input: string) => void;
  kill?: () => void;
  isActive?: () => boolean;
}

type ManagedExecutionState = VirtualExecutionState | ExternalExecutionState;

/**
 * Central owner for execution backgrounding lifecycle across shell and tools.
 */
export class ExecutionLifecycleService {
  private static readonly EXIT_INFO_TTL_MS = 5 * 60 * 1000;
  private static nextVirtualExecutionId = 2_000_000_000;

  private static activeExecutions = new Map<number, ManagedExecutionState>();
  private static activeResolvers = new Map<
    number,
    (result: ExecutionResult) => void
  >();
  private static activeListeners = new Map<
    number,
    Set<(event: ExecutionOutputEvent) => void>
  >();
  private static exitedExecutionInfo = new Map<
    number,
    { exitCode: number; signal?: number }
  >();

  private static storeExitInfo(
    executionId: number,
    exitCode: number,
    signal?: number,
  ): void {
    this.exitedExecutionInfo.set(executionId, {
      exitCode,
      signal,
    });
    setTimeout(() => {
      this.exitedExecutionInfo.delete(executionId);
    }, this.EXIT_INFO_TTL_MS).unref();
  }

  private static allocateVirtualExecutionId(): number {
    let executionId = ++this.nextVirtualExecutionId;
    while (this.activeExecutions.has(executionId)) {
      executionId = ++this.nextVirtualExecutionId;
    }
    return executionId;
  }

  private static createPendingResult(
    executionId: number,
  ): Promise<ExecutionResult> {
    return new Promise<ExecutionResult>((resolve) => {
      this.activeResolvers.set(executionId, resolve);
    });
  }

  private static createAbortedResult(
    executionId: number,
    execution: ManagedExecutionState,
  ): ExecutionResult {
    const output = execution.getBackgroundOutput?.() ?? execution.output;
    return {
      rawOutput: Buffer.from(output, 'utf8'),
      output,
      exitCode: 130,
      signal: null,
      error: new Error('Operation cancelled by user.'),
      aborted: true,
      pid: executionId,
      executionMethod: execution.executionMethod,
    };
  }

  /**
   * Resets lifecycle state for isolated unit tests.
   */
  static resetForTest(): void {
    this.activeExecutions.clear();
    this.activeResolvers.clear();
    this.activeListeners.clear();
    this.exitedExecutionInfo.clear();
    this.nextVirtualExecutionId = 2_000_000_000;
  }

  static registerExecution(
    executionId: number,
    registration: ExternalExecutionRegistration,
  ): ExecutionHandle {
    if (this.activeExecutions.has(executionId) || this.activeResolvers.has(executionId)) {
      throw new Error(`Execution ${executionId} is already registered.`);
    }
    this.exitedExecutionInfo.delete(executionId);

    this.activeExecutions.set(executionId, {
      executionMethod: registration.executionMethod,
      output: registration.initialOutput ?? '',
      kind: 'external',
      getBackgroundOutput: registration.getBackgroundOutput,
      getSubscriptionSnapshot: registration.getSubscriptionSnapshot,
      writeInput: registration.writeInput,
      kill: registration.kill,
      isActive: registration.isActive,
    });

    return {
      pid: executionId,
      result: this.createPendingResult(executionId),
    };
  }

  static createVirtualExecution(
    initialOutput = '',
    onKill?: () => void,
  ): ExecutionHandle {
    const executionId = this.allocateVirtualExecutionId();

    this.activeExecutions.set(executionId, {
      executionMethod: 'none',
      output: initialOutput,
      kind: 'virtual',
      onKill,
      getBackgroundOutput: () => {
        const state = this.activeExecutions.get(executionId);
        return state?.output ?? initialOutput;
      },
      getSubscriptionSnapshot: () => {
        const state = this.activeExecutions.get(executionId);
        return state?.output ?? initialOutput;
      },
    });

    return {
      pid: executionId,
      result: this.createPendingResult(executionId),
    };
  }

  /**
   * @deprecated Use createVirtualExecution() for new call sites.
   */
  static createExecution(
    initialOutput = '',
    onKill?: () => void,
  ): ExecutionHandle {
    return this.createVirtualExecution(initialOutput, onKill);
  }

  static appendOutput(executionId: number, chunk: string): void {
    const execution = this.activeExecutions.get(executionId);
    if (!execution || chunk.length === 0) {
      return;
    }

    execution.output += chunk;
    this.emitEvent(executionId, { type: 'data', chunk });
  }

  static emitEvent(executionId: number, event: ExecutionOutputEvent): void {
    const listeners = this.activeListeners.get(executionId);
    if (listeners) {
      listeners.forEach((listener) => listener(event));
    }
  }

  private static resolvePending(
    executionId: number,
    result: ExecutionResult,
  ): void {
    const resolve = this.activeResolvers.get(executionId);
    if (!resolve) {
      return;
    }

    resolve(result);
    this.activeResolvers.delete(executionId);
  }

  private static settleExecution(
    executionId: number,
    result: ExecutionResult,
  ): void {
    if (!this.activeExecutions.has(executionId)) {
      return;
    }

    this.resolvePending(executionId, result);
    this.emitEvent(executionId, {
      type: 'exit',
      exitCode: result.exitCode,
      signal: result.signal,
    });

    this.activeListeners.delete(executionId);
    this.activeExecutions.delete(executionId);
    this.storeExitInfo(
      executionId,
      result.exitCode ?? 0,
      result.signal ?? undefined,
    );
  }

  static completeVirtualExecution(
    executionId: number,
    options?: ExecutionCompletionOptions,
  ): void {
    const execution = this.activeExecutions.get(executionId);
    if (!execution) {
      return;
    }

    const {
      error = null,
      aborted = false,
      exitCode = error ? 1 : 0,
      signal = null,
    } = options ?? {};

    const output = execution.getBackgroundOutput?.() ?? execution.output;

    this.settleExecution(executionId, {
      rawOutput: Buffer.from(output, 'utf8'),
      output,
      exitCode,
      signal,
      error,
      aborted,
      pid: executionId,
      executionMethod: execution.executionMethod,
    });
  }

  /**
   * @deprecated Use completeVirtualExecution() for new call sites.
   */
  static completeExecution(
    executionId: number,
    options?: ExecutionCompletionOptions,
  ): void {
    this.completeVirtualExecution(executionId, options);
  }

  static completeWithResult(
    executionId: number,
    result: ExecutionResult,
  ): void {
    this.settleExecution(executionId, result);
  }

  /**
   * @deprecated Use completeWithResult() for new call sites.
   */
  static finalizeExecution(
    executionId: number,
    result: ExecutionResult,
  ): void {
    this.completeWithResult(executionId, result);
  }

  static background(executionId: number): void {
    const resolve = this.activeResolvers.get(executionId);
    if (!resolve) {
      return;
    }

    const execution = this.activeExecutions.get(executionId);
    if (!execution) {
      return;
    }

    const output = execution.getBackgroundOutput?.() ?? execution.output;

    resolve({
      rawOutput: Buffer.from(''),
      output,
      exitCode: null,
      signal: null,
      error: null,
      aborted: false,
      pid: executionId,
      executionMethod: execution.executionMethod,
      backgrounded: true,
    });

    this.activeResolvers.delete(executionId);
  }

  static subscribe(
    executionId: number,
    listener: (event: ExecutionOutputEvent) => void,
  ): () => void {
    if (!this.activeListeners.has(executionId)) {
      this.activeListeners.set(executionId, new Set());
    }
    this.activeListeners.get(executionId)?.add(listener);

    const execution = this.activeExecutions.get(executionId);
    if (execution) {
      const snapshot =
        execution.getSubscriptionSnapshot?.() ??
        (execution.output.length > 0 ? execution.output : undefined);
      if (snapshot && (typeof snapshot !== 'string' || snapshot.length > 0)) {
        listener({ type: 'data', chunk: snapshot });
      }
    }

    return () => {
      this.activeListeners.get(executionId)?.delete(listener);
      if (this.activeListeners.get(executionId)?.size === 0) {
        this.activeListeners.delete(executionId);
      }
    };
  }

  static onExit(
    executionId: number,
    callback: (exitCode: number, signal?: number) => void,
  ): () => void {
    if (this.activeExecutions.has(executionId)) {
      const listener = (event: ExecutionOutputEvent) => {
        if (event.type === 'exit') {
          callback(event.exitCode ?? 0, event.signal ?? undefined);
          unsubscribe();
        }
      };
      const unsubscribe = this.subscribe(executionId, listener);
      return unsubscribe;
    }

    const exitedInfo = this.exitedExecutionInfo.get(executionId);
    if (exitedInfo) {
      callback(exitedInfo.exitCode, exitedInfo.signal);
    }

    return () => {};
  }

  static kill(executionId: number): void {
    const execution = this.activeExecutions.get(executionId);
    if (!execution) {
      return;
    }

    if (execution.kind === 'virtual') {
      execution.onKill?.();
    }

    if (execution.kind === 'external') {
      execution.kill?.();
    }

    this.completeWithResult(
      executionId,
      this.createAbortedResult(executionId, execution),
    );
  }

  static isActive(executionId: number): boolean {
    const execution = this.activeExecutions.get(executionId);
    if (!execution) {
      try {
        return process.kill(executionId, 0);
      } catch {
        return false;
      }
    }

    if (execution.kind === 'virtual') {
      return true;
    }

    if (execution.kind === 'external' && execution.isActive) {
      try {
        return execution.isActive();
      } catch {
        return false;
      }
    }

    try {
      return process.kill(executionId, 0);
    } catch {
      return false;
    }
  }

  static writeInput(executionId: number, input: string): void {
    const execution = this.activeExecutions.get(executionId);
    if (execution?.kind === 'external') {
      execution.writeInput?.(input);
    }
  }
}