@opengeni/db 0.27.7 → 0.27.9

package/src/codex-token-resolver.ts

@@ -3,7 +3,7 @@
 // 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
+// worker re-exports buildCodexTokenResolver from @opengeni/db for back-compat, so
 // the agent-turn.ts call site is unchanged.
 //
 // Why @opengeni/db is the right home: the resolver only orchestrates accessors
@@ -35,15 +35,52 @@ import {
   refreshCodexToken,
 } from "@opengeni/codex";
 import { encryptEnvironmentValue } from "./environment-crypto";
-import {
-  loadCodexCredentialForRun,
-  recordCodexAccountUsage,
-  recordCodexTokenRefresh,
-  setCodexCredentialStatus,
-  withCodexCredentialRefreshLock,
-  type CodexCredentialForRun,
-  type Database,
-} from "./index";
+import type { Database } from "./database";
+
+export type CodexCredentialTokens = {
+  accessToken: string;
+  refreshToken: string;
+  idToken: string;
+};
+
+export type CodexCredentialForRun = {
+  id: string;
+  version: number;
+  workspaceId: string;
+  tokens: CodexCredentialTokens;
+  chatgptAccountId: string | null;
+  scopes: string | null;
+  planType: string | null;
+  isFedramp: boolean;
+  expiresAt: Date | null;
+  lastRefreshAt: Date | null;
+  status: string;
+  lastError: string | null;
+};
+
+export type CodexAccountUsageSnapshot = {
+  primaryUsedPercent?: number | null;
+  primaryResetAt?: Date | null;
+  secondaryUsedPercent?: number | null;
+  secondaryResetAt?: Date | null;
+  checkedAt?: Date;
+  resetCreditAvailableCount?: number | null;
+  resetCreditsCheckedAt?: Date | null;
+};
+
+type CodexCredentialRefreshInput = {
+  id: string;
+  version: number;
+  workspaceId: string;
+  credentialEncrypted: string;
+  expiresAt: Date | null;
+  lastRefreshAt: Date;
+};
+
+type CodexCredentialStatusTarget = {
+  id: string;
+  version: number;
+};
 
 // 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) —
