TJWaterAgent/src/cli/executeCliCommand.ts

import { spawn } from "node:child_process";

import { type RuntimeSessionContext } from "../runtime/sessionContext.js";

type OutputStream = "stdout" | "stderr";

export type CliExecutionResult = {
  outcome: "completed" | "timeout" | "output_limit";
  exitCode: number | null;
  signal: NodeJS.Signals | null;
  status: number;
  stderr: string;
  stdout: string;
  exceededStream?: OutputStream;
};

type ExecuteCliCommandOptions = {
  apiBaseUrl: string;
  cliPath: string;
  maxOutputBytes: number;
  terminationGraceMs?: number;
};

const getCompletedStatus = (exitCode: number | null, stdout: string) => {
  let errorCode = "";
  try {
    const payload = JSON.parse(stdout) as { error?: { code?: unknown } };
    errorCode =
      typeof payload.error?.code === "string" ? payload.error.code : "";
  } catch {
    errorCode = "";
  }

  if (errorCode === "HTTP_401" || errorCode === "UNAUTHENTICATED") {
    return 401;
  }
  if (errorCode === "HTTP_403") {
    return 403;
  }
  return exitCode === 0 ? 200 : 502;
};

export const executeCliCommand = async (
  context: RuntimeSessionContext,
  command: string,
  timeoutSec: number,
  options: ExecuteCliCommandOptions,
): Promise<CliExecutionResult> => {
  const maxOutputBytes = options.maxOutputBytes;
  if (!Number.isSafeInteger(maxOutputBytes) || maxOutputBytes <= 0) {
    throw new Error("maxOutputBytes must be a positive safe integer");
  }

  const child = spawn(
    options.cliPath,
    ["--auth-stdin", ...command.split(/\s+/).filter(Boolean)],
    { stdio: ["pipe", "pipe", "pipe"] },
  );
  const stdoutChunks: Buffer[] = [];
  const stderrChunks: Buffer[] = [];
  let stdoutBytes = 0;
  let stderrBytes = 0;
  let terminationReason:
    | "timeout"
    | "output_limit"
    | "execution_error"
    | null = null;
  let exceededStream: OutputStream | undefined;
  let terminationStarted = false;
  let settled = false;
  let executionError: Error | null = null;
  let forceKillTimer: ReturnType<typeof setTimeout> | undefined;

  const result = await new Promise<CliExecutionResult>((resolve, reject) => {
    const cleanup = () => {
      clearTimeout(timeoutTimer);
      if (forceKillTimer) {
        clearTimeout(forceKillTimer);
      }
    };

    const terminate = (
      reason: "timeout" | "output_limit" | "execution_error",
    ) => {
      if (terminationStarted) {
        return;
      }
      terminationStarted = true;
      terminationReason = reason;

      if (child.exitCode === null && child.signalCode === null) {
        child.kill("SIGTERM");
      }
      forceKillTimer = setTimeout(() => {
        if (child.exitCode === null && child.signalCode === null) {
          child.kill("SIGKILL");
        }
      }, options.terminationGraceMs ?? 1500);
    };

    const capture = (stream: OutputStream, data: Buffer) => {
      if (terminationReason) {
        return;
      }
      const chunks = stream === "stdout" ? stdoutChunks : stderrChunks;
      const bytes = stream === "stdout" ? stdoutBytes : stderrBytes;
      if (bytes + data.length > maxOutputBytes) {
        exceededStream = stream;
        terminate("output_limit");
        return;
      }
      chunks.push(data);
      if (stream === "stdout") {
        stdoutBytes += data.length;
      } else {
        stderrBytes += data.length;
      }
    };

    const timeoutTimer = setTimeout(() => {
      if (child.exitCode === null && child.signalCode === null) {
        terminate("timeout");
      }
    }, timeoutSec * 1000);

    child.stdout.on("data", (data: Buffer) => capture("stdout", data));
    child.stderr.on("data", (data: Buffer) => capture("stderr", data));
    child.stdin.on("error", (error) => {
      if (terminationReason === null) {
        executionError = error;
        terminate("execution_error");
      }
    });
    child.on("error", (error) => {
      if (terminationReason === null) {
        executionError = error;
        terminate("execution_error");
      }
    });
    child.on("close", (exitCode, signal) => {
      if (settled) {
        return;
      }
      settled = true;
      cleanup();
      if (terminationReason === "timeout") {
        resolve({
          outcome: "timeout",
          exitCode,
          signal,
          status: 504,
          stderr: "",
          stdout: "",
        });
        return;
      }
      if (executionError) {
        reject(executionError);
        return;
      }
      if (terminationReason === "output_limit") {
        resolve({
          outcome: "output_limit",
          exceededStream,
          exitCode,
          signal,
          status: 502,
          stderr: "",
          stdout: "",
        });
        return;
      }

      const stdout = Buffer.concat(stdoutChunks, stdoutBytes).toString("utf-8");
      const stderr = Buffer.concat(stderrChunks, stderrBytes).toString("utf-8");
      resolve({
        outcome: "completed",
        exitCode,
        signal,
        status: getCompletedStatus(exitCode, stdout),
        stderr,
        stdout,
      });
    });

    child.stdin.end(
      JSON.stringify({
        server: options.apiBaseUrl,
        access_token: context.accessToken,
        project_id: context.projectId,
      }),
    );
  });

  return result;
};