openclaw/src/cli/one-shot-exit.ts

import type { RuntimeEnv } from "../runtime.js";
import { defaultRuntime } from "../runtime.js";

type VitestWorkerMarkers = {
  tinypoolState?: unknown;
  vitestWorker?: unknown;
};

let requestedExitCode: number | undefined;

function resolveVitestWorkerMarkers(): VitestWorkerMarkers {
  const processMarkers = process as NodeJS.Process & Record<string, unknown>;
  const globalMarkers = globalThis as typeof globalThis & Record<string, unknown>;
  return {
    tinypoolState: processMarkers["__tinypool_state__"],
    vitestWorker: globalMarkers["__vitest_worker__"],
  };
}

function isVitestWorker(
  env: NodeJS.ProcessEnv,
  markers: VitestWorkerMarkers = resolveVitestWorkerMarkers(),
): boolean {
  const hasVitestEnv =
    env.VITEST === "true" ||
    env.VITEST === "1" ||
    env.VITEST_POOL_ID !== undefined ||
    env.VITEST_WORKER_ID !== undefined;
  return (
    hasVitestEnv && (markers.tinypoolState !== undefined || markers.vitestWorker !== undefined)
  );
}

export function requestExitAfterOneShotOutput(
  runtime: RuntimeEnv = defaultRuntime,
  exitCode = 0,
): boolean {
  if (runtime !== defaultRuntime) {
    return false;
  }
  requestedExitCode = exitCode;
  return true;
}

export function flushExitAfterOneShotOutput(
  runtime: RuntimeEnv = defaultRuntime,
  env: NodeJS.ProcessEnv = process.env,
  markers: VitestWorkerMarkers = resolveVitestWorkerMarkers(),
): void {
  const exitCode = requestedExitCode;
  requestedExitCode = undefined;
  if (exitCode === undefined || runtime !== defaultRuntime || isVitestWorker(env, markers)) {
    return;
  }

  const exit = () => runtime.exit(exitCode);
  let pendingStreams = 2;

  const drain = (stream: NodeJS.WriteStream) => {
    stream.write("", () => {
      pendingStreams -= 1;
      if (pendingStreams === 0) {
        setImmediate(exit);
      }
    });
  };

  drain(process.stdout);
  drain(process.stderr);
}