@remodex/rmx 1.0.2

package/src/codex/user-identity.ts

@@ -0,0 +1,374 @@
+/**
+ * Environment-independent identity and namespace resolution for Codex writes.
+ *
+ * Bun 1.3.14 made the obvious implementation unsafe: both `os.homedir()` and
+ * `os.userInfo().homedir` follow HOME. A service and CLI for the same account
+ * could therefore coordinate through different databases. The namespace is
+ * keyed only by the effective uid/SID and the canonical CODEX_HOME.
+ *
+ * Design record: devlog/_fin/260804_codex_write_substrate/005_contract.md §7.
+ */
+import { createHash } from "node:crypto";
+import {
+  lstatSync,
+  mkdirSync,
+  realpathSync,
+  statSync,
+} from "node:fs";
+import { isAbsolute, join, resolve } from "node:path";
+
+import type {
+  ResolveCodexCoordinatorDatabasePath,
+  ResolveCodexCatalogSerializationDatabasePath,
+  ResolveCodexHistorySerializationDatabasePath,
+  ResolveEffectiveUserIdentity,
+  UserIdentity,
+} from "./convergence-types";
+
+const POSIX_PRIVATE_MODE = 0o700;
+const POSIX_TMP_REQUIRED_MODE = 0o1003;
+const POSIX_TMP_PATH = "/tmp";
+const SID_PATTERN = /^S-1-(?:\d+-)+\d+$/i;
+
+export class CodexUserIdentityRefusal extends Error {
+  readonly code = "CODEX_USER_IDENTITY_REFUSED";
+
+  constructor(message: string, options?: ErrorOptions) {
+    super(message, options);
+    this.name = "CodexUserIdentityRefusal";
+  }
+}
+
+function refuse(message: string, cause?: unknown): never {
+  throw new CodexUserIdentityRefusal(message, cause === undefined ? undefined : { cause });
+}
+
+function powershellValue(expression: string): string {
+  let result: ReturnType<typeof Bun.spawnSync>;
+  try {
+    result = Bun.spawnSync([
+      "powershell.exe",
+      "-NoLogo",
+      "-NoProfile",
+      "-NonInteractive",
+      "-Command",
+      expression,
+    ], {
+      stdin: "ignore",
+      stdout: "pipe",
+      stderr: "pipe",
+    });
+  } catch (cause) {
+    refuse("Windows effective-account lookup could not start.", cause);
+  }
+  if (result.exitCode !== 0) refuse("Windows effective-account lookup failed.");
+  const value = new TextDecoder().decode(result.stdout).trim();
+  if (!value) refuse("Windows effective-account lookup returned an empty value.");
+  return value;
+}
+
+function resolveWindowsSid(): string {
+  const sid = powershellValue(
+    "[System.Security.Principal.WindowsIdentity]::GetCurrent().User.Value",
+  );
+  if (!SID_PATTERN.test(sid)) refuse("Windows effective-account lookup returned an invalid SID.");
+  return sid.toUpperCase();
+}
+
+export const resolveEffectiveUserIdentity: ResolveEffectiveUserIdentity = () => {
+  if (process.platform === "win32") {
+    return { platform: "win32", sid: resolveWindowsSid() };
+  }
+
+  const getuid = process.getuid;
+  if (typeof getuid !== "function") {
+    refuse("The runtime does not expose the effective POSIX uid.");
+  }
+  let uid: number;
+  try {
+    uid = getuid.call(process);
+  } catch (cause) {
+    refuse("The effective POSIX uid lookup failed.", cause);
+  }
+  if (!Number.isSafeInteger(uid) || uid < 0) {
+    refuse("The runtime returned an invalid effective POSIX uid.");
+  }
+  return { platform: "posix", uid };
+};
+
+function assertPrivatePosixDirectory(path: string, uid: number): void {
+  let entry;
+  try {
+    entry = lstatSync(path);
+  } catch (cause) {
+    refuse("The Codex coordinator namespace cannot be inspected.", cause);
+  }
+  if (entry.isSymbolicLink() || !entry.isDirectory()) {
+    refuse("The Codex coordinator namespace is not a real directory.");
+  }
+  if (entry.uid !== uid || (entry.mode & 0o777) !== POSIX_PRIVATE_MODE) {
+    refuse("The Codex coordinator namespace has unsafe ownership or permissions.");
+  }
+}
+
+function ensurePrivatePosixDirectory(path: string, uid: number): void {
+  try {
+    mkdirSync(path, { mode: POSIX_PRIVATE_MODE });
+  } catch (cause) {
+    const code = cause && typeof cause === "object" && "code" in cause
+      ? String((cause as { code?: unknown }).code)
+      : "";
+    if (code !== "EEXIST") refuse("The Codex coordinator namespace cannot be created.", cause);
+  }
+  assertPrivatePosixDirectory(path, uid);
+}
+
+function resolveTrustedPosixTmp(): string {
+  let realTmp: string;
+  try {
+    realTmp = realpathSync.native(POSIX_TMP_PATH);
+    const entry = statSync(realTmp);
+    if (!entry.isDirectory() || entry.uid !== 0) {
+      refuse("The system temporary directory has unsafe ownership.");
+    }
+    if ((entry.mode & POSIX_TMP_REQUIRED_MODE) !== POSIX_TMP_REQUIRED_MODE) {
+      refuse("The system temporary directory lacks sticky world write/search permissions.");
+    }
+  } catch (cause) {
+    if (cause instanceof CodexUserIdentityRefusal) throw cause;
+    refuse("The system temporary directory cannot be trusted.", cause);
+  }
+  return realTmp;
+}
+
+function resolvePosixRuntimeRoot(uid: number): string {
+  const realTmp = resolveTrustedPosixTmp();
+  const root = join(realTmp, `opencodex-runtime-v1-${uid}`);
+  ensurePrivatePosixDirectory(root, uid);
+  return root;
+}
+
+export type CoordinatorNamespaceProbe =
+  | { readonly status: "ok"; readonly root: string }
+  | { readonly status: "missing" };
+
+/**
+ * Read-only namespace probe for diagnostics (`rmx doctor`).
+ *
+ * Unlike the runtime resolvers, this never creates the root or the lock
+ * directories: a doctor run must observe the namespace, not initialize it.
+ * A missing namespace is reported as `missing` instead of refused, so a fresh
+ * machine does not read as a broken one; an existing but unsafe namespace is
+ * refused exactly like the creating path would refuse it.
+ */
+export function probeCodexCoordinatorNamespace(identity: UserIdentity): CoordinatorNamespaceProbe {
+  if (identity.platform === "posix") {
+    const root = join(resolveTrustedPosixTmp(), `opencodex-runtime-v1-${identity.uid}`);
+    let entry;
+    try {
+      entry = lstatSync(root);
+    } catch (cause) {
+      const code = cause && typeof cause === "object" && "code" in cause
+        ? String((cause as { code?: unknown }).code)
+        : "";
+      if (code === "ENOENT") return { status: "missing" };
+      refuse("The Codex coordinator namespace cannot be inspected.", cause);
+    }
+    if (entry.isSymbolicLink() || !entry.isDirectory()) {
+      refuse("The Codex coordinator namespace is not a real directory.");
+    }
+    if (entry.uid !== identity.uid || (entry.mode & 0o777) !== POSIX_PRIVATE_MODE) {
+      refuse("The Codex coordinator namespace has unsafe ownership or permissions.");
+    }
+    return { status: "ok", root };
+  }
+
+  if (!SID_PATTERN.test(identity.sid)) refuse("The coordinator identity contains an invalid SID.");
+  const localAppData = powershellValue(
+    "[Environment]::GetFolderPath([Environment+SpecialFolder]::LocalApplicationData)",
+  );
+  if (!isAbsolute(localAppData)) refuse("Windows LocalAppData resolution returned a relative path.");
+  const root = resolve(localAppData, "OpenCodex", "Runtime", "v1", identity.sid.toUpperCase());
+  let entry;
+  try {
+    entry = lstatSync(root);
+  } catch (cause) {
+    const code = cause && typeof cause === "object" && "code" in cause
+      ? String((cause as { code?: unknown }).code)
+      : "";
+    if (code === "ENOENT") return { status: "missing" };
+    refuse("The Windows coordinator namespace cannot be inspected.", cause);
+  }
+  if (entry.isSymbolicLink() || !entry.isDirectory()) {
+    refuse("The Windows coordinator namespace is not a real directory.");
+  }
+  try {
+    const real = realpathSync.native(root);
+    if (!samePathIdentity(real, root, "win32")) {
+      refuse("The Windows coordinator namespace is redirected by a junction or reparse point.");
+    }
+    return { status: "ok", root: real };
+  } catch (cause) {
+    if (cause instanceof CodexUserIdentityRefusal) throw cause;
+    refuse("The Windows coordinator namespace cannot be resolved.", cause);
+  }
+}
+
+function resolveWindowsRuntimeRoot(identity: Extract<UserIdentity, { platform: "win32" }>): string {
+  if (!SID_PATTERN.test(identity.sid)) refuse("The coordinator identity contains an invalid SID.");
+  const localAppData = powershellValue(
+    "[Environment]::GetFolderPath([Environment+SpecialFolder]::LocalApplicationData)",
+  );
+  if (!isAbsolute(localAppData)) refuse("Windows LocalAppData resolution returned a relative path.");
+
+  // The SID and known-folder values come from the effective token/.NET OS APIs,
+  // never USERPROFILE or LOCALAPPDATA. WP11 adds descriptor/reparse/ACL checks at
+  // the stable-database open boundary where those checks can cover SQLite too.
+  const root = resolve(localAppData, "OpenCodex", "Runtime", "v1", identity.sid.toUpperCase());
+  try {
+    mkdirSync(root, { recursive: true });
+  } catch (cause) {
+    refuse("The Windows coordinator namespace cannot be created.", cause);
+  }
+  // Canonicalize before anything is keyed on the path: a junctioned or
+  // differently-cased LocalAppData must land on ONE namespace, or two processes
+  // that share the real directory would build different lock paths and never
+  // contend. The lock modules also compare this path against realpath, so a
+  // non-canonical spelling here would read as "unsafe" on every acquisition.
+  //
+  // Canonicalizing must only ever fold spelling (case, separators). If the
+  // realpath lands somewhere else entirely, a component of the namespace is a
+  // junction/reparse redirect; accepting the target would convert the old
+  // refusal into silently opening the redirected location, so refuse instead.
+  try {
+    const real = realpathSync.native(root);
+    if (!samePathIdentity(real, root, "win32")) {
+      refuse("The Windows coordinator namespace is redirected by a junction or reparse point.");
+    }
+    return real;
+  } catch (cause) {
+    if (cause instanceof CodexUserIdentityRefusal) throw cause;
+    refuse("The Windows coordinator namespace cannot be resolved.", cause);
+  }
+}
+
+/**
+ * Windows path identity is case-insensitive; everywhere else it is exact.
+ *
+ * The lock modules compare a requested lock path against its own realpath, and
+ * byte equality refuses legitimate Windows spellings (drive-letter case, mixed
+ * component casing) as "unsafe". This is the same semantics
+ * `history-provider.ts` already applies to manifest paths; the platform
+ * argument exists so both branches are testable on any host.
+ */
+export function samePathIdentity(
+  a: string,
+  b: string,
+  platform: NodeJS.Platform = process.platform,
+): boolean {
+  const left = resolve(a);
+  const right = resolve(b);
+  return platform === "win32" ? left.toLowerCase() === right.toLowerCase() : left === right;
+}
+
+export const resolveCodexCoordinatorDatabasePath: ResolveCodexCoordinatorDatabasePath = (
+  identity,
+  canonicalCodexHome,
+) => {
+  if (!isAbsolute(canonicalCodexHome)) {
+    refuse("The canonical CODEX_HOME must be an absolute path.");
+  }
+  const root = identity.platform === "posix"
+    ? resolvePosixRuntimeRoot(identity.uid)
+    : resolveWindowsRuntimeRoot(identity);
+  const locks = join(root, "native-write-locks");
+  if (identity.platform === "posix") ensurePrivatePosixDirectory(locks, identity.uid);
+  else {
+    try {
+      mkdirSync(locks, { recursive: true });
+    } catch (cause) {
+      refuse("The Windows coordinator lock directory cannot be created.", cause);
+    }
+  }
+
+  const homeDigest = createHash("sha256").update(canonicalCodexHome).digest("hex");
+  return join(locks, `${homeDigest}.sqlite`);
+};
+
+/**
+ * K's FINAL database path. Never the native coordinator path.
+ *
+ * Catalog serialization is a different ownership surface from the native
+ * coordinator N: `K -> C` is a legal order and `N -> K` nests, so sharing one
+ * database would make the required nesting self-contend. The two live in
+ * sibling directories under the same per-user runtime root — same identity
+ * namespace, same environment-independent parent, distinct exclusion.
+ *
+ * Consumers use the returned path verbatim and append nothing
+ * (`005_contract.md:1256-1330`).
+ */
+export const resolveCodexCatalogSerializationDatabasePath:
+  ResolveCodexCatalogSerializationDatabasePath = (identity, canonicalCodexHome) => {
+    if (!isAbsolute(canonicalCodexHome)) {
+      refuse("The canonical CODEX_HOME must be an absolute path.");
+    }
+    const root = identity.platform === "posix"
+      ? resolvePosixRuntimeRoot(identity.uid)
+      : resolveWindowsRuntimeRoot(identity);
+    const locks = join(root, "catalog-write-locks");
+    if (identity.platform === "posix") ensurePrivatePosixDirectory(locks, identity.uid);
+    else {
+      try {
+        mkdirSync(locks, { recursive: true });
+      } catch (cause) {
+        refuse("The Windows catalog serialization directory cannot be created.", cause);
+      }
+    }
+
+    const homeDigest = createHash("sha256").update(canonicalCodexHome).digest("hex");
+    return join(locks, `${homeDigest}.sqlite`);
+  };
+
+/**
+ * H's FINAL database path.
+ *
+ * Keyed by the canonical state DB in addition to the canonical home, unlike N and
+ * K. One `CODEX_HOME` can name a different `state_5.sqlite` — `model_catalog_json`
+ * has the same shape of indirection for catalogs — and two operations against
+ * different history databases are not the same exclusion. Hashing only the home
+ * would serialize them together, and hashing the raw request path would let two
+ * spellings of one database take different locks.
+ */
+export const resolveCodexHistorySerializationDatabasePath:
+  ResolveCodexHistorySerializationDatabasePath = (
+    identity,
+    canonicalCodexHome,
+    canonicalStateDbPath,
+  ) => {
+    if (!isAbsolute(canonicalCodexHome)) {
+      refuse("The canonical CODEX_HOME must be an absolute path.");
+    }
+    if (!isAbsolute(canonicalStateDbPath)) {
+      refuse("The canonical Codex state database must be an absolute path.");
+    }
+    const root = identity.platform === "posix"
+      ? resolvePosixRuntimeRoot(identity.uid)
+      : resolveWindowsRuntimeRoot(identity);
+    const locks = join(root, "history-write-locks");
+    if (identity.platform === "posix") ensurePrivatePosixDirectory(locks, identity.uid);
+    else {
+      try {
+        mkdirSync(locks, { recursive: true });
+      } catch (cause) {
+        refuse("The Windows history serialization directory cannot be created.", cause);
+      }
+    }
+
+    // Both components are length-prefixed so no pair of (home, stateDb) values can
+    // collide by concatenation.
+    const digest = createHash("sha256")
+      .update(`${canonicalCodexHome.length}:${canonicalCodexHome}`)
+      .update(`${canonicalStateDbPath.length}:${canonicalStateDbPath}`)
+      .digest("hex");
+    return join(locks, `${digest}.sqlite`);
+  };

