@vellumai/assistant 0.8.2 → 0.8.4

package/src/providers/inference/codex-token-refresh.ts

@@ -0,0 +1,128 @@
+/**
+ * Automatic token refresh for ChatGPT Codex OAuth (subscription auth).
+ *
+ * OpenAI rotates refresh tokens on every use — concurrent refreshes will
+ * invalidate one token. A module-level mutex prevents this race.
+ */
+
+import { refreshOAuth2Token } from "../../security/oauth2.js";
+import {
+  getSecureKeyAsync,
+  setSecureKeyAsync,
+} from "../../security/secure-keys.js";
+import { getLogger } from "../../util/logger.js";
+
+const log = getLogger("codex-token-refresh");
+
+const CODEX_TOKEN_URL = "https://auth.openai.com/oauth/token";
+const CODEX_CLIENT_ID = "app_EMoamEEZ73f0CkXaXp7hrann";
+
+/** Refresh 5 minutes before expiry to avoid using a nearly-expired token. */
+const REFRESH_MARGIN_SECONDS = 300;
+
+/**
+ * Module-level mutex to prevent concurrent refresh races.
+ * OpenAI rotates refresh tokens on every use — two concurrent refreshes
+ * will invalidate one token.
+ */
+let refreshInFlight: Promise<string | null> | null = null;
+
+/**
+ * Return a valid Codex access token, refreshing transparently if expired.
+ *
+ * @param credentialPrefix - Credential key prefix, e.g. `"credential/chatgpt"`.
+ *   The function reads `<prefix>/access_token`, `<prefix>/refresh_token`,
+ *   and `<prefix>/expires_at` from the credential store.
+ * @returns The access token string, or `null` if no token is stored.
+ */
+export async function getValidCodexAccessToken(
+  credentialPrefix: string,
+): Promise<string | null> {
+  const accessToken = await getSecureKeyAsync(
+    `${credentialPrefix}/access_token`,
+  );
+  if (!accessToken) return null;
+
+  const expiresAtStr = await getSecureKeyAsync(
+    `${credentialPrefix}/expires_at`,
+  );
+  if (!expiresAtStr) return accessToken; // no expiry info — use token as-is
+
+  const expiresAt = Number(expiresAtStr);
+  const now = Date.now() / 1000;
+
+  if (now < expiresAt - REFRESH_MARGIN_SECONDS) {
+    return accessToken; // token is still fresh
+  }
+
+  // Token is expired or about to expire — refresh it.
+  // Use mutex to prevent concurrent refresh races.
+  if (refreshInFlight) {
+    return await refreshInFlight;
+  }
+
+  refreshInFlight = doRefresh(credentialPrefix);
+  try {
+    return await refreshInFlight;
+  } finally {
+    refreshInFlight = null;
+  }
+}
+
+async function doRefresh(credentialPrefix: string): Promise<string | null> {
+  const refreshToken = await getSecureKeyAsync(
+    `${credentialPrefix}/refresh_token`,
+  );
+  if (!refreshToken) {
+    log.warn("No refresh token available for Codex OAuth");
+    // Return the existing access token — it might still work even if expired
+    return (
+      (await getSecureKeyAsync(`${credentialPrefix}/access_token`)) ?? null
+    );
+  }
+
+  try {
+    const result = await refreshOAuth2Token(
+      CODEX_TOKEN_URL,
+      CODEX_CLIENT_ID,
+      refreshToken,
+    );
+
+    // Store the new tokens
+    await setSecureKeyAsync(
+      `${credentialPrefix}/access_token`,
+      result.accessToken,
+    );
+    if (result.refreshToken) {
+      await setSecureKeyAsync(
+        `${credentialPrefix}/refresh_token`,
+        result.refreshToken,
+      );
+    }
+    if (result.expiresIn) {
+      const newExpiresAt = Math.floor(Date.now() / 1000 + result.expiresIn);
+      await setSecureKeyAsync(
+        `${credentialPrefix}/expires_at`,
+        String(newExpiresAt),
+      );
+    }
+
+    log.info("Codex OAuth token refreshed successfully");
+    return result.accessToken;
+  } catch (err) {
+    log.error({ err }, "Codex OAuth token refresh failed");
+    // Return the existing access token as fallback
+    return (
+      (await getSecureKeyAsync(`${credentialPrefix}/access_token`)) ?? null
+    );
+  }
+}
+
+// ---------------------------------------------------------------------------
+// Test helpers
+// ---------------------------------------------------------------------------
+
+/** @internal Test-only: reset the in-flight refresh mutex. */
+export function _resetRefreshMutex(): void {
+  refreshInFlight = null;
+}

