pikasTech-HWLAB/internal/cloud/codex-stdio-session-turn-state.ts

import { CODEX_STDIO_FIRST_TOKEN_PROGRESS_MS } from "./codex-stdio-session.ts";
import {
  appServerActivityTimeoutError,
  appServerCommandStatus,
  appServerHardTimeoutError,
  appServerTerminalStatus,
  commandExecutionCommandText,
  commandExecutionErrorText,
  commandExecutionOutputText,
  extractAppServerRecord,
  extractAppServerString,
  firstNonEmpty,
  optionalId,
  redactText,
  tailText
} from "./codex-stdio-session-helpers.ts";

export function createAppServerTurnState({ traceRecorder, session } = {}) {
  let threadId = optionalId(session?.threadId);
  let turnId = null;
  let assistantText = "";
  let finalResponse = "";
  let terminal = null;
  let terminalResolve;
  const terminalPromise = new Promise((resolve) => {
    terminalResolve = resolve;
  });
  let lastActivityAt = Date.now();
  let lastActivityLabel = "turn-state:created";
  let lastWaitingFor = "app-server-turn";

  function setThreadId(value) {
    threadId = optionalId(value) ?? threadId;
  }

  function setTurnId(value) {
    turnId = optionalId(value) ?? turnId;
  }

  function activitySnapshot(referenceNow = Date.now()) {
    return {
      lastActivityAt: new Date(lastActivityAt).toISOString(),
      lastActivityLabel,
      idleMs: Math.max(0, referenceNow - lastActivityAt),
      waitingFor: lastWaitingFor
    };
  }

  function observeActivity({ label = "app-server:activity", waitingFor = null } = {}) {
    lastActivityAt = Date.now();
    lastActivityLabel = label;
    if (waitingFor) lastWaitingFor = waitingFor;
    return activitySnapshot(lastActivityAt);
  }

  function appendTrace(event = {}) {
    const appended = traceRecorder?.append(event);
    observeActivity({
      label: appended?.label ?? event.label ?? event.type ?? "trace:event",
      waitingFor: appended?.waitingFor ?? event.waitingFor ?? null
    });
    return appended;
  }

  function appendProgressTrace(event = {}) {
    return traceRecorder?.append(event);
  }

  function hasAssistantOutput() {
    return Boolean(firstNonEmpty(finalResponse, assistantText));
  }

  function appendFirstTokenProgressTrace(referenceNow = Date.now()) {
    const activity = activitySnapshot(referenceNow);
    if (hasAssistantOutput() || !turnId || terminal) return null;
    if (!["assistant-message", "turn/completed"].includes(activity.waitingFor)) return null;
    return appendProgressTrace({
      type: "turn",
      status: "running",
      label: "turn:waiting:first_assistant_token",
      message: "Codex app-server is waiting for the first assistant token; this progress trace does not reset the backend idle timer.",
      progressOnly: true,
      idleMs: activity.idleMs,
      lastActivityAt: activity.lastActivityAt,
      sessionId: session?.sessionId,
      sessionStatus: session?.status,
      turn: session?.turn,
      threadId,
      turnId,
      waitingFor: "assistant-message"
    });
  }

  function finish(status, error = null, extra = {}) {
    if (terminal) return;
    terminal = {
      terminalStatus: appServerTerminalStatus(status),
      terminalError: error ? redactText(error) : null,
      threadId,
      turnId,
      ...activitySnapshot(),
      ...extra
    };
    terminalResolve(terminal);
  }

  function handle(message) {
    const method = typeof message?.method === "string" ? message.method : "unknown";
    observeActivity({
      label: `app-server:${method}`,
      waitingFor: appServerWaitingForMethod(method)
    });
    const params = extractAppServerRecord(message?.params);
    const item = extractAppServerRecord(params?.item);
    const turn = extractAppServerRecord(params?.turn);
    if (method === "thread/started") {
      setThreadId(extractAppServerString(extractAppServerRecord(params?.thread), "id"));
      appendTrace({
        type: "thread",
        status: "completed",
        label: "thread:started",
        sessionId: session?.sessionId,
        sessionStatus: session?.status,
        turn: session?.turn,
        threadId,
        waitingFor: "turn/start"
      });
      return;
    }
    if (method === "turn/started") {
      setTurnId(extractAppServerString(turn, "id"));
      appendTrace({
        type: "turn",
        status: "started",
        label: "turn:started",
        sessionId: session?.sessionId,
        sessionStatus: session?.status,
        turn: session?.turn,
        threadId,
        turnId,
        waitingFor: "assistant-message"
      });
      return;
    }
    if (method === "item/agentMessage/delta") {
      const delta = String(params?.delta ?? "");
      assistantText += delta;
      if (delta) {
        traceRecorder?.appendAssistantDelta?.({
          itemId: optionalId(params?.itemId),
          chunk: delta,
          sessionId: session?.sessionId,
          sessionStatus: session?.status,
          turn: session?.turn,
          threadId,
          turnId,
          waitingFor: "turn/completed"
        });
      }
      return;
    }
    if (method === "item/completed" && item?.type === "agentMessage") {
      if (typeof item.text === "string") finalResponse = item.text;
      appendTrace({
        type: "assistant_message",
        status: "completed",
        label: "assistant:item_completed",
        sessionId: session?.sessionId,
        sessionStatus: session?.status,
        turn: session?.turn,
        threadId,
        turnId,
        waitingFor: "turn/completed"
      });
      return;
    }
    if (method === "item/started" && item?.type === "commandExecution") {
      appendCommandExecutionTrace(item, {
        status: "started",
        label: "item/commandExecution:started"
      });
      return;
    }
    if (method === "item/completed" && item?.type === "commandExecution") {
      appendCommandExecutionTrace(item, {
        status: appServerCommandStatus(item),
        label: "item/commandExecution:completed"
      });
      return;
    }
    if (method === "item/commandExecution/outputDelta" || method === "item/reasoning/summaryTextDelta" || method === "item/reasoning/textDelta") {
      appendTrace({
        type: method.startsWith("item/reasoning/") ? "reasoning" : "tool_call",
        status: "output_chunk",
        label: method,
        outputSummary: String(params?.delta ?? "").slice(0, 400),
        sessionId: session?.sessionId,
        sessionStatus: session?.status,
        turn: session?.turn,
        threadId,
        turnId,
        waitingFor: "turn/completed"
      });
      return;
    }
    if (method === "error") {
      const error = extractAppServerRecord(params?.error);
      const message = typeof error?.message === "string" ? error.message : "Codex app-server error";
      const willRetry = params?.willRetry === true;
      setThreadId(extractAppServerString(params, "threadId"));
      setTurnId(extractAppServerString(params, "turnId"));
      appendTrace({
        type: willRetry ? "provider_retry" : "error",
        status: willRetry ? "retrying" : "failed",
        label: willRetry ? "app-server:retrying" : "app-server:error",
        errorCode: willRetry ? "codex_stdio_provider_retry" : "codex_stdio_failed",
        message: redactText(message),
        sessionId: session?.sessionId,
        sessionStatus: session?.status,
        turn: session?.turn,
        threadId,
        turnId,
        waitingFor: willRetry ? "turn/completed" : "user-retry",
        terminal: !willRetry
      });
      if (!willRetry) finish("failed", message);
      return;
    }
    if (method === "turn/completed") {
      const status = appServerTerminalStatus(turn?.status);
      const error = extractAppServerRecord(turn?.error);
      const message = typeof error?.message === "string" ? error.message : null;
      appendTrace({
        type: "turn",
        status,
        label: `turn:completed:${status ?? "unknown"}`,
        sessionId: session?.sessionId,
        sessionStatus: session?.status,
        turn: session?.turn,
        threadId,
        turnId,
        errorCode: status === "completed" ? null : "codex_stdio_failed",
        message: message ? redactText(message) : undefined,
        terminal: status !== "completed"
      });
      finish(status, message);
    }
  }

  function appendCommandExecutionTrace(item, { status, label }) {
    const output = commandExecutionOutputText(item);
    appendTrace({
      type: "tool_call",
      stage: "tool_call",
      status,
      label,
      itemId: optionalId(item.id),
      toolName: "commandExecution",
      command: redactText(commandExecutionCommandText(item)),
      exitCode: Number.isInteger(item.exitCode) ? item.exitCode : undefined,
      durationMs: typeof item.durationMs === "number" ? item.durationMs : undefined,
      outputBytes: output.length,
      stdoutSummary: output ? tailText(redactText(output), 600) : undefined,
      stderrSummary: commandExecutionErrorText(item),
      sessionId: session?.sessionId,
      sessionStatus: session?.status,
      turn: session?.turn,
      threadId,
      turnId,
      waitingFor: status === "completed" ? "turn/completed" : "commandExecution/completed"
    });
  }

  async function wait(timeoutMs, closedPromise, { hardTimeoutMs = null } = {}) {
    let activityTimer;
    let hardTimer;
    let progressTimer;
    const timeoutPromise = new Promise((_, reject) => {
      const checkActivity = () => {
        const now = Date.now();
        const activity = activitySnapshot(now);
        if (activity.idleMs >= timeoutMs) {
          const error = appServerActivityTimeoutError({
            timeoutMs,
            idleMs: activity.idleMs,
            lastActivityAt: activity.lastActivityAt,
            waitingFor: activity.waitingFor,
            threadId,
            turnId,
            partialAssistant: hasAssistantOutput()
          });
          appendTrace({
            type: "timeout",
            status: "failed",
            label: "timeout:no_activity",
            errorCode: error.code,
            message: error.message,
            timeoutMs,
            idleMs: activity.idleMs,
            lastActivityAt: activity.lastActivityAt,
            sessionId: session?.sessionId,
            sessionStatus: session?.status,
            turn: session?.turn,
            threadId,
            turnId,
            waitingFor: activity.waitingFor,
            terminal: true
          });
          reject(error);
          return;
        }
        activityTimer = setTimeout(checkActivity, Math.max(1, timeoutMs - activity.idleMs));
      };
      activityTimer = setTimeout(checkActivity, timeoutMs);
    });
    const progressPromise = new Promise(() => {
      const tick = () => {
        appendFirstTokenProgressTrace();
        progressTimer = setTimeout(tick, CODEX_STDIO_FIRST_TOKEN_PROGRESS_MS);
      };
      progressTimer = setTimeout(tick, CODEX_STDIO_FIRST_TOKEN_PROGRESS_MS);
    });
    const hardTimeoutPromise = hardTimeoutMs
      ? new Promise((_, reject) => {
          hardTimer = setTimeout(() => {
            const activity = activitySnapshot();
            const error = appServerHardTimeoutError({
              hardTimeoutMs,
              timeoutMs,
              idleMs: activity.idleMs,
              lastActivityAt: activity.lastActivityAt,
              waitingFor: activity.waitingFor,
              threadId,
              turnId
            });
            appendTrace({
              type: "timeout",
              status: "failed",
              label: "timeout:hard_cap",
              errorCode: error.code,
              message: error.message,
              timeoutMs,
              hardTimeoutMs,
              idleMs: activity.idleMs,
              lastActivityAt: activity.lastActivityAt,
              sessionId: session?.sessionId,
              sessionStatus: session?.status,
              turn: session?.turn,
              threadId,
              turnId,
              waitingFor: activity.waitingFor,
              terminal: true
            });
            reject(error);
          }, hardTimeoutMs);
        })
      : null;
    const candidates = [terminalPromise];
    if (closedPromise && typeof closedPromise.then === "function") {
      candidates.push(closedPromise.then((exit) => {
        if (terminal) return terminal;
        return {
          terminalStatus: "failed",
          terminalError: hasAssistantOutput()
            ? "Codex app-server transport closed after partial assistant output but before turn/completed."
            : "Codex app-server transport closed before turn/completed.",
          threadId,
          turnId,
          ...activitySnapshot(),
          appServerExit: exit
        };
      }));
    }
    try {
      return await Promise.race([...candidates, timeoutPromise, hardTimeoutPromise, progressPromise].filter(Boolean));
    } finally {
      clearTimeout(activityTimer);
      clearTimeout(hardTimer);
      clearTimeout(progressTimer);
    }
  }

  function snapshot() {
    const activity = activitySnapshot();
    return {
      threadId,
      turnId,
      assistantText: redactText(assistantText),
      finalResponse: redactText(firstNonEmpty(finalResponse, assistantText)),
      terminalStatus: terminal?.terminalStatus ?? null,
      terminalError: terminal?.terminalError ?? null,
      lastActivityAt: terminal?.lastActivityAt ?? activity.lastActivityAt,
      idleMs: terminal?.idleMs ?? activity.idleMs,
      waitingFor: terminal?.waitingFor ?? activity.waitingFor
    };
  }

  return {
    appendTrace,
    handle,
    setThreadId,
    setTurnId,
    wait,
    snapshot
  };
}

function appServerWaitingForMethod(method) {
  if (method === "thread/started") return "turn/start";
  if (method === "turn/started") return "assistant-message";
  if (method === "item/agentMessage/delta") return "turn/completed";
  if (method === "item/completed") return "turn/completed";
  if (method === "item/commandExecution/outputDelta") return "turn/completed";
  if (method === "item/reasoning/summaryTextDelta" || method === "item/reasoning/textDelta") return "turn/completed";
  if (method === "error") return "turn/completed";
  if (method === "turn/completed") return "turn/completed";
  return "app-server-notification";
}