bot-agents/src/agent-codex.ts

// Signs Codex in for one profile, runs it, and returns its result file path.
import process from "node:process";
import type { Readable } from "node:stream";
import { text } from "node:stream/consumers";
import { setTimeout as sleep } from "node:timers/promises";

import { isCodexAuthFailure, retryInvalidToken } from "./auth.ts";
import type { CodexAuthentication } from "./auth.ts";
import { chunks, endWith, writeChunk } from "./command.ts";
import type { CommandStatus } from "./command.ts";
import { agentCommand } from "./isolation.ts";
import { AgentAuthError, agentEnvironment, withProfile } from "./profiles.ts";
import {
  makeProfileFile,
  postAuthorComment,
  PROCESS_ENV,
  runAgentCommand,
  WORKSPACE,
} from "./run-context.ts";
import type { RuntimeProfile } from "./runtime.ts";

interface CodexAttempt {
  status: CommandStatus;
  error: string;
}

async function collectText(stream: Readable, append: (chunk: string) => void): Promise<void> {
  const decoder = new TextDecoder();
  for await (const chunk of chunks(stream)) {
    append(decoder.decode(chunk, { stream: true }));
  }
  append(decoder.decode());
}

// Polls every 100 ms until the login prints something, failing if it exits first.
async function pollForInstructions(
  shownText: () => string,
  status: Promise<CommandStatus>,
): Promise<void> {
  if (shownText().trim() !== "") {
    return;
  }
  const exited = await Promise.race([status.then(() => true), sleep(100, false)]);
  if (exited) {
    throw new AgentAuthError("AI bot login produced no instructions");
  }
  await pollForInstructions(shownText, status);
}

async function codexDeviceLogin(
  profile: RuntimeProfile,
  environment: Record<string, string>,
): Promise<void> {
  const login = agentCommand(profile, "codex", ["login", "--device-auth"], {
    cwd: WORKSPACE,
    env: environment,
    clearEnv: true,
    stdout: "piped",
    stderr: "piped",
  }).spawn();
  let shown = "";
  const append = (chunk: string): void => {
    shown += chunk;
  };
  const drained = Promise.all([
    collectText(login.stdout, append),
    collectText(login.stderr, append),
  ]);
  await sleep(3000);
  await pollForInstructions(() => shown, login.status);
  await postAuthorComment(withProfile(shown, profile));
  const [result] = await Promise.all([login.status, drained]);
  if (!result.success) {
    throw new AgentAuthError(`codex login exited with ${result.code}`);
  }
}

async function codexApiLogin(
  profile: RuntimeProfile,
  environment: Record<string, string>,
  token: string,
): Promise<void> {
  const login = agentCommand(profile, "codex", ["login", "--with-api-key"], {
    cwd: WORKSPACE,
    env: environment,
    clearEnv: true,
    stdin: "piped",
    stderr: "piped",
  }).spawn();
  await endWith(login.stdin, `${token}\n`);
  const [status, error] = await Promise.all([login.status, text(login.stderr)]);
  if (!status.success) {
    throw new Error(error.trim() || `codex login exited with ${status.code}`);
  }
}

async function signIn(
  profile: RuntimeProfile,
  environment: Record<string, string>,
  authentication: CodexAuthentication,
): Promise<void> {
  if (authentication.kind === "api-key") {
    await codexApiLogin(profile, environment, authentication.token);
    return;
  }
  const loggedIn = await agentCommand(profile, "codex", ["login", "status"], {
    cwd: WORKSPACE,
    env: environment,
    clearEnv: true,
  }).output();
  if (!loggedIn.success) {
    await codexDeviceLogin(profile, environment);
  }
}

async function execCodex(
  profile: RuntimeProfile,
  environment: Record<string, string>,
  prompt: string,
  outputPath: string,
): Promise<CodexAttempt> {
  const args = [
    "exec",
    "--model",
    profile.model,
    "--dangerously-bypass-approvals-and-sandbox",
    "--output-last-message",
    outputPath,
    prompt,
  ];
  const codex = agentCommand(profile, "codex", args, {
    cwd: WORKSPACE,
    env: environment,
    clearEnv: true,
    stdout: "inherit",
    stderr: "piped",
  }).spawn();
  const decoder = new TextDecoder();
  let error = "";
  for await (const chunk of chunks(codex.stderr)) {
    await writeChunk(process.stderr, chunk);
    error += decoder.decode(chunk, { stream: true });
  }
  error += decoder.decode();
  // Codex echoes the prompt, whose comments may quote auth errors.
  return { status: await codex.status, error: error.replaceAll(prompt, "") };
}

export async function runCodex(
  profile: RuntimeProfile,
  prompt: string,
  authentication: CodexAuthentication,
): Promise<string> {
  const token = authentication.kind === "api-key" ? authentication.token : undefined;
  const environment = agentEnvironment(PROCESS_ENV, profile, token);
  await signIn(profile, environment, authentication);
  const outputPath = await makeProfileFile(profile, "result.md");
  const attempt = (): Promise<CodexAttempt> => execCodex(profile, environment, prompt, outputPath);
  const result =
    authentication.kind === "api-key"
      ? await attempt()
      : await retryInvalidToken(attempt, async () => {
          await runAgentCommand(profile, "codex", ["logout"]);
          await codexDeviceLogin(profile, environment);
        });
  if (!result.status.success) {
    const failure = new Error(result.error.trim() || `codex exited with ${result.status.code}`);
    if (isCodexAuthFailure(result.error)) {
      throw new AgentAuthError("codex sign-in failed", { cause: failure });
    }
    throw failure;
  }
  return outputPath;
}