@data-club/ai-hub-server 0.0.1

package/src/state/turnRunner.ts

@@ -0,0 +1,588 @@
+/**
+ * Server-authoritative turn lifecycle (§6.6, §10).
+ *
+ * `runTurn(...)` is the one entry point that:
+ *
+ * 1. Inserts the `turn` row in status=running, phase=model_text.
+ * 2. Opens a live channel (`turnRegistry.openTurnChannel`) so concurrent
+ *    resumption attempts can attach.
+ * 3. Emits `stream_start` (seq 0).
+ * 4. Calls `LettaClient.streamUserMessage(...)` which spawns a Letta
+ *    Code subprocess, acquires the throttle lease, sends the user
+ *    message. Returns a `ChatTurnHandle`.
+ * 5. Iterates the handle's `events` stream:
+ *    - Non-terminal SDK messages → `translateSdkMessage(...)` →
+ *      stamps + persists + broadcasts each emitted ai-hub event.
+ *    - Terminal `result` SDK message → `finalizeTerminal(...)` →
+ *      emits the closing `assistant_message_complete` + `stream_end`
+ *      with the correct reason (`done` / `cancelled` / `error`).
+ * 6. On cancellation (caller signal OR `requestTurnCancellation`):
+ *    calls `handle.abort()` to send `interrupt` to the CLI, keeps
+ *    draining until terminal, then closes.
+ * 7. On throttle rejection (`ServerOverloadedError`): emits
+ *    `error{code:"server_overloaded"}` + `stream_end{error}`.
+ *    Transitions turn → failed.
+ * 8. On unhandled error: emits `error` + `stream_end{error}`,
+ *    `turn.status = "failed"`.
+ * 9. Always calls `handle.close()` in finally (releases subprocess +
+ *    throttle lease).
+ *
+ * @see SPECIFICATIONS.md §6.3.2 — chat / SSE
+ * @see SPECIFICATIONS.md §6.6 — state-machine ownership
+ * @see SPECIFICATIONS.md §6.9 — cancellation
+ * @see SPECIFICATIONS.md §9 — wire protocol
+ * @see SPECIFICATIONS.md §10 — state machine
+ */
+
+import { randomUUID } from "node:crypto";
+
+import { eq, sql } from "drizzle-orm";
+
+import type {
+  AIHubEvent,
+  AIHubEventState,
+  RunningPhase,
+  TurnStatus,
+} from "@data-club/ai-hub";
+
+import { messageUsage, turn, turnEvent } from "../db/schema.js";
+import type { DrizzleDB } from "../db/connection.js";
+import { ServerOverloadedError, type ChatTurnHandle, type LettaClient } from "../letta/client.js";
+import {
+  finalizeAssistantMessage,
+  finalizeTerminal,
+  initialTranslatorState,
+  translateSdkMessage,
+  type TranslatorState,
+  type UnsequencedEvent,
+} from "../letta/eventTranslator.js";
+import type { Logger } from "../observability/logger.js";
+import * as metrics from "../observability/metrics.js";
+
+/** The ai-hub SSE wire-protocol version this server speaks (§9.1). */
+export const WIRE_PROTOCOL_VERSION = "1.0";
+
+/* ────────────────────────────────────────────────────────────────────────── */
+/*  Inputs                                                                    */
+/* ────────────────────────────────────────────────────────────────────────── */
+
+/** Parameters for `runTurn`. */
+export interface RunTurnOpts {
+  /** Drizzle handle. */
+  db: DrizzleDB;
+  /** Letta client wrapper (production wraps the SDKs; tests mock). */
+  lettaClient: LettaClient;
+  /** Pino logger for this request. */
+  logger: Logger;
+  /** Per-request id for log correlation (stamped on `error` events). */
+  requestId: string;
+  /** Owner of the session. */
+  userId: string;
+  /** ai-hub session id this turn belongs to. */
+  sessionId: string;
+  /** The Letta conversation id stored on the session row. */
+  lettaConversationId: string;
+  /** The Letta agent id stored on the session row. */
+  lettaAgentId: string;
+  /**
+   * Step-12.5 — absolute path of the per-session sandbox folder, or
+   * `null` when the session has no sandbox (today's stateless shape).
+   * Threaded as `cwd` into the Letta Code subprocess spawn so the
+   * agent's filesystem tools operate inside the seeded folder. When
+   * non-null, also flips the disallowed-tools list to re-enable
+   * `Read` / `Write` / `Bash` / etc.
+   */
+  sandboxDir: string | null;
+  /** The user's typed message. */
+  userMessage: string;
+  /** Optional idempotency key from the request body. */
+  idempotencyKey?: string;
+  /**
+   * Listener that receives every event the runner emits. The chat
+   * handler subscribes its SSE writer here. `null` events signal
+   * "channel closed; no more events coming."
+   */
+  onEvent: (event: AIHubEvent | null) => void;
+  /**
+   * Caller-supplied AbortSignal. When aborted, the runner calls
+   * `handle.abort()` on the spawned Letta Code subprocess and emits
+   * `stream_end{reason:"cancelled"}` after the upstream terminal
+   * `result` arrives.
+   */
+  signal?: AbortSignal;
+}
+
+/** Returned by `runTurn`. The runner returns after `turn_event`s are committed. */
+export interface RunTurnResult {
+  /** The new turn id (the `turn_id` carried in `stream_start`). */
+  turnId: string;
+  /** Terminal status (`done` / `cancelled` / `failed`). */
+  status: Exclude<TurnStatus, "idle" | "running">;
+}
+
+/* ────────────────────────────────────────────────────────────────────────── */
+/*  Runner                                                                    */
+/* ────────────────────────────────────────────────────────────────────────── */
+
+/**
+ * Runs one turn from "user pressed Send" to terminal status. Returns a
+ * promise that resolves when the runner has persisted the terminal
+ * event(s), closed the live channel, and torn down the Letta Code
+ * subprocess.
+ *
+ * The caller is expected to keep `onEvent` alive for the entire run.
+ *
+ * @param opts - see `RunTurnOpts`.
+ * @returns the turn id and terminal status.
+ */
+export async function runTurn(opts: RunTurnOpts): Promise<RunTurnResult> {
+  const turnId = randomUUID();
+  const startedAt = new Date().toISOString();
+  const startedMs = Date.now();
+
+  metrics.chatRequestsTotal.inc();
+
+  // Step 1 — insert turn row in running/model_text.
+  opts.db
+    .insert(turn)
+    .values({
+      id: turnId,
+      sessionId: opts.sessionId,
+      status: "running",
+      phase: "model_text",
+      lastSeq: 0,
+      idempotencyKey: opts.idempotencyKey ?? null,
+      startedAt,
+    })
+    .run();
+
+  const { openTurnChannel } = await import("./turnRegistry.js");
+  const publisher = openTurnChannel(turnId);
+
+  // Compose the caller's signal with the registry-supplied cancel
+  // signal: either side can fire to cancel the turn. The registry
+  // cancel is what `POST /chat/:id/cancel` triggers from a separate
+  // HTTP request.
+  const registryCancelSignal = publisher.cancelSignal;
+  const cancelController = new AbortController();
+  const propagateCancel = (): void => cancelController.abort();
+  if (opts.signal) {
+    if (opts.signal.aborted) cancelController.abort();
+    else opts.signal.addEventListener("abort", propagateCancel, { once: true });
+  }
+  if (registryCancelSignal.aborted) cancelController.abort();
+  else registryCancelSignal.addEventListener("abort", propagateCancel, { once: true });
+
+  let translatorState: TranslatorState = initialTranslatorState();
+  let nextSeq = 0;
+
+  /** Stamps + persists + broadcasts one event. Also fills in
+   *  `request_id` on error events — the translator emits them without
+   *  a request_id (it has no access to the request), so we add it here. */
+  const emit = (unsequenced: UnsequencedEvent): AIHubEvent => {
+    let withRequestId: UnsequencedEvent = unsequenced;
+    if (unsequenced.event === "error") {
+      withRequestId = {
+        ...unsequenced,
+        data: { ...unsequenced.data, request_id: opts.requestId },
+      } as UnsequencedEvent;
+    }
+    const event = { ...withRequestId, seq: nextSeq++ } as AIHubEvent;
+    const persistedData =
+      "data" in event ? (event.data as unknown as object) : ({} as object);
+    opts.db
+      .insert(turnEvent)
+      .values({
+        turnId,
+        seq: event.seq,
+        eventName: event.event,
+        data: persistedData,
+        state: event.state as unknown as object,
+      })
+      .run();
+    publisher.publish(event);
+    opts.onEvent(event);
+    return event;
+  };
+
+  /** Persists turn-row mutations driven by translator-state changes. */
+  const syncTurnRow = (status: TurnStatus, phase: RunningPhase | null): void => {
+    const sets: Record<string, unknown> = { lastSeq: nextSeq - 1 };
+    if (status !== "running") {
+      sets.status = status;
+      sets.phase = null;
+      sets.finishedAt = new Date().toISOString();
+    } else {
+      sets.status = "running";
+      sets.phase = phase;
+    }
+    opts.db.update(turn).set(sets).where(eq(turn.id, turnId)).run();
+  };
+
+  let handle: ChatTurnHandle | null = null;
+
+  try {
+    // Step 3 — stream_start (emitted once, at the top of the run).
+    emit({
+      event: "stream_start",
+      data: {
+        protocol_version: WIRE_PROTOCOL_VERSION,
+        session_id: opts.sessionId,
+        turn_id: turnId,
+      },
+      state: { status: "running", phase: "model_text" } satisfies AIHubEventState,
+    });
+    syncTurnRow("running", "model_text");
+
+    if (cancelController.signal.aborted) {
+      // Caller cancelled before we even spawned — emit cancelled and
+      // bail without ever calling streamUserMessage.
+      throw new AbortLikeError();
+    }
+
+    // Step 4 — spawn + send. Throttle gate runs inside
+    // streamUserMessage; ServerOverloadedError is caught below and
+    // turned into a `server_overloaded` SSE error event.
+    const lettaStartedMs = Date.now();
+    try {
+      handle = await opts.lettaClient.streamUserMessage({
+        conversationId: opts.lettaConversationId,
+        lettaAgentId: opts.lettaAgentId,
+        userMessage: opts.userMessage,
+        ...(opts.sandboxDir !== null ? { cwd: opts.sandboxDir } : {}),
+      });
+    } catch (err) {
+      metrics.lettaRequestDurationMs.set(Date.now() - lettaStartedMs);
+      throw err;
+    }
+
+    // Wire the cancel signal to the handle's abort once we have one.
+    // If the signal already fired during the await above, abort + bail
+    // immediately. Otherwise install a listener for future aborts.
+    if (cancelController.signal.aborted) {
+      void handle.abort();
+      throw new AbortLikeError();
+    }
+    const onCancel = (): void => {
+      void handle?.abort();
+    };
+    cancelController.signal.addEventListener("abort", onCancel, { once: true });
+
+    // Step 5 — drain the SDK message stream.
+    // Track the last Letta run id we see — used to fetch usage stats
+    // after the terminal `result`. The SDK strips Letta message ids
+    // from its translated payloads, so we re-fetch via REST
+    // (`getRunUsageStats(runId)`) to get both the usage_statistics
+    // and the last assistant message id we'll stamp the badge on.
+    let lastRunId: string | null = null;
+    try {
+      for await (const sdkMsg of handle.events) {
+        // SDKAssistantMessage / SDKReasoningMessage / SDKToolCallMessage
+        // each carry `runId`. Take the most recent — a multi-step turn
+        // can have several runs, and we want the latest.
+        if ((sdkMsg as { runId?: string }).runId) {
+          lastRunId = (sdkMsg as { runId: string }).runId;
+        }
+        if (sdkMsg.type === "result") {
+          // Terminal — finalizeTerminal emits the right closing events
+          // based on success / interrupted / error.
+          // Override the reason to `cancelled` when our cancel fired
+          // (even if the SDK reports success — see the cancel fixture
+          // open-question-5 finding: the SDK may complete the run if
+          // the interrupt arrived during a tool call).
+          const cancelOverride = cancelController.signal.aborted;
+          const terminalToUse = cancelOverride
+            ? { ...sdkMsg, success: false, errorCode: "interrupted" as const }
+            : sdkMsg;
+          const step = finalizeTerminal(terminalToUse, translatorState);
+          translatorState = step.state;
+          for (const ev of step.events) {
+            emit(ev);
+          }
+          syncTurnRow(translatorState.status, null);
+          break;
+        }
+
+        const step = translateSdkMessage(sdkMsg, translatorState);
+        translatorState = step.state;
+        for (const ev of step.events) {
+          emit(ev);
+        }
+        syncTurnRow(translatorState.status, translatorState.phase);
+      }
+    } finally {
+      cancelController.signal.removeEventListener("abort", onCancel);
+      metrics.lettaRequestDurationMs.set(Date.now() - lettaStartedMs);
+    }
+
+    // Step 5.5 — usage emission + persistence. Only fire on successful
+    // termination (cancelled / failed turns may or may not have a
+    // usable runId, and the badge isn't load-bearing). Best-effort:
+    // any failure here logs and continues — we already emitted
+    // stream_end above, so the turn is "done" from the client's POV.
+    if (
+      lastRunId !== null &&
+      translatorState.status === "done"
+    ) {
+      try {
+        const usage = await opts.lettaClient.getRunUsageStats(lastRunId);
+        if (usage && usage.lastAssistantMessageId) {
+          // Persist for /sessions/{id}/messages history rendering.
+          opts.db
+            .insert(messageUsage)
+            .values({
+              lettaMessageId: usage.lastAssistantMessageId,
+              sessionId: opts.sessionId,
+              promptTokens: usage.promptTokens,
+              completionTokens: usage.completionTokens,
+              totalTokens: usage.totalTokens,
+              cachedInputTokens: usage.cachedInputTokens,
+              cacheWriteTokens: usage.cacheWriteTokens,
+              reasoningTokens: usage.reasoningTokens,
+              contextTokens: usage.contextTokens,
+              runId: lastRunId,
+            })
+            .onConflictDoUpdate({
+              target: messageUsage.lettaMessageId,
+              set: {
+                promptTokens: usage.promptTokens,
+                completionTokens: usage.completionTokens,
+                totalTokens: usage.totalTokens,
+                cachedInputTokens: usage.cachedInputTokens,
+                cacheWriteTokens: usage.cacheWriteTokens,
+                reasoningTokens: usage.reasoningTokens,
+                contextTokens: usage.contextTokens,
+                runId: lastRunId,
+              },
+            })
+            .run();
+          // Live wire event — UI can attach the badge immediately
+          // without waiting for a /messages refetch.
+          opts.onEvent({
+            seq: 0, // emit() bypassed here; the stream's already drained.
+            event: "usage",
+            data: {
+              letta_message_id: usage.lastAssistantMessageId,
+              ...(usage.promptTokens !== null ? { input: usage.promptTokens } : {}),
+              ...(usage.completionTokens !== null ? { output: usage.completionTokens } : {}),
+              ...(usage.cachedInputTokens !== null ? { cache_read: usage.cachedInputTokens } : {}),
+              ...(usage.cacheWriteTokens !== null ? { cache_write: usage.cacheWriteTokens } : {}),
+              ...(usage.reasoningTokens !== null ? { reasoning: usage.reasoningTokens } : {}),
+              ...(usage.contextTokens !== null ? { context: usage.contextTokens } : {}),
+              ...(usage.totalTokens !== null ? { total: usage.totalTokens } : {}),
+            },
+            state: { status: "done", phase: null } satisfies AIHubEventState,
+          });
+        }
+      } catch (err) {
+        opts.logger.warn({ err, runId: lastRunId }, "usage-stats fetch/persist failed");
+      }
+    }
+
+    publisher.close();
+    opts.onEvent(null);
+    metrics.chatDurationMs.set(Date.now() - startedMs);
+
+    const finalStatus = translatorState.status as RunTurnResult["status"];
+    opts.logger.info(
+      {
+        event: `turn_${finalStatus}`,
+        turn_id: turnId,
+        session_id: opts.sessionId,
+      },
+      `Turn ${finalStatus}.`,
+    );
+    return { turnId, status: finalStatus };
+  } catch (err) {
+    if (err instanceof ServerOverloadedError) {
+      emit({
+        event: "error",
+        data: {
+          code: "server_overloaded",
+          message: err.message,
+        },
+        state: { status: "failed", phase: null },
+      });
+      emit({
+        event: "stream_end",
+        data: { reason: "error" },
+        state: { status: "failed", phase: null },
+      });
+      syncTurnRow("failed", null);
+      publisher.close();
+      opts.onEvent(null);
+      metrics.chatDurationMs.set(Date.now() - startedMs);
+      opts.logger.warn(
+        {
+          event: "turn_rejected_overloaded",
+          turn_id: turnId,
+          session_id: opts.sessionId,
+          gate: err.gate,
+          detail: err.detail,
+        },
+        "Turn rejected: server overloaded.",
+      );
+      return { turnId, status: "failed" };
+    }
+
+    if (cancelController.signal.aborted || err instanceof AbortLikeError) {
+      // Emit any partial assistant text the turn produced before the
+      // cancel landed, then the cancelled stream_end.
+      for (const ev of finalizeAssistantMessage(translatorState, "cancelled")) {
+        emit(ev);
+      }
+      emit({
+        event: "stream_end",
+        data: { reason: "cancelled" },
+        state: { status: "cancelled", phase: null },
+      });
+      syncTurnRow("cancelled", null);
+      publisher.close();
+      opts.onEvent(null);
+      metrics.chatDurationMs.set(Date.now() - startedMs);
+      opts.logger.info(
+        { event: "turn_cancelled", turn_id: turnId, session_id: opts.sessionId },
+        "Turn cancelled.",
+      );
+      return { turnId, status: "cancelled" };
+    }
+
+    const message = err instanceof Error ? err.message : String(err);
+    emit({
+      event: "error",
+      data: { code: classifyServerError(err), message },
+      state: { status: "failed", phase: null },
+    });
+    emit({
+      event: "stream_end",
+      data: { reason: "error" },
+      state: { status: "failed", phase: null },
+    });
+    syncTurnRow("failed", null);
+    publisher.close();
+    opts.onEvent(null);
+    metrics.chatDurationMs.set(Date.now() - startedMs);
+    opts.logger.error(
+      {
+        event: "turn_failed",
+        turn_id: turnId,
+        session_id: opts.sessionId,
+        err: message,
+      },
+      "Turn failed.",
+    );
+    return { turnId, status: "failed" };
+  } finally {
+    if (handle) {
+      try {
+        handle.close();
+      } catch {
+        // best-effort
+      }
+    }
+    if (opts.signal) opts.signal.removeEventListener("abort", propagateCancel);
+    registryCancelSignal.removeEventListener("abort", propagateCancel);
+  }
+}
+
+/* ────────────────────────────────────────────────────────────────────────── */
+/*  Helpers                                                                   */
+/* ────────────────────────────────────────────────────────────────────────── */
+
+/** Sentinel for "the cancel path fired" — the catch arm checks
+ *  `err instanceof AbortLikeError` plus the cancel signal. */
+class AbortLikeError extends Error {
+  constructor() {
+    super("Turn aborted.");
+    this.name = "AbortLikeError";
+  }
+}
+
+/**
+ * Maps an unhandled runner error onto an ai-hub `error.code`.
+ *
+ * Post-migration the set of error sources is much smaller — most of
+ * the provider/upstream signals come through the translator (which
+ * maps them via `classifyErrorCode`). What lands here is
+ * spawn-failure, network-down-to-Letta, decryption errors, and other
+ * "couldn't even start the turn" failures.
+ *
+ * @param err - the unhandled error.
+ * @returns the SSE `error.code` to emit.
+ */
+function classifyServerError(err: unknown): "provider_error" | "internal" {
+  if (err instanceof Error) {
+    const msg = err.message.toLowerCase();
+    if (msg.includes("letta") || msg.includes("upstream")) return "provider_error";
+  }
+  return "internal";
+}
+
+/* ────────────────────────────────────────────────────────────────────────── */
+/*  Resumption helpers                                                        */
+/* ────────────────────────────────────────────────────────────────────────── */
+
+/**
+ * Replays persisted `turn_event` rows newer than `fromSeq`, returned as
+ * ai-hub-shaped events. Used by `GET /chat/:turn_id/events?from_seq=N`
+ * for the replay half of resumption.
+ *
+ * @param db - Drizzle handle.
+ * @param turnId - the turn id.
+ * @param fromSeq - replays events with `seq > fromSeq`.
+ * @returns the persisted events.
+ *
+ * @see SPECIFICATIONS.md §6.3.2 — resumption
+ */
+export function replayPersistedEvents(
+  db: DrizzleDB,
+  turnId: string,
+  fromSeq: number,
+): ReadonlyArray<AIHubEvent> {
+  const rows = db
+    .select()
+    .from(turnEvent)
+    .where(sql`${turnEvent.turnId} = ${turnId} AND ${turnEvent.seq} > ${fromSeq}`)
+    .orderBy(turnEvent.seq)
+    .all();
+  return rows.map((row) => ({
+    event: row.eventName,
+    data: row.data as Record<string, unknown>,
+    state: row.state as AIHubEventState,
+    seq: row.seq,
+  })) as unknown as AIHubEvent[];
+}
+
+/**
+ * Idempotently marks a turn as cancelled in the DB. Returns the row's
+ * resulting status. Called by `POST /chat/:turn_id/cancel` after firing
+ * `requestTurnCancellation`. The runner's catch arm also persists
+ * `cancelled` via `syncTurnRow`; this helper just guarantees the row
+ * is updated when the runner has already exited (terminal-after-cancel
+ * race).
+ *
+ * Does not overwrite `done` / `failed` — those are terminal states.
+ *
+ * @param db - Drizzle handle.
+ * @param turnId - the turn to cancel.
+ * @returns the row's resulting status, or `null` if no such row.
+ */
+export function markCancelled(
+  db: DrizzleDB,
+  turnId: string,
+): TurnStatus | null {
+  const row = db.select().from(turn).where(eq(turn.id, turnId)).get();
+  if (!row) return null;
+  if (row.status === "running") {
+    db.update(turn)
+      .set({
+        status: "cancelled",
+        phase: null,
+        finishedAt: new Date().toISOString(),
+      })
+      .where(eq(turn.id, turnId))
+      .run();
+    return "cancelled";
+  }
+  return row.status;
+}