package/src/providers/inference/connections.ts

@@ -1,4 +1,5 @@
 import { and, eq, isNull } from "drizzle-orm";
+import { z } from "zod";
 
 import type { DrizzleDb } from "../../memory/db-connection.js";
 import { providerConnections } from "../../memory/schema/inference.js";
@@ -6,6 +7,8 @@ import { clearConnectionProviderCache } from "../registry.js";
 import {
   type Auth,
   AuthSchema,
+  type ConnectionModel,
+  ConnectionModelSchema,
   type ConnectionProvider,
   ConnectionProviderSchema,
   type ConnectionStatus,
@@ -14,6 +17,23 @@ import {
   VALID_CONNECTION_PROVIDERS,
 } from "./auth.js";
 
+// ---------------------------------------------------------------------------
+// Helpers
+// ---------------------------------------------------------------------------
+
+function parseModelsColumn(raw: string | null): ConnectionModel[] | null {
+  if (raw === null || raw === "") return null;
+  try {
+    const parsed = z.array(ConnectionModelSchema).safeParse(JSON.parse(raw));
+    return parsed.success ? parsed.data : null;
+  } catch {
+    return null;
+  }
+}
+
+export const PROVIDERS_REQUIRING_BASE_URL_AND_MODELS: ReadonlySet<string> =
+  new Set(["openai-compatible"]);
+
 // ---------------------------------------------------------------------------
 // Read
 // ---------------------------------------------------------------------------
@@ -46,6 +66,8 @@ export function listConnections(
         provider: provider.data,
         status,
         label: row.label ?? null,
+        baseUrl: row.baseUrl ?? null,
+        models: parseModelsColumn(row.models),
         isManaged: MANAGED_CONNECTION_NAMES.has(row.name),
       },
     ];
@@ -77,6 +99,8 @@ export function getConnection(
     provider: provider.data,
     status,
     label: row.label ?? null,
+    baseUrl: row.baseUrl ?? null,
+    models: parseModelsColumn(row.models),
     isManaged: MANAGED_CONNECTION_NAMES.has(row.name),
   };
 }
@@ -91,22 +115,30 @@ export type CreateConnectionInput = {
   auth: Auth;
   status?: ConnectionStatus;
   label?: string | null;
+  baseUrl?: string | null;
+  models?: ConnectionModel[] | null;
 };
 
 export type UpdateConnectionInput = {
   auth: Auth;
   status?: ConnectionStatus;
   label?: string | null;
+  baseUrl?: string | null;
+  models?: ConnectionModel[] | null;
 };
 
 export type ConnectionCreateError =
   | { code: "already_exists" }
   | { code: "invalid_provider"; provider: string }
-  | { code: "invalid_auth" };
+  | { code: "invalid_auth" }
+  | { code: "base_url_required" }
+  | { code: "models_required" };
 
 export type ConnectionUpdateError =
   | { code: "not_found" }
-  | { code: "invalid_auth" };
+  | { code: "invalid_auth" }
+  | { code: "base_url_required" }
+  | { code: "models_required" };
 
 export type ConnectionDeleteError =
   | { code: "not_found" }
@@ -143,6 +175,15 @@ export function createConnection(
 
   const status = input.status ?? "active";
   const label = input.label ?? null;
+  const baseUrl = input.baseUrl ?? null;
+  const models = input.models ?? null;
+
+  if (PROVIDERS_REQUIRING_BASE_URL_AND_MODELS.has(provider)) {
+    if (!baseUrl) return { ok: false, error: { code: "base_url_required" } };
+    if (!models || models.length === 0) {
+      return { ok: false, error: { code: "models_required" } };
+    }
+  }
 
   const now = Date.now();
   db.insert(providerConnections)
@@ -152,6 +193,8 @@ export function createConnection(
       auth: JSON.stringify(authResult.data),
       status,
       label,
+      baseUrl,
+      models: models === null ? null : JSON.stringify(models),
       createdAt: now,
       updatedAt: now,
     })
