fix(agent): bound CLI subprocess execution

2026-08-18 17:00:23 +08:00
parent 9aa5a96e60
commit 18e8b25f48
7 changed files with 384 additions and 65 deletions
@@ -0,0 +1,196 @@
+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;
+};
@@ -1,5 +1,4 @@
 import { randomUUID } from "node:crypto";
-import { spawn } from "node:child_process";
 import cors from "cors";
 import express from "express";

@@ -11,6 +10,7 @@ import {
  runWithCredentialRefresh,
 } from "./auth/credentialRefresh.js";
 import { SessionTranscriptStore } from "./sessions/transcriptStore.js";
+import { executeCliCommand } from "./cli/executeCliCommand.js";
 import { ChatSessionBridge } from "./chat/sessionBridge.js";
 import { config } from "./config.js";
 import { SessionUiStateStore } from "./sessions/uiStateStore.js";
@@ -229,7 +229,12 @@ app.post("/internal/tools/tjwater-cli-call", async (req, res) => {
    result = await runWithCredentialRefresh(
      credentialRefreshCoordinator,
      context,
-      (activeContext) => executeCliCommand(activeContext, command, timeoutSec),
+      (activeContext) =>
+        executeCliCommand(activeContext, command, timeoutSec, {
+          apiBaseUrl: config.TJWATER_API_BASE_URL,
+          cliPath: config.TJWATER_CLI_PATH,
+          maxOutputBytes: config.MAX_INLINE_RESULT_BYTES,
+        }),
    );
  } catch (error) {
    if (!(error instanceof CredentialRefreshError)) {
@@ -266,6 +271,20 @@ app.post("/internal/tools/tjwater-cli-call", async (req, res) => {
    return;
  }

+  if (result.outcome === "output_limit") {
+    res.status(502).json({
+      ok: false,
+      schema_version: "tjwater-cli/v1",
+      summary: "CLI 输出超过安全限制",
+      error: {
+        code: "OUTPUT_LIMIT_EXCEEDED",
+        message: `${result.exceededStream ?? "output"} exceeded ${config.MAX_INLINE_RESULT_BYTES} bytes`,
+        retryable: false,
+      },
+    });
+    return;
+  }
+
  if (result.status === 401) {
    markAuthExpired(
      getRuntimeSessionContext(sessionId) ?? context,
@@ -300,68 +319,6 @@ app.post("/internal/tools/tjwater-cli-call", async (req, res) => {
  }
 });

-const executeCliCommand = async (
-  context: RuntimeSessionContext,
-  command: string,
-  timeoutSec: number,
-) => {
-  const child = spawn(
-    config.TJWATER_CLI_PATH,
-    ["--auth-stdin", ...command.split(/\s+/).filter(Boolean)],
-    { stdio: ["pipe", "pipe", "pipe"] },
-  );
-  let stdout = "";
-  let stderr = "";
-  child.stdout.on("data", (data: Buffer) => {
-    stdout += data.toString("utf-8");
-  });
-  child.stderr.on("data", (data: Buffer) => {
-    stderr += data.toString("utf-8");
-  });
-  child.stdin.write(
-    JSON.stringify({
-      server: config.TJWATER_API_BASE_URL,
-      access_token: context.accessToken,
-      project_id: context.projectId,
-    }),
-  );
-  child.stdin.end();
-
-  const exitCode = await new Promise<number | null>((resolve, reject) => {
-    const timer = setTimeout(() => {
-      child.kill("SIGTERM");
-      resolve(-1);
-    }, timeoutSec * 1000);
-    child.on("close", (code) => {
-      clearTimeout(timer);
-      resolve(code);
-    });
-    child.on("error", (error) => {
-      clearTimeout(timer);
-      reject(error);
-    });
-  });
-  let errorCode = "";
-  try {
-    const payload = JSON.parse(stdout) as { error?: { code?: unknown } };
-    errorCode =
-      typeof payload.error?.code === "string" ? payload.error.code : "";
-  } catch {
-    errorCode = "";
-  }
-  const status =
-    exitCode === -1
-      ? 504
-      : errorCode === "HTTP_401" || errorCode === "UNAUTHENTICATED"
-        ? 401
-        : errorCode === "HTTP_403"
-          ? 403
-          : exitCode === 0
-            ? 200
-            : 502;
-  return { exitCode, status, stderr, stdout };
-};
-
 app.post("/internal/tools/store-render-ref", async (req, res) => {
  if (req.header("x-agent-internal-token") !== internalToken) {
    res.status(403).json({ message: "forbidden" });