@opengeni/db 0.2.0

package/src/codex-token-resolver.ts

@@ -0,0 +1,247 @@
+// Shared, id-addressed, REFRESHING Codex token resolver (P2).
+//
+// Hoisted here from apps/worker/src/activities/codex-auth.ts so BOTH the worker
+// (turn-time bearer for the streamed run) AND the api (the /wham/usage quota-bar
+// reads) drive ONE resolver — no duplicated refresh/CAS/single-flight logic. The
+// worker re-exports buildCodexTokenResolver from this module for back-compat, so
+// the agent-turn.ts call site is unchanged.
+//
+// Why @opengeni/db is the right home: the resolver only orchestrates accessors
+// this package already owns (loadCodexCredentialForRun / recordCodexTokenRefresh /
+// setCodexCredentialStatus / encryptEnvironmentValue) plus pure @opengeni/codex
+// refresh helpers and the @opengeni/config key — keeping the refresh-CAS + RLS
+// invariants co-located with the rows they protect.
+//
+// CROSS-PROCESS SAFETY is preserved unchanged: the single-flight `inflight` map is
+// process-module-scoped, so worker and api each get their own — that is CORRECT
+// (each process coalesces its own concurrent refreshes). The real cross-process
+// guard is the (id, version) CAS inside recordCodexTokenRefresh: if the api
+// refreshes a token while a worker turn refreshes the same account, the loser's
+// CAS writes 0 rows (stale version) and it re-reads the winner's token, so the
+// one-time refresh token is never double-spent. RLS is untouched (every accessor
+// wraps withWorkspaceRls internally).
+
+import { environmentsEncryptionKeyBytes, type Settings } from "@opengeni/config";
+import {
+  accessTokenExpiry,
+  CODEX_CLIENT_VERSION,
+  CODEX_REFRESH_FALLBACK_MS,
+  CODEX_REFRESH_WINDOW_MS,
+  CodexReloginRequired,
+  type CodexTokenSnapshot,
+  type CodexUsagePayload,
+  fetchCodexUsage,
+  normalizeCodexUsage,
+  refreshCodexToken,
+} from "@opengeni/codex";
+import { encryptEnvironmentValue } from "./environment-crypto";
+import {
+  loadCodexCredentialForRun,
+  recordCodexAccountUsage,
+  recordCodexTokenRefresh,
+  setCodexCredentialStatus,
+  type CodexCredentialForRun,
+  type Database,
+} from "./index";
+
+// Single-flight per CREDENTIAL INSTANCE (row id + version), process-module scope.
+// Keying by the loaded credential's id+version — NOT by workspaceId alone (P1-b) —
+// is what makes a disconnect→reconnect safe: a post-reconnect getToken loads a
+// DIFFERENT row (new uuid id) and so gets a distinct key, instead of coalescing
+// onto the OLD in-flight refresh and writing stale rotated tokens over the freshly
+// connected credential. Concurrent calls for the SAME credential still coalesce,
+// so the one-time refresh token is never double-spent.
+const inflight = new Map<string, Promise<CodexTokenSnapshot>>();
+
+// Dependencies are injectable so the lifecycle logic (single-flight, staleness,
+// needs_relogin transition) is unit-testable without a database. Production uses
+// the real db + codex functions via the default bag.
+export type CodexAuthDeps = {
+  loadCredential: typeof loadCodexCredentialForRun;
+  recordRefresh: typeof recordCodexTokenRefresh;
+  setStatus: typeof setCodexCredentialStatus;
+  refresh: typeof refreshCodexToken;
+  encrypt: typeof encryptEnvironmentValue;
+  keyBytes: typeof environmentsEncryptionKeyBytes;
+};
+
+const defaultDeps: CodexAuthDeps = {
+  loadCredential: loadCodexCredentialForRun,
+  recordRefresh: recordCodexTokenRefresh,
+  setStatus: setCodexCredentialStatus,
+  refresh: refreshCodexToken,
+  encrypt: encryptEnvironmentValue,
+  keyBytes: environmentsEncryptionKeyBytes,
+};
+
+export function buildCodexTokenResolver(
+  db: Database,
+  settings: Settings,
+  workspaceId: string,
+  // The RESOLVED effective credential id (pin > workspace active), threaded from
+  // the worker. A mid-turn switch loads a DIFFERENT row id, gets a distinct
+  // single-flight key, and the (id, version) CAS in recordCodexTokenRefresh writes
+  // 0 rows against the now-inactive row — so a refresh racing a switch can never
+  // clobber the newly-active account. The single-flight map needs zero change.
+  credentialId: string,
+  deps: CodexAuthDeps = defaultDeps,
+): { getToken: () => Promise<CodexTokenSnapshot>; refresh: () => Promise<CodexTokenSnapshot> } {
+  const snapshot = (cred: CodexCredentialForRun): CodexTokenSnapshot => ({
+    accessToken: cred.tokens.accessToken,
+    chatgptAccountId: cred.chatgptAccountId,
+    isFedramp: cred.isFedramp,
+  });
+
+  const performRefresh = async (cred: CodexCredentialForRun): Promise<CodexTokenSnapshot> => {
+    try {
+      const next = await deps.refresh(cred.tokens.refreshToken);
+      const tokens = {
+        access_token: next.accessToken ?? cred.tokens.accessToken,
+        refresh_token: next.refreshToken ?? cred.tokens.refreshToken,
+        id_token: next.idToken ?? cred.tokens.idToken,
+      };
+      const key = deps.keyBytes(settings);
+      if (!key) {
+        throw new Error("OPENGENI_ENVIRONMENTS_ENCRYPTION_KEY is not configured");
+      }
+      // Compare-and-set on the loaded (id, version): if a disconnect→reconnect
+      // replaced the row mid-refresh, this writes 0 rows and we must NOT clobber
+      // the new credential with our now-defunct rotated tokens.
+      const persisted = await deps.recordRefresh(db, {
+        id: cred.id,
+        version: cred.version,
+        workspaceId,
+        credentialEncrypted: deps.encrypt(key, JSON.stringify(tokens)),
+        expiresAt: accessTokenExpiry(tokens.access_token),
+        lastRefreshAt: new Date(),
+      });
+      if (!persisted) {
+        // The row changed under us. Our rotated tokens belong to a stale family;
+        // fall back to whatever is connected NOW (a reconnect leaves an active
+        // row). If nothing active remains, a relogin is genuinely required.
+        const current = await deps.loadCredential(db, settings, workspaceId, credentialId);
+        if (current && current.status === "active") {
+          return snapshot(current);
+        }
+        throw new CodexReloginRequired("Codex credential changed during token refresh; reconnect required.");
+      }
+      return { accessToken: tokens.access_token, chatgptAccountId: cred.chatgptAccountId, isFedramp: cred.isFedramp };
+    } catch (error) {
+      if (error instanceof CodexReloginRequired) {
+        // Stamp needs_relogin ONLY if the row we refreshed is STILL current
+        // (compare-and-set on the loaded id+version). A relogin triggered by the
+        // OLD token family must never stamp needs_relogin onto a freshly
+        // reconnected credential.
+        await deps.setStatus(db, workspaceId, "needs_relogin", error.message, { id: cred.id, version: cred.version });
+      }
+      throw error;
+    }
+  };
+
+  // ALL refreshes — whether triggered by proactive staleness (getToken) or by a
+  // 401 retry (refresh) — coalesce onto one in-flight promise per credential
+  // instance, so concurrent calls can never double-spend the one-time refresh
+  // token (which would trigger refresh_token_reused -> needs_relogin).
+  const doRefresh = (cred: CodexCredentialForRun): Promise<CodexTokenSnapshot> => {
+    const key = `${cred.id}:${cred.version}`;
+    const existing = inflight.get(key);
+    if (existing) {
+      return existing;
+    }
+    const promise = performRefresh(cred).finally(() => {
+      if (inflight.get(key) === promise) {
+        inflight.delete(key);
+      }
+    });
+    inflight.set(key, promise);
+    return promise;
+  };
+
+  const resolve = async (force: boolean): Promise<CodexTokenSnapshot> => {
+    const cred = await deps.loadCredential(db, settings, workspaceId, credentialId);
+    if (!cred) {
+      throw new CodexReloginRequired("No Codex subscription is connected for this workspace.");
+    }
+    const exp = cred.expiresAt ?? accessTokenExpiry(cred.tokens.accessToken);
+    const stale =
+      force ||
+      (exp
+        ? exp.getTime() <= Date.now() + CODEX_REFRESH_WINDOW_MS
+        : cred.lastRefreshAt
+          ? cred.lastRefreshAt.getTime() < Date.now() - CODEX_REFRESH_FALLBACK_MS
+          : true);
+    return stale ? doRefresh(cred) : snapshot(cred);
+  };
+
+  return { getToken: () => resolve(false), refresh: () => resolve(true) };
+}
+
+function errorUsagePayload(reason?: "needs_relogin"): CodexUsagePayload {
+  return {
+    status: "error",
+    planType: null,
+    fiveHour: null,
+    weekly: null,
+    limitReached: false,
+    fetchedAt: new Date().toISOString(),
+    ...(reason ? { reason } : {}),
+  };
+}
+
+/**
+ * THE single per-account usage path both the api route and an (optional) worker
+ * poll call, so the refresh discipline and the cache-write can never drift.
+ *
+ *   1. resolve a REFRESHING bearer for THIS account (proactive staleness refresh,
+ *      single-flight, (id,version) CAS-persist) — this is what stops an idle
+ *      account's expired JWT from 401-ing the usage read.
+ *   2. fetch GET /wham/usage with that bearer.
+ *   3. normalize (§3) into the P2/P3 contract.
+ *   4. on any windows present, write the five usage-cache columns (the TTL clock).
+ *
+ * A refresh that stamps needs_relogin returns { status:"error", reason } and never
+ * hits the provider; a transient refresh error returns a plain error payload.
+ */
+export async function fetchCodexUsageForAccount(
+  db: Database,
+  settings: Settings,
+  workspaceId: string,
+  credentialId: string,
+): Promise<CodexUsagePayload> {
+  const resolver = buildCodexTokenResolver(db, settings, workspaceId, credentialId);
+  let token: CodexTokenSnapshot;
+  try {
+    token = await resolver.getToken();
+  } catch (error) {
+    return errorUsagePayload(error instanceof CodexReloginRequired ? "needs_relogin" : undefined);
+  }
+
+  let normalized: CodexUsagePayload;
+  try {
+    const usage = await fetchCodexUsage({
+      accessToken: token.accessToken,
+      chatgptAccountId: token.chatgptAccountId,
+      isFedramp: token.isFedramp,
+      clientVersion: CODEX_CLIENT_VERSION,
+    });
+    normalized = normalizeCodexUsage(usage.status, usage.payload);
+  } catch {
+    // A network throw on the /wham/usage read must surface as an error PAYLOAD
+    // ({status:"error"} at 200), never an unhandled 500 from the route.
+    return errorUsagePayload();
+  }
+
+  if (normalized.fiveHour || normalized.weekly) {
+    // Cache-write is best-effort: a disconnect under us (false) or a transient
+    // write error must NOT sink the freshly-read usage we are about to return.
+    await recordCodexAccountUsage(db, workspaceId, credentialId, {
+      primaryUsedPercent: normalized.fiveHour?.percent ?? null,
+      primaryResetAt: normalized.fiveHour?.resetAt ? new Date(normalized.fiveHour.resetAt) : null,
+      secondaryUsedPercent: normalized.weekly?.percent ?? null,
+      secondaryResetAt: normalized.weekly?.resetAt ? new Date(normalized.weekly.resetAt) : null,
+      checkedAt: new Date(),
+    }).catch(() => undefined);
+  }
+
+  return normalized;
+}