@@ -169,6 +212,8 @@ export function createConnection(
       auth: authResult.data,
       status,
       label,
+      baseUrl,
+      models,
       createdAt: now,
       updatedAt: now,
       isManaged: MANAGED_CONNECTION_NAMES.has(input.name),
@@ -193,15 +238,34 @@ export function updateConnection(
     return { ok: false, error: { code: "invalid_auth" } };
   }
 
+  const nextBaseUrl =
+    input.baseUrl !== undefined ? input.baseUrl : existing.baseUrl;
+  const nextModels =
+    input.models !== undefined ? input.models : existing.models;
+
+  if (PROVIDERS_REQUIRING_BASE_URL_AND_MODELS.has(existing.provider)) {
+    if (!nextBaseUrl)
+      return { ok: false, error: { code: "base_url_required" } };
+    if (!nextModels || nextModels.length === 0) {
+      return { ok: false, error: { code: "models_required" } };
+    }
+  }
+
   const now = Date.now();
   const setClause: {
     auth: string;
     updatedAt: number;
     status?: string;
     label?: string | null;
+    baseUrl?: string | null;
+    models?: string | null;
   } = { auth: JSON.stringify(authResult.data), updatedAt: now };
   if (input.status !== undefined) setClause.status = input.status;
   if (input.label !== undefined) setClause.label = input.label;
+  if (input.baseUrl !== undefined) setClause.baseUrl = input.baseUrl;
+  if (input.models !== undefined)
+    setClause.models =
+      input.models === null ? null : JSON.stringify(input.models);
 
   db.update(providerConnections)
     .set(setClause)
@@ -218,6 +282,8 @@ export function updateConnection(
       auth: authResult.data,
       status: input.status !== undefined ? input.status : existing.status,
       label: input.label !== undefined ? input.label : existing.label,
+      baseUrl: nextBaseUrl,
+      models: nextModels,
       updatedAt: now,
     },
   };
@@ -294,6 +360,12 @@ const CANONICAL_CONNECTIONS: Array<{
     auth: { type: "platform" },
     label: "Google Gemini",
   },
+  {
+    name: "fireworks-managed",
+    provider: "fireworks",
+    auth: { type: "platform" },
+    label: "Fireworks",
+  },
 ];
 
 /**
@@ -330,7 +402,7 @@ export const MANAGED_CONNECTION_NAMES: ReadonlySet<string> = new Set(
  *
  * Status handling: the upsert never touches `status` so user customization
  * is preserved across reboots. New rows default to `status: "active"` via the
- * column default. Off-platform installs flip the three canonical rows to
+ * column default. Off-platform installs flip canonical managed rows to
  * `status: "disabled"` ONCE at hatch time via
  * `disableManagedConnectionsForByokHatch` (called from `seedInferenceProfiles`
  * when `isHatch && !isPlatform`); subsequent boots leave whatever the user
@@ -374,7 +446,7 @@ export function seedCanonicalConnections(db: DrizzleDb): void {
 }
 
 /**
- * Flip the three canonical managed connections to `status: "disabled"` at
+ * Flip canonical managed connections to `status: "disabled"` at
  * hatch time on BYOK (off-platform) installs.
  *
  * Why hatch-time only: managed connections need platform auth that a fresh
@@ -389,10 +461,18 @@ export function seedCanonicalConnections(db: DrizzleDb): void {
  *
  * Idempotent: a second hatch (workspace reset) re-disables the rows, which
  * is the right call — re-hatch means re-onboard.
+ *
+ * When onboarding explicitly selected a managed profile, callers may exclude
+ * that selected connection so the managed route remains usable for the first
+ * post-onboarding message.
  */