package/src/codex/warmup.ts

@@ -0,0 +1,192 @@
+export class CodexWarmupError extends Error {
+  code: "http_status" | "missing_body" | "stream_failed" | "stream_incomplete" | "stream_error" | "invalid_sse" | "no_terminal" | "transport";
+  status?: number;
+  /** Upstream error detail extracted from the response body (truncated to 512 chars). */
+  upstreamDetail?: string;
+
+  constructor(
+    code: CodexWarmupError["code"],
+    message = "Codex warmup failed",
+    options: { status?: number; cause?: unknown; upstreamDetail?: string } = {},
+  ) {
+    super(message);
+    this.name = "CodexWarmupError";
+    this.code = code;
+    this.status = options.status;
+    this.upstreamDetail = options.upstreamDetail;
+    if (options.cause !== undefined) this.cause = options.cause;
+  }
+}
+
+export interface CodexWarmupOptions {
+  accessToken: string;
+  chatgptAccountId: string;
+  model?: string;
+  timeoutMs?: number;
+}
+
+const CODEX_RESPONSES_URL = "https://chatgpt.com/backend-api/codex/responses";
+const DEFAULT_MODEL = "gpt-5.4-mini";
+const FALLBACK_MODELS = ["gpt-5.5"];
+const DEFAULT_TIMEOUT_MS = 30_000;
+const MAX_ERROR_BODY_BYTES = 2048;
+
+/** Read the first MAX_ERROR_BODY_BYTES of a response body and extract an error message. */
+async function readErrorDetail(res: Response): Promise<string | undefined> {
+  try {
+    const text = await res.text();
+    const trimmed = text.slice(0, MAX_ERROR_BODY_BYTES);
+    try {
+      const json = JSON.parse(trimmed) as Record<string, unknown>;
+      // ChatGPT backend error shape: { error: { message: "..." } } or { detail: "..." }
+      const nested = json.error;
+      if (nested && typeof nested === "object" && typeof (nested as Record<string, unknown>).message === "string") {
+        return ((nested as Record<string, unknown>).message as string).slice(0, 512);
+      }
+      if (typeof json.detail === "string") return json.detail.slice(0, 512);
+      if (typeof json.error === "string") return (json.error as string).slice(0, 512);
+      if (typeof json.message === "string") return json.message.slice(0, 512);
+    } catch {
+      // Non-JSON response body may contain sensitive data (tokens, credentials).
+      // Only surface structured error messages, never raw text.
+    }
+    return undefined;
+  } catch {
+    return undefined;
+  }
+}
+
+function safeWarmupReason(err: unknown): string {
+  if (err instanceof CodexWarmupError) {
+    const base = err.status ? `${err.code}:${err.status}` : err.code;
+    return err.upstreamDetail ? `${base} — ${err.upstreamDetail}` : base;
+  }
+  return "transport";
+}
+
+export function codexWarmupFailureReason(err: unknown): string {
+  return safeWarmupReason(err);
+}
+
+function eventTypeFromData(data: unknown): string | undefined {
+  if (!data || typeof data !== "object") return undefined;
+  const record = data as Record<string, unknown>;
+  return typeof record.type === "string" ? record.type : undefined;
+}
+
+function parseSseFrame(frame: string): unknown | null {
+  const dataLines = frame
+    .split(/\r?\n/)
+    .filter(line => line.startsWith("data:"))
+    .map(line => line.slice(5).trimStart());
+  if (dataLines.length === 0) return null;
+  const data = dataLines.join("\n").trim();
+  if (!data || data === "[DONE]") return null;
+  try {
+    return JSON.parse(data) as unknown;
+  } catch (err) {
+    throw new CodexWarmupError("invalid_sse", "Codex warmup received invalid SSE", { cause: err });
+  }
+}
+
+async function drainWarmupSse(body: ReadableStream<Uint8Array>): Promise<void> {
+  const reader = body.getReader();
+  const decoder = new TextDecoder();
+  let buffer = "";
+
+  try {
+    for (;;) {
+      const { done, value } = await reader.read();
+      if (done) break;
+      buffer += decoder.decode(value, { stream: true });
+
+      for (;;) {
+        const frameEnd = buffer.search(/\r?\n\r?\n/);
+        if (frameEnd < 0) break;
+        const frame = buffer.slice(0, frameEnd);
+        const delimiterLength = buffer[frameEnd] === "\r" ? 4 : 2;
+        buffer = buffer.slice(frameEnd + delimiterLength);
+        const parsed = parseSseFrame(frame);
+        const type = eventTypeFromData(parsed);
+        if (type === "response.completed") return;
+        if (type === "response.failed") throw new CodexWarmupError("stream_failed");
+        if (type === "response.incomplete") throw new CodexWarmupError("stream_incomplete");
+        if (type === "error") throw new CodexWarmupError("stream_error");
+      }
+    }
+
+    if (buffer.trim()) {
+      const parsed = parseSseFrame(buffer);
+      const type = eventTypeFromData(parsed);
+      if (type === "response.completed") return;
+      if (type === "response.failed") throw new CodexWarmupError("stream_failed");
+      if (type === "response.incomplete") throw new CodexWarmupError("stream_incomplete");
+      if (type === "error") throw new CodexWarmupError("stream_error");
+    }
+
+    throw new CodexWarmupError("no_terminal", "Codex warmup ended before completion");
+  } finally {
+    reader.releaseLock();
+  }
+}
+
+async function tryWarmup(options: CodexWarmupOptions, model: string): Promise<void> {
+  let res: Response;
+  try {
+    res = await fetch(CODEX_RESPONSES_URL, {
+      method: "POST",
+      headers: {
+        Authorization: `Bearer ${options.accessToken}`,
+        "ChatGPT-Account-Id": options.chatgptAccountId,
+        "Content-Type": "application/json",
+      },
+      body: JSON.stringify({
+        model,
+        instructions: "Reply with OK.",
+        input: [{ type: "message", role: "user", content: [{ type: "input_text", text: "hi" }] }],
+        stream: true,
+        store: false,
+      }),
+      signal: AbortSignal.timeout(options.timeoutMs ?? DEFAULT_TIMEOUT_MS),
+    });
+  } catch (err) {
+    throw new CodexWarmupError("transport", "Codex warmup request failed", { cause: err });
+  }
+
+  if (!res.ok) {
+    const upstreamDetail = await readErrorDetail(res);
+    throw new CodexWarmupError("http_status", "Codex warmup was rejected", {
+      status: res.status,
+      upstreamDetail,
+    });
+  }
+  if (!res.body) throw new CodexWarmupError("missing_body");
+
+  try {
+    await drainWarmupSse(res.body);
+  } finally {
+    await res.body?.cancel().catch(() => {});
+  }
+}
+
+export async function warmCodexAccount(options: CodexWarmupOptions): Promise<void> {
+  const primaryModel = options.model?.trim() || DEFAULT_MODEL;
+  try {
+    await tryWarmup(options, primaryModel);
+    return;
+  } catch (err) {
+    // Retry with fallback models on 400 (model may not be available for this account).
+    if (!(err instanceof CodexWarmupError) || err.status !== 400) throw err;
+    let lastErr = err;
+    for (const fallback of FALLBACK_MODELS) {
+      if (fallback === primaryModel) continue;
+      try {
+        await tryWarmup(options, fallback);
+        return;
+      } catch (retryErr) {
+        if (retryErr instanceof CodexWarmupError) lastErr = retryErr;
+      }
+    }
+    throw lastErr;
+  }
+}

package/src/codex/websocket-registry.ts

@@ -0,0 +1,100 @@
+import type { ServerWebSocket } from "bun";
+import type { WsData } from "../server/ws-bridge";
+import { createAdmissionGate, type AdmissionMetrics, type AdmissionReservation } from "../lib/admission";
+
+export const MAX_TRACKED_CODEX_WEBSOCKETS = 128;
+const websocketGate = createAdmissionGate("codex_websockets", MAX_TRACKED_CODEX_WEBSOCKETS);
+
+export function tryReserveCodexWebSocket(): AdmissionReservation<ServerWebSocket<WsData>> | null {
+  const gateLease = websocketGate.tryAcquire();
+  if (!gateLease) return null;
+  let socket: ServerWebSocket<WsData> | undefined;
+  let active = true;
+  return {
+    bind(value) {
+      if (!active) return;
+      socket = value;
+    },
+    release() {
+      if (!active) return;
+      active = false;
+      socket = undefined;
+      gateLease.release();
+    },
+  };
+}
+
+export function codexWebSocketAdmissionMetrics(): AdmissionMetrics {
+  return websocketGate.metrics();
+}
+
+const socketsByAccount = new Map<string, Set<ServerWebSocket<WsData>>>();
+
+function trackedAccountId(ws: ServerWebSocket<WsData>): string | null {
+  const ctx = ws.data.authContext;
+  return ctx?.kind === "pool" || ctx?.kind === "main-pool" ? ctx.accountId : null;
+}
+
+function addSocketForAccount(accountId: string | null, ws: ServerWebSocket<WsData>): void {
+  if (!accountId) return;
+  let sockets = socketsByAccount.get(accountId);
+  if (!sockets) {
+    sockets = new Set();
+    socketsByAccount.set(accountId, sockets);
+  }
+  sockets.add(ws);
+}
+
+function removeSocketForAccount(accountId: string | null, ws: ServerWebSocket<WsData>): void {
+  if (!accountId) return;
+  const sockets = socketsByAccount.get(accountId);
+  if (!sockets) return;
+  sockets.delete(ws);
+  if (sockets.size === 0) socketsByAccount.delete(accountId);
+}
+
+export function registerCodexWebSocket(ws: ServerWebSocket<WsData>): void {
+  addSocketForAccount(trackedAccountId(ws), ws);
+}
+
+export function unregisterCodexWebSocket(ws: ServerWebSocket<WsData>): void {
+  removeSocketForAccount(trackedAccountId(ws), ws);
+}
+
+export function updateCodexWebSocketAuthContext(
+  ws: ServerWebSocket<WsData>,
+  authContext: WsData["authContext"],
+): void {
+  const before = trackedAccountId(ws);
+  removeSocketForAccount(before, ws);
+  ws.data.authContext = authContext;
+  addSocketForAccount(trackedAccountId(ws), ws);
+}
+
+export function invalidateCodexWebSocketsForAccount(accountId: string): number {
+  const sockets = socketsByAccount.get(accountId);
+  if (!sockets) return 0;
+  const snapshot = [...sockets];
+  socketsByAccount.delete(accountId);
+  for (const ws of snapshot) {
+    try {
+      ws.data.cancel?.();
+    } catch {
+      /* ignore cancel callbacks during invalidation */
+    }
+    try {
+      ws.close(4001, "Codex account invalidated");
+    } catch {
+      /* socket may already be closing */
+    }
+  }
+  return snapshot.length;
+}
+
+export function getTrackedCodexWebSocketCountForAccount(accountId: string): number {
+  return socketsByAccount.get(accountId)?.size ?? 0;
+}
+
+export function clearCodexWebSocketRegistry(): void {
+  socketsByAccount.clear();
+}