package/src/environment-crypto.ts

@@ -0,0 +1,51 @@
+import { createCipheriv, createDecipheriv, randomBytes } from "node:crypto";
+
+const VERSION_PREFIX = "v1";
+const IV_BYTES = 12;
+const GCM_TAG_BYTES = 16;
+const KEY_BYTES = 32;
+
+/**
+ * Encrypts one workspace environment variable value with AES-256-GCM under an
+ * operator key held outside Postgres. Output format: `v1:<b64 iv>:<b64 ciphertext||tag>`.
+ */
+export function encryptEnvironmentValue(key: Uint8Array, plaintext: string): string {
+  assertKey(key);
+  const iv = randomBytes(IV_BYTES);
+  const cipher = createCipheriv("aes-256-gcm", key, iv);
+  const ciphertext = Buffer.concat([cipher.update(plaintext, "utf8"), cipher.final(), cipher.getAuthTag()]);
+  return `${VERSION_PREFIX}:${iv.toString("base64")}:${ciphertext.toString("base64")}`;
+}
+
+/**
+ * Decrypts a stored `v1:` value. Error messages never echo plaintext or
+ * ciphertext: unknown versions throw "unsupported environment value format",
+ * auth-tag mismatches throw "environment value decryption failed".
+ */
+export function decryptEnvironmentValue(key: Uint8Array, stored: string): string {
+  assertKey(key);
+  const parts = stored.split(":");
+  if (parts.length !== 3 || parts[0] !== VERSION_PREFIX) {
+    throw new Error("unsupported environment value format");
+  }
+  const iv = Buffer.from(parts[1]!, "base64");
+  const payload = Buffer.from(parts[2]!, "base64");
+  if (iv.length !== IV_BYTES || payload.length <= GCM_TAG_BYTES) {
+    throw new Error("unsupported environment value format");
+  }
+  const tag = payload.subarray(payload.length - GCM_TAG_BYTES);
+  const ciphertext = payload.subarray(0, payload.length - GCM_TAG_BYTES);
+  const decipher = createDecipheriv("aes-256-gcm", key, iv);
+  decipher.setAuthTag(tag);
+  try {
+    return Buffer.concat([decipher.update(ciphertext), decipher.final()]).toString("utf8");
+  } catch {
+    throw new Error("environment value decryption failed");
+  }
+}
+
+function assertKey(key: Uint8Array): void {
+  if (key.length !== KEY_BYTES) {
+    throw new Error("environment encryption key must be exactly 32 bytes");
+  }
+}