@@ -142,25 +179,37 @@ export async function withCodexTokenDeadline<T>(
 
 // 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.
+// the root composition wrapper supplies the real db + codex functions.
 export type CodexAuthDeps = {
-  loadCredential: typeof loadCodexCredentialForRun;
-  recordRefresh: typeof recordCodexTokenRefresh;
-  setStatus: typeof setCodexCredentialStatus;
+  loadCredential: (
+    db: Database,
+    settings: Settings,
+    workspaceId: string,
+    credentialId: string,
+  ) => Promise<CodexCredentialForRun | null>;
+  recordRefresh: (db: Database, input: CodexCredentialRefreshInput) => Promise<boolean>;
+  setStatus: (
+    db: Database,
+    workspaceId: string,
+    status: "active" | "needs_relogin" | "error",
+    lastError: string | null,
+    target: CodexCredentialStatusTarget,
+  ) => Promise<boolean>;
   refresh: typeof refreshCodexToken;
   encrypt: typeof encryptEnvironmentValue;
   keyBytes: typeof environmentsEncryptionKeyBytes;
-  withRefreshLock: typeof withCodexCredentialRefreshLock;
-};
-
-const defaultDeps: CodexAuthDeps = {
-  loadCredential: loadCodexCredentialForRun,
-  recordRefresh: recordCodexTokenRefresh,
-  setStatus: setCodexCredentialStatus,
-  refresh: refreshCodexToken,
-  encrypt: encryptEnvironmentValue,
-  keyBytes: environmentsEncryptionKeyBytes,
-  withRefreshLock: withCodexCredentialRefreshLock,
+  withRefreshLock: <T>(
+    db: Database,
+    workspaceId: string,
+    credentialId: string,
+    fn: (lockedDb: Database) => Promise<T>,
+  ) => Promise<T>;
+  recordUsage?: (
+    db: Database,
+    workspaceId: string,
+    credentialId: string,
+    snapshot: CodexAccountUsageSnapshot,
+  ) => Promise<boolean>;
 };
 
 export function buildCodexTokenResolver(
@@ -173,7 +222,7 @@ export function buildCodexTokenResolver(
   // 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,
+  deps: CodexAuthDeps,
 ): { getToken: () => Promise<CodexTokenSnapshot>; refresh: () => Promise<CodexTokenSnapshot> } {
   const snapshot = (cred: CodexCredentialForRun): CodexTokenSnapshot => ({
     accessToken: cred.tokens.accessToken,
@@ -336,9 +385,10 @@ export async function fetchCodexUsageForAccount(
   settings: Settings,
   workspaceId: string,
   credentialId: string,
+  deps: CodexAuthDeps,
   fetchImpl: CodexFetch = fetch,
 ): Promise<CodexUsagePayload> {
-  const resolver = buildCodexTokenResolver(db, settings, workspaceId, credentialId);
+  const resolver = buildCodexTokenResolver(db, settings, workspaceId, credentialId, deps);
   let token: CodexTokenSnapshot;
   try {
     token = await resolver.getToken();
@@ -373,27 +423,29 @@ export async function fetchCodexUsageForAccount(
     // 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 result we are about to return.
-    await recordCodexAccountUsage(db, workspaceId, credentialId, {
-      ...(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);
+    await deps
+      .recordUsage?.(db, workspaceId, credentialId, {
+        ...(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;
@@ -418,9 +470,10 @@ export async function fetchCodexRateLimitResetCreditsForAccount(
   settings: Settings,
   workspaceId: string,
   credentialId: string,
+  deps: CodexAuthDeps,
   fetchImpl: CodexFetch = fetch,
 ): Promise<CodexRateLimitResetCreditsAccountResult> {
-  const resolver = buildCodexTokenResolver(db, settings, workspaceId, credentialId);
+  const resolver = buildCodexTokenResolver(db, settings, workspaceId, credentialId, deps);
   let token: CodexTokenSnapshot;
   try {
     token = await resolver.getToken();

package/src/connection-token-resolver.ts

@@ -4,7 +4,9 @@ import {
   type Settings,
 } from "@opengeni/config";
 import type {
+  ConnectionKind,
   ConnectionCredentialsPort,
+  ConnectionStatus,
   McpConnectionResourceScope,
   McpCredentialAuthNeededReason,
   McpCredentialsRequest,
@@ -24,14 +26,49 @@ export { isPrivateAddress } from "@opengeni/network";
 import { Buffer } from "node:buffer";
 import { isIP } from "node:net";
 import { encryptEnvironmentValue } from "./environment-crypto";
-import {
-  loadConnectionCredentialForBroker,
-  recordConnectionTokenRefresh,
-  recordConnectionUsed,
-  setConnectionStatus,
-  type ConnectionCredentialForBroker,
-  type Database,
-} from "./index";
+import type { Database } from "./database";
+
+export type ConnectionCredentialForBroker = {
+  id: string;
+  accountId: string;
+  workspaceId: string;
+  subjectId: string | null;
+  providerDomain: string;
+  kind: ConnectionKind;
+  status: ConnectionStatus;
+  credential: Record<string, unknown>;
+  grantedScopes: string[];
+  expiresAt: Date | null;
+  lastRefreshAt: Date | null;
+  version: number;
+  metadata: Record<string, unknown>;
+};
+
+export type ConnectionCredentialLookupInput = {
+  workspaceId: string;
+  connectionId?: string;
+  providerDomain: string;
+  kind?: ConnectionKind;
+  subjectId?: string | null;
+  allowSubjectOwned?: boolean;
+};
+
+export type ConnectionTokenRefreshInput = {
+  id: string;
+  version: number;
+  workspaceId: string;
+  credentialEncrypted: string;
+  expiresAt: Date | null;
+  grantedScopes?: string[];
+  lastRefreshAt: Date;
+  subjectId?: string | null;
+};
+
+export type ConnectionStatusGuard = {
+  id: string;
+  version: number;
+  subjectId?: string | null;
+};
 
 export type ResolveConnectionCredentialResult =
   | {
@@ -334,10 +371,25 @@ function parseHostCredentialExpiry(value: string | null | undefined): Date | nul
 }
 
 export type ConnectionBrokerDeps = {
-  loadCredential: typeof loadConnectionCredentialForBroker;
-  recordRefresh: typeof recordConnectionTokenRefresh;
-  setStatus: typeof setConnectionStatus;
-  recordUsed: typeof recordConnectionUsed;
+  loadCredential: (
+    db: Database,
+    settings: Settings,
+    input: ConnectionCredentialLookupInput,
+  ) => Promise<ConnectionCredentialForBroker | null>;
+  recordRefresh: (db: Database, input: ConnectionTokenRefreshInput) => Promise<boolean>;
+  setStatus: (
+    db: Database,
+    workspaceId: string,
+    status: ConnectionStatus,
+    lastError: string | null,
+    guard: ConnectionStatusGuard,
+  ) => Promise<boolean>;
+  recordUsed: (
+    db: Database,
+    workspaceId: string,
+    connectionId: string,
+    subjectId?: string | null,
+  ) => Promise<void>;
   refresh: typeof refreshOAuthConnectionCredential;
   encrypt: typeof encryptEnvironmentValue;
   keyBytes: typeof environmentsEncryptionKeyBytes;
@@ -372,17 +424,6 @@ export type RefreshTransportOptions = {
   dnsLookup?: DnsLookup;
 };
 
-const defaultDeps: ConnectionBrokerDeps = {
-  loadCredential: loadConnectionCredentialForBroker,
-  recordRefresh: recordConnectionTokenRefresh,
-  setStatus: setConnectionStatus,
-  recordUsed: recordConnectionUsed,
-  refresh: refreshOAuthConnectionCredential,
-  encrypt: encryptEnvironmentValue,
-  keyBytes: environmentsEncryptionKeyBytes,
-  now: () => new Date(),
-};
-
 const inflight = new Map<string, Promise<ConnectionCredentialForBroker>>();
 const REFRESH_WINDOW_MS = 60_000;
 const CONNECTION_REFRESH_TIMEOUT_MS = 10_000;
@@ -390,7 +431,7 @@ const CONNECTION_REFRESH_TIMEOUT_MS = 10_000;
 export function buildConnectionTokenResolver(
   db: Database,
   settings: Settings,
-  deps: ConnectionBrokerDeps = defaultDeps,
+  deps: ConnectionBrokerDeps,
   options: ConnectionTokenResolverOptions = {},
 ): (input: ResolveConnectionCredentialInput) => Promise<ResolveConnectionCredentialResult> {
   type CredentialLookupInput = Pick<
@@ -404,7 +445,7 @@ export function buildConnectionTokenResolver(
     if (subjectOwned && !input.subjectId) {
       return null;
     }
-    const request: Parameters<typeof loadConnectionCredentialForBroker>[2] = {
+    const request: ConnectionCredentialLookupInput = {
       workspaceId: input.workspaceId,
       providerDomain: input.connectionRef.providerDomain,
       allowSubjectOwned: subjectOwned,
@@ -484,7 +525,7 @@ export function buildConnectionTokenResolver(
       throw new Error("OPENGENI_ENVIRONMENTS_ENCRYPTION_KEY is not configured");
     }
     const refreshed = await deps.refresh(cred, ref, settings, options.refreshTransport);
-    const refreshRecord: Parameters<typeof recordConnectionTokenRefresh>[1] = {
+    const refreshRecord: ConnectionTokenRefreshInput = {
       id: cred.id,
       version: cred.version,
       workspaceId: cred.workspaceId,