-export function disableManagedConnectionsForByokHatch(db: DrizzleDb): void {
+export function disableManagedConnectionsForByokHatch(
+  db: DrizzleDb,
+  options: { excludeConnection?: string } = {},
+): void {
   const now = Date.now();
   for (const name of MANAGED_CONNECTION_NAMES) {
+    if (name === options.excludeConnection) continue;
     db.update(providerConnections)
       .set({ status: "disabled", updatedAt: now })
       .where(eq(providerConnections.name, name))

package/src/providers/inference/resolve-auth.ts

@@ -2,10 +2,11 @@
  * Resolves an `Auth` config into a `ResolvedAuth` that adapters consume.
  *
  * Resolution rules:
- *   - api_key  → fetch credential from vault → inject as bearer header
- *   - platform → build managed proxy URL and fetch the platform API key
- *   - none     → pass through with no auth headers
- *   - oauth_subscription / service_account → reject (v2 not yet shipped)
+ *   - api_key              → fetch credential from vault → inject as bearer header
+ *   - platform             → build managed proxy URL and fetch the platform API key
+ *   - none                 → pass through with no auth headers
+ *   - oauth_subscription   → fetch OAuth token from vault (with auto-refresh) → inject as bearer header
+ *   - service_account      → reject (v2 not yet shipped)
  */
 
 import {
@@ -13,7 +14,11 @@ import {
   resolveManagedProxyContext,
 } from "../../providers/platform-proxy/context.js";
 import { getSecureKeyAsync } from "../../security/secure-keys.js";
+import { getLogger } from "../../util/logger.js";
 import type { Auth, ResolvedAuth } from "./auth.js";
+import { PROVIDERS_REQUIRING_BASE_URL_AND_MODELS } from "./connections.js";
+
+const log = getLogger("resolve-auth");
 
 export type ResolveAuthError =
   | { code: "credential_not_found"; credential: string }
@@ -23,9 +28,23 @@ export type ResolveAuthError =
 export async function resolveAuth(
   auth: Auth,
   provider: string,
+  opts: { baseUrl?: string | null } = {},
 ): Promise<
   { ok: true; resolved: ResolvedAuth } | { ok: false; error: ResolveAuthError }
 > {
+  // Defense-in-depth: strip baseUrl for providers that should not accept one.
+  // The route layer rejects base_url for non-openai-compatible providers, but
+  // this guard catches any code path that bypasses route validation (e.g.
+  // corrupted DB rows, direct calls from internal code).
+  let safeBaseUrl = opts.baseUrl;
+  if (safeBaseUrl && !PROVIDERS_REQUIRING_BASE_URL_AND_MODELS.has(provider)) {
+    log.warn(
+      { provider, baseUrl: safeBaseUrl },
+      `Stripping baseUrl for provider "${provider}" — base_url is only valid for openai-compatible providers.`,
+    );
+    safeBaseUrl = null;
+  }
+
   switch (auth.type) {
     case "api_key": {
       const value = await getSecureKeyAsync(auth.credential);
@@ -40,6 +59,7 @@ export async function resolveAuth(
         resolved: {
           kind: "header",
           headers: { Authorization: `Bearer ${value}` },
+          ...(safeBaseUrl ? { baseUrl: safeBaseUrl } : {}),
         },
       };
     }
@@ -63,7 +83,32 @@ export async function resolveAuth(
     case "none":
       return { ok: true, resolved: { kind: "none" } };
 
-    case "oauth_subscription":
+    case "oauth_subscription": {
+      // Extract the credential prefix from the credential key.
+      // The credential field stores "credential/openai-codex/access_token";
+      // we need the prefix "credential/openai-codex" for the refresh logic.
+      const credentialPrefix = auth.credential.replace(/\/access_token$/, "");
+
+      const { getValidCodexAccessToken } = await import(
+        "./codex-token-refresh.js"
+      );
+      const token = await getValidCodexAccessToken(credentialPrefix);
+
+      if (!token) {
+        return {
+          ok: false,
+          error: { code: "credential_not_found", credential: auth.credential },
+        };
+      }
+      return {
+        ok: true,
+        resolved: {
+          kind: "header",
+          headers: { Authorization: `Bearer ${token}` },
+        },
+      };
+    }
+
     case "service_account":
       return {
         ok: false,