package/src/event-payload-sanitizer.ts

@@ -0,0 +1,89 @@
+/**
+ * Last line of defense against a session event crashing a whole turn.
+ *
+ * Postgres `text`/`jsonb` cannot store a NUL byte (U+0000) nor lone UTF-16
+ * surrogates. Raw exec output routinely carries both -- chrome/crashpad logs,
+ * `cat` of a binary, random bytes -- and the worker persists that output verbatim
+ * inside `agent.toolCall.output` / `sandbox.command.output` event payloads. When
+ * such a payload reaches `INSERT INTO session_events`, the driver rejects it
+ * ("Failed query: insert into session_events") and the turn dies.
+ *
+ * `sanitizeEventPayload` deep-walks any payload value (objects, arrays, nested)
+ * and, for every string, strips NUL and rewrites invalid/lone UTF-16 surrogates
+ * to the Unicode replacement char (U+FFFD), so the result is always valid UTF-8
+ * that jsonb can store. It is cheap and total: only strings are touched, and only
+ * the two disallowed classes of code unit -- no meaningful text is lost, no
+ * truncation (truncation is handled elsewhere).
+ */
+
+const REPLACEMENT = "�";
+
+/**
+ * Strip NUL and repair invalid/lone UTF-16 surrogates in a single string.
+ * Returns the input unchanged (same reference) when it is already clean, so the
+ * common case allocates nothing.
+ */
+export function sanitizeEventString(value: string): string {
+  // Fast path: no NUL and no surrogate code unit at all -> nothing to do.
+  // Surrogates live in U+D800..U+DFFF; a quick scan avoids the rebuild cost.
+  let needsWork = false;
+  for (let i = 0; i < value.length; i++) {
+    const code = value.charCodeAt(i);
+    if (code === 0x0000 || (code >= 0xd800 && code <= 0xdfff)) {
+      needsWork = true;
+      break;
+    }
+  }
+  if (!needsWork) {
+    return value;
+  }
+
+  let out = "";
+  for (let i = 0; i < value.length; i++) {
+    const code = value.charCodeAt(i);
+    if (code === 0x0000) {
+      // Drop NUL entirely.
+      continue;
+    }
+    if (code >= 0xd800 && code <= 0xdbff) {
+      // High surrogate: valid only when immediately followed by a low surrogate.
+      const next = i + 1 < value.length ? value.charCodeAt(i + 1) : 0;
+      if (next >= 0xdc00 && next <= 0xdfff) {
+        out += value[i]! + value[i + 1]!;
+        i += 1;
+        continue;
+      }
+      out += REPLACEMENT;
+      continue;
+    }
+    if (code >= 0xdc00 && code <= 0xdfff) {
+      // Lone low surrogate (a valid pair would have been consumed above).
+      out += REPLACEMENT;
+      continue;
+    }
+    out += value[i]!;
+  }
+  return out;
+}
+
+/**
+ * Deep-walk a session event payload and sanitize every string value. Mirrors the
+ * shape of the worker redaction deep-walk: objects, arrays, and nested
+ * combinations are traversed; non-string leaves pass through untouched. Object
+ * keys are sanitized too -- they are jsonb-constrained the same as values.
+ */
+export function sanitizeEventPayload<T>(payload: T): T {
+  if (typeof payload === "string") {
+    return sanitizeEventString(payload) as unknown as T;
+  }
+  if (Array.isArray(payload)) {
+    return payload.map((item) => sanitizeEventPayload(item)) as unknown as T;
+  }
+  if (payload && typeof payload === "object") {
+    const entries = Object.entries(payload as Record<string, unknown>).map(
+      ([key, value]) => [sanitizeEventString(key), sanitizeEventPayload(value)] as const,
+    );
+    return Object.fromEntries(entries) as unknown as T;
+  }
+  return payload;
+}