@opengeni/db 0.9.3 → 0.10.7

package/drizzle/0108_fence_invalidated_warming_epochs.sql

@@ -0,0 +1,76 @@
+-- deployment-mode: rolling
+-- Fence provider creates that outlive a failed or expired warming acquisition.
+--
+-- A warming row can be invalidated while a provider create is still
+-- non-abortably in flight. Advancing lease_epoch in the same transaction that
+-- exposes cold/draining means a successor can never reuse the old acquisition
+-- epoch, so late record/commit/fail/cleanup callbacks fail their existing CAS.
+
+CREATE OR REPLACE FUNCTION opengeni_private.reap_sandbox_leases(
+  p_viewer_holder_ttl_ms bigint,
+  p_turn_holder_ttl_ms   bigint,
+  p_idle_grace_ms        bigint
+)
+RETURNS TABLE (workspace_id uuid, sandbox_group_id uuid, instance_id text, lease_epoch integer)
+LANGUAGE plpgsql
+SECURITY DEFINER
+AS $$
+BEGIN
+  DELETE FROM sandbox_lease_holders h
+  WHERE h.kind = 'viewer'
+    AND h.last_heartbeat_at < now() - make_interval(secs => p_viewer_holder_ttl_ms / 1000.0);
+
+  IF p_turn_holder_ttl_ms > 0 THEN
+    DELETE FROM sandbox_lease_holders h
+    WHERE h.kind = 'turn'
+      AND h.last_heartbeat_at < now() - make_interval(secs => p_turn_holder_ttl_ms / 1000.0);
+  END IF;
+
+  UPDATE sandbox_leases L SET
+    refcount       = c.total,
+    turn_holders   = c.turns,
+    viewer_holders = c.viewers,
+    liveness = CASE WHEN L.liveness = 'warm' AND c.total = 0 AND c.turns = 0
+                    THEN 'draining' ELSE L.liveness END,
+    expires_at = CASE WHEN L.liveness = 'warm' AND c.total = 0 AND c.turns = 0
+                    THEN now() + make_interval(secs => p_idle_grace_ms / 1000.0)
+                    ELSE L.expires_at END,
+    updated_at = now()
+  FROM (
+    SELECT L2.id,
+           (SELECT count(*) FROM sandbox_lease_holders h WHERE h.lease_id = L2.id)::int                       AS total,
+           (SELECT count(*) FROM sandbox_lease_holders h WHERE h.lease_id = L2.id AND h.kind = 'turn')::int   AS turns,
+           (SELECT count(*) FROM sandbox_lease_holders h WHERE h.lease_id = L2.id AND h.kind = 'viewer')::int AS viewers
+    FROM sandbox_leases L2
+  ) c
+  WHERE L.id = c.id;
+
+  -- The old warming epoch is permanently closed before a successor can acquire.
+  UPDATE sandbox_leases AS L SET
+    liveness = 'cold', instance_id = NULL,
+    resume_backend_id = NULL, resume_state = NULL,
+    data_plane_url = NULL, terminal_data_plane_url = NULL,
+    lease_epoch = L.lease_epoch + 1,
+    updated_at = now()
+  WHERE L.liveness = 'warming' AND L.expires_at < now() AND L.instance_id IS NULL;
+
+  -- Keep an attributed provider id for the drain/terminate path, but fence the
+  -- expired creator before exposing the row as drainable.
+  UPDATE sandbox_leases AS L SET
+    liveness = 'draining',
+    refcount = 0,
+    turn_holders = 0,
+    viewer_holders = 0,
+    data_plane_url = NULL,
+    terminal_data_plane_url = NULL,
+    lease_epoch = L.lease_epoch + 1,
+    expires_at = now() - interval '1 millisecond',
+    updated_at = now()
+  WHERE L.liveness = 'warming' AND L.expires_at < now() AND L.instance_id IS NOT NULL;
+
+  RETURN QUERY
+    SELECT L.workspace_id, L.sandbox_group_id, L.instance_id, L.lease_epoch
+    FROM sandbox_leases L
+    WHERE L.liveness = 'draining' AND L.expires_at < now() AND L.refcount = 0;
+END;
+$$;

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@opengeni/db",
-  "version": "0.9.3",
+  "version": "0.10.7",
   "description": "OpenGeni persistence: Drizzle schema, RLS-scoped query layer, the SQL migration runner, and role provisioning.",
   "license": "Apache-2.0",
   "repository": {
@@ -49,9 +49,10 @@
     "prepublishOnly": "bash ../../scripts/prepublish-guard"
   },
   "dependencies": {
-    "@opengeni/codex": "^0.2.5",
-    "@opengeni/config": "^0.6.2",
-    "@opengeni/contracts": "^0.15.0",
+    "@opengeni/codex": "^0.2.7",
+    "@opengeni/config": "^0.6.9",
+    "@opengeni/contracts": "^0.18.0",
+    "@opengeni/network": "^0.1.1",
     "drizzle-orm": "^0.45.2",
     "postgres": "^3.4.7"
   },

package/src/codex-token-resolver.ts

@@ -26,6 +26,10 @@ import {
   CodexReloginRequired,
   type CodexTokenSnapshot,
   type CodexUsagePayload,
+  type CodexFetch,
+  type CodexRateLimitResetCreditsDetails,
+  type ResetCreditFetchFailureReason,
+  fetchCodexRateLimitResetCredits,
   fetchCodexUsage,
   normalizeCodexUsage,
   refreshCodexToken,
@@ -49,6 +53,92 @@ import {
 // 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>>();
+const CODEX_TOKEN_REFRESH_TIMEOUT_MS = 6_000;
+
+export type CodexTokenDeadlineClock = {
+  setTimeout: (callback: () => void, delayMs: number) => ReturnType<typeof globalThis.setTimeout>;
+  clearTimeout: (handle: ReturnType<typeof globalThis.setTimeout>) => void;
+};
+
+export type CodexTokenDeadlineOptions = {
+  timeoutMs?: number | undefined;
+  signal?: AbortSignal | undefined;
+  clock?: CodexTokenDeadlineClock | undefined;
+};
+
+const systemCodexTokenDeadlineClock: CodexTokenDeadlineClock = {
+  setTimeout: (callback, delayMs) => globalThis.setTimeout(callback, delayMs),
+  clearTimeout: (handle) => globalThis.clearTimeout(handle),
+};
+
+/**
+ * Bound a refresh promise without abandoning its rejection handler when the
+ * deadline or cancellation wins. The provider promise is observed exactly
+ * once, while the observer itself always fulfills, so a late provider failure
+ * cannot become an unhandled rejection or replace the authoritative outcome.
+ */
+export async function withCodexTokenDeadline<T>(
+  operation: Promise<T>,
+  options: CodexTokenDeadlineOptions = {},
+): Promise<T> {
+  const timeoutMs = options.timeoutMs ?? CODEX_TOKEN_REFRESH_TIMEOUT_MS;
+  if (!Number.isFinite(timeoutMs) || timeoutMs <= 0) {
+    throw new Error("Codex token refresh timeout must be positive");
+  }
+  const clock = options.clock ?? systemCodexTokenDeadlineClock;
+  const signal = options.signal;
+
+  return await new Promise<T>((resolve, reject) => {
+    let settled = false;
+    let timeout: ReturnType<typeof globalThis.setTimeout> | undefined;
+
+    const cleanup = (): void => {
+      if (timeout !== undefined) {
+        clock.clearTimeout(timeout);
+        timeout = undefined;
+      }
+      signal?.removeEventListener("abort", onAbort);
+    };
+
+    const settle = (
+      outcome: { kind: "resolve"; value: T } | { kind: "reject"; error: unknown },
+    ) => {
+      if (settled) return;
+      settled = true;
+      cleanup();
+      if (outcome.kind === "resolve") {
+        resolve(outcome.value);
+      } else {
+        reject(outcome.error);
+      }
+    };
+
+    const onAbort = (): void => {
+      settle({
+        kind: "reject",
+        error: signal?.reason ?? new Error("Codex token refresh cancelled"),
+      });
+    };
+
+    if (signal?.aborted) {
+      onAbort();
+    } else {
+      signal?.addEventListener("abort", onAbort, { once: true });
+      timeout = clock.setTimeout(
+        () => settle({ kind: "reject", error: new Error("Codex token refresh timed out") }),
+        timeoutMs,
+      );
+    }
+
+    // Do not use Promise.race here. Its derived promise can obscure which
+    // branch owns settlement, while this fulfillment-only observer makes the
+    // losing provider branch explicitly consumed after timeout/cancellation.
+    void Promise.resolve(operation).then(
+      (value) => settle({ kind: "resolve", value }),
+      (error) => settle({ kind: "reject", error }),
+    );
+  });
+}
 
 // Dependencies are injectable so the lifecycle logic (single-flight, staleness,
 // needs_relogin transition) is unit-testable without a database. Production uses
@@ -96,7 +186,10 @@ export function buildCodexTokenResolver(
     cred: CodexCredentialForRun,
   ): Promise<CodexTokenSnapshot> => {
     try {
-      const next = await deps.refresh(cred.tokens.refreshToken);
+      // Bound even injected/custom refresh implementations that ignore abort
+      // signals. The provider client has its own AbortController timeout; this
+      // outer fence ensures the DB advisory transaction cannot be held forever.
+      const next = await withCodexTokenDeadline(deps.refresh(cred.tokens.refreshToken));
       const tokens = {
         access_token: next.accessToken ?? cred.tokens.accessToken,
         refresh_token: next.refreshToken ?? cred.tokens.refreshToken,
@@ -219,6 +312,7 @@ function errorUsagePayload(reason?: "needs_relogin"): CodexUsagePayload {
     weekly: null,
     limitReached: false,
     fetchedAt: new Date().toISOString(),
+    rateLimitResetCredits: null,
     ...(reason ? { reason } : {}),
   };
 }
@@ -242,6 +336,7 @@ export async function fetchCodexUsageForAccount(
   settings: Settings,
   workspaceId: string,
   credentialId: string,
+  fetchImpl: CodexFetch = fetch,
 ): Promise<CodexUsagePayload> {
   const resolver = buildCodexTokenResolver(db, settings, workspaceId, credentialId);
   let token: CodexTokenSnapshot;
@@ -253,12 +348,15 @@ export async function fetchCodexUsageForAccount(
 
   let normalized: CodexUsagePayload;
   try {
-    const usage = await fetchCodexUsage({
-      accessToken: token.accessToken,
-      chatgptAccountId: token.chatgptAccountId,
-      isFedramp: token.isFedramp,
-      clientVersion: CODEX_CLIENT_VERSION,
-    });
+    const usage = await fetchCodexUsage(
+      {
+        accessToken: token.accessToken,
+        chatgptAccountId: token.chatgptAccountId,
+        isFedramp: token.isFedramp,
+        clientVersion: CODEX_CLIENT_VERSION,
+      },
+      fetchImpl,
+    );
     normalized = normalizeCodexUsage(usage.status, usage.payload);
   } catch {
     // A network throw on the /wham/usage read must surface as an error PAYLOAD
@@ -266,17 +364,80 @@ export async function fetchCodexUsageForAccount(
     return errorUsagePayload();
   }
 
-  if (normalized.fiveHour || normalized.weekly) {
+  const parsedQuota =
+    normalized.status !== "error" && (normalized.fiveHour != null || normalized.weekly != null);
+  if (parsedQuota || normalized.rateLimitResetCredits) {
+    const checkedAt = new Date();
+    // Quota windows and reset-summary freshness are independent. A malformed
+    // usage body can still carry a syntactically valid count; that count may be
+    // cached without erasing or falsely refreshing the last valid quota truth.
     // 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.
+    // write error must NOT sink the freshly-read result 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(),
+      ...(parsedQuota
+        ? {
+            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,
+          }
+        : {}),
+      ...(normalized.rateLimitResetCredits
+        ? {
+            resetCreditAvailableCount: normalized.rateLimitResetCredits.availableCount,
+            resetCreditsCheckedAt: checkedAt,
+          }
+        : {}),
     }).catch(() => undefined);
   }
 
   return normalized;
 }
+
+export type CodexRateLimitResetCreditsAccountResult =
+  | { ok: true; status: number; details: CodexRateLimitResetCreditsDetails }
+  | {
+      ok: false;
+      status: number;
+      reason: ResetCreditFetchFailureReason | "needs_relogin";
+    };
+
+/**
+ * Fresh detailed reset-credit inventory for one exact workspace credential.
+ * The token is refreshed through the same resolver as usage and never escapes
+ * this server-side function. Detailed rows are returned to the route only and
+ * are never persisted as redemption authority.
+ */
+export async function fetchCodexRateLimitResetCreditsForAccount(
+  db: Database,
+  settings: Settings,
+  workspaceId: string,
+  credentialId: string,
+  fetchImpl: CodexFetch = fetch,
+): Promise<CodexRateLimitResetCreditsAccountResult> {
+  const resolver = buildCodexTokenResolver(db, settings, workspaceId, credentialId);
+  let token: CodexTokenSnapshot;
+  try {
+    token = await resolver.getToken();
+  } catch (error) {
+    return {
+      ok: false,
+      status: 0,
+      reason: error instanceof CodexReloginRequired ? "needs_relogin" : "network_error",
+    };
+  }
+  return await fetchCodexRateLimitResetCredits(
+    {
+      accessToken: token.accessToken,
+      chatgptAccountId: token.chatgptAccountId,
+      isFedramp: token.isFedramp,
+      clientVersion: CODEX_CLIENT_VERSION,
+    },
+    fetchImpl,
+  );
+}