@vellumai/assistant 0.8.2 → 0.8.4

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

@@ -0,0 +1,254 @@
+import {
+  afterEach,
+  beforeEach,
+  describe,
+  expect,
+  mock,
+  test,
+} from "bun:test";
+
+// ---------------------------------------------------------------------------
+// Mocks
+// ---------------------------------------------------------------------------
+
+mock.module("../../../util/logger.js", () => ({
+  getLogger: () =>
+    new Proxy({} as Record<string, unknown>, {
+      get: () => () => {},
+    }),
+}));
+
+const mockGetSecureKey = mock<(key: string) => Promise<string | undefined>>(
+  async () => undefined,
+);
+const mockSetSecureKey = mock<(key: string, value: string) => Promise<boolean>>(
+  async () => true,
+);
+
+mock.module("../../../security/secure-keys.js", () => ({
+  getSecureKeyAsync: mockGetSecureKey,
+  setSecureKeyAsync: mockSetSecureKey,
+}));
+
+const mockRefreshOAuth2Token = mock<
+  (...args: unknown[]) => Promise<{
+    accessToken: string;
+    refreshToken?: string;
+    expiresIn?: number;
+  }>
+>(async () => ({
+  accessToken: "new-access-token",
+  refreshToken: "new-refresh-token",
+  expiresIn: 3600,
+}));
+
+mock.module("../../../security/oauth2.js", () => ({
+  refreshOAuth2Token: mockRefreshOAuth2Token,
+}));
+
+// ---------------------------------------------------------------------------
+// Import under test (after mocks)
+// ---------------------------------------------------------------------------
+
+import {
+  _resetRefreshMutex,
+  getValidCodexAccessToken,
+} from "../codex-token-refresh.js";
+
+const PREFIX = "credential/openai-codex";
+
+// ---------------------------------------------------------------------------
+// Helpers
+// ---------------------------------------------------------------------------
+
+function setSecureKeyMap(map: Record<string, string>): void {
+  mockGetSecureKey.mockImplementation(async (key: string) => map[key]);
+}
+
+function futureTimestamp(secondsFromNow: number): string {
+  return String(Math.floor(Date.now() / 1000) + secondsFromNow);
+}
+
+function pastTimestamp(secondsAgo: number): string {
+  return String(Math.floor(Date.now() / 1000) - secondsAgo);
+}
+
+// ---------------------------------------------------------------------------
+// Tests
+// ---------------------------------------------------------------------------
+
+describe("getValidCodexAccessToken", () => {
+  beforeEach(() => {
+    mockGetSecureKey.mockReset();
+    mockSetSecureKey.mockReset();
+    mockRefreshOAuth2Token.mockReset();
+    _resetRefreshMutex();
+
+    mockSetSecureKey.mockImplementation(async () => true);
+    mockRefreshOAuth2Token.mockImplementation(async () => ({
+      accessToken: "new-access-token",
+      refreshToken: "new-refresh-token",
+      expiresIn: 3600,
+    }));
+  });
+
+  afterEach(() => {
+    _resetRefreshMutex();
+  });
+
+  test("returns null when no access token is stored", async () => {
+    setSecureKeyMap({});
+    const result = await getValidCodexAccessToken(PREFIX);
+    expect(result).toBeNull();
+  });
+
+  test("returns access token when no expires_at is stored (graceful degradation)", async () => {
+    setSecureKeyMap({
+      [`${PREFIX}/access_token`]: "my-token",
+    });
+    const result = await getValidCodexAccessToken(PREFIX);
+    expect(result).toBe("my-token");
+    expect(mockRefreshOAuth2Token).not.toHaveBeenCalled();
+  });
+
+  test("returns access token when not expired", async () => {
+    setSecureKeyMap({
+      [`${PREFIX}/access_token`]: "my-token",
+      [`${PREFIX}/expires_at`]: futureTimestamp(600), // 10 minutes from now
+    });
+    const result = await getValidCodexAccessToken(PREFIX);
+    expect(result).toBe("my-token");
+    expect(mockRefreshOAuth2Token).not.toHaveBeenCalled();
+  });
+
+  test("refreshes token when expired", async () => {
+    const keys: Record<string, string> = {
+      [`${PREFIX}/access_token`]: "old-token",
+      [`${PREFIX}/refresh_token`]: "old-refresh",
+      [`${PREFIX}/expires_at`]: pastTimestamp(60), // expired 1 minute ago
+    };
+    setSecureKeyMap(keys);
+
+    const result = await getValidCodexAccessToken(PREFIX);
+
+    expect(result).toBe("new-access-token");
+    expect(mockRefreshOAuth2Token).toHaveBeenCalledTimes(1);
+    expect(mockRefreshOAuth2Token).toHaveBeenCalledWith(
+      "https://auth.openai.com/oauth/token",
+      "app_EMoamEEZ73f0CkXaXp7hrann",
+      "old-refresh",
+    );
+
+    // Verify new tokens are stored
+    expect(mockSetSecureKey).toHaveBeenCalledWith(
+      `${PREFIX}/access_token`,
+      "new-access-token",
+    );
+    expect(mockSetSecureKey).toHaveBeenCalledWith(
+      `${PREFIX}/refresh_token`,
+      "new-refresh-token",
+    );
+    // expires_at should be stored as well
+    const expiresAtCall = mockSetSecureKey.mock.calls.find(
+      (c) => c[0] === `${PREFIX}/expires_at`,
+    );
+    expect(expiresAtCall).toBeDefined();
+    const storedExpiresAt = Number(expiresAtCall![1]);
+    const now = Math.floor(Date.now() / 1000);
+    // Should be approximately now + 3600 (within 5 seconds tolerance)
+    expect(storedExpiresAt).toBeGreaterThan(now + 3590);
+    expect(storedExpiresAt).toBeLessThanOrEqual(now + 3610);
+  });
+
+  test("refreshes token when about to expire (within 5-minute margin)", async () => {
+    setSecureKeyMap({
+      [`${PREFIX}/access_token`]: "old-token",
+      [`${PREFIX}/refresh_token`]: "old-refresh",
+      [`${PREFIX}/expires_at`]: futureTimestamp(60), // only 1 minute left
+    });
+
+    const result = await getValidCodexAccessToken(PREFIX);
+
+    expect(result).toBe("new-access-token");
+    expect(mockRefreshOAuth2Token).toHaveBeenCalledTimes(1);
+  });
+
+  test("concurrent refresh calls are deduplicated (mutex)", async () => {
+    setSecureKeyMap({
+      [`${PREFIX}/access_token`]: "old-token",
+      [`${PREFIX}/refresh_token`]: "old-refresh",
+      [`${PREFIX}/expires_at`]: pastTimestamp(60),
+    });
+
+    // Use a deferred promise to control when the refresh completes.
+    // We create it upfront so the mock captures it synchronously.
+    let resolveRefresh!: (v: {
+      accessToken: string;
+      refreshToken: string;
+      expiresIn: number;
+    }) => void;
+    const refreshPromise = new Promise<{
+      accessToken: string;
+      refreshToken: string;
+      expiresIn: number;
+    }>((resolve) => {
+      resolveRefresh = resolve;
+    });
+
+    mockRefreshOAuth2Token.mockImplementation(() => refreshPromise);
+
+    // Fire two concurrent refreshes
+    const p1 = getValidCodexAccessToken(PREFIX);
+    const p2 = getValidCodexAccessToken(PREFIX);
+
+    // Allow the async get-key calls to settle before resolving refresh
+    await new Promise((r) => setTimeout(r, 10));
+
+    // Resolve the single in-flight refresh
+    resolveRefresh({
+      accessToken: "shared-new-token",
+      refreshToken: "shared-new-refresh",
+      expiresIn: 3600,
+    });
+
+    const [r1, r2] = await Promise.all([p1, p2]);
+
+    expect(r1).toBe("shared-new-token");
+    expect(r2).toBe("shared-new-token");
+    // Only one refresh call should have been made
+    expect(mockRefreshOAuth2Token).toHaveBeenCalledTimes(1);
+  });
+
+  test("falls back to existing token when refresh fails", async () => {
+    setSecureKeyMap({
+      [`${PREFIX}/access_token`]: "old-token",
+      [`${PREFIX}/refresh_token`]: "old-refresh",
+      [`${PREFIX}/expires_at`]: pastTimestamp(60),
+    });
+
+    mockRefreshOAuth2Token.mockImplementation(async () => {
+      throw new Error("OAuth2 token refresh failed (HTTP 400: invalid_grant)");
+    });
+
+    const result = await getValidCodexAccessToken(PREFIX);
+
+    // Should fall back to the existing access token
+    expect(result).toBe("old-token");
+    expect(mockRefreshOAuth2Token).toHaveBeenCalledTimes(1);
+  });
+
+  test("falls back to existing token when no refresh token available", async () => {
+    setSecureKeyMap({
+      [`${PREFIX}/access_token`]: "old-token",
+      // no refresh_token
+      [`${PREFIX}/expires_at`]: pastTimestamp(60),
+    });
+
+    const result = await getValidCodexAccessToken(PREFIX);
+
+    // Should return the existing access token
+    expect(result).toBe("old-token");
+    // Should not attempt a refresh
+    expect(mockRefreshOAuth2Token).not.toHaveBeenCalled();
+  });
+});

package/src/providers/inference/__tests__/connections-openai-compatible.test.ts

@@ -0,0 +1,175 @@
+import { Database } from "bun:sqlite";
+import { describe, expect, test } from "bun:test";
+
+import { drizzle } from "drizzle-orm/bun-sqlite";
+
+import { migrateCreateProviderConnections } from "../../../memory/migrations/243-provider-connections.js";
+import { migrateProviderConnectionStatusLabel } from "../../../memory/migrations/244-provider-connection-status-label.js";
+import { migrateProviderConnectionBaseUrlAndModels } from "../../../memory/migrations/250-provider-connection-base-url-and-models.js";
+import * as schema from "../../../memory/schema.js";
+import {
+  createConnection,
+  getConnection,
+  listConnections,
+  updateConnection,
+} from "../connections.js";
+
+function createTestDb() {
+  const sqlite = new Database(":memory:");
+  sqlite.exec("PRAGMA journal_mode=WAL");
+  return drizzle(sqlite, { schema });
+}
+
+function bootDb() {
+  const db = createTestDb();
+  migrateCreateProviderConnections(db);
+  migrateProviderConnectionStatusLabel(db);
+  migrateProviderConnectionBaseUrlAndModels(db);
+  return db;
+}
+
+describe("openai-compatible connection CRUD", () => {
+  test("create requires base_url", () => {
+    const db = bootDb();
+    const result = createConnection(db, {
+      name: "my-vllm",
+      provider: "openai-compatible",
+      auth: { type: "api_key", credential: "cred-vllm" },
+      models: [{ id: "my-model" }],
+    });
+    expect(result.ok).toBe(false);
+    if (!result.ok) {
+      expect(result.error.code).toBe("base_url_required");
+    }
+  });
+
+  test("create requires at least one model", () => {
+    const db = bootDb();
+    const result = createConnection(db, {
+      name: "my-vllm",
+      provider: "openai-compatible",
+      auth: { type: "api_key", credential: "cred-vllm" },
+      baseUrl: "http://localhost:8080/v1",
+    });
+    expect(result.ok).toBe(false);
+    if (!result.ok) {
+      expect(result.error.code).toBe("models_required");
+    }
+  });
+
+  test("create rejects empty models array", () => {
+    const db = bootDb();
+    const result = createConnection(db, {
+      name: "my-vllm",
+      provider: "openai-compatible",
+      auth: { type: "api_key", credential: "cred-vllm" },
+      baseUrl: "http://localhost:8080/v1",
+      models: [],
+    });
+    expect(result.ok).toBe(false);
+    if (!result.ok) {
+      expect(result.error.code).toBe("models_required");
+    }
+  });
+
+  test("create persists baseUrl and models; round-trips via getConnection", () => {
+    const db = bootDb();
+    const models = [
+      { id: "llama-3-70b" },
+      { id: "mistral-7b", displayName: "Mistral 7B" },
+    ];
+    const result = createConnection(db, {
+      name: "my-vllm",
+      provider: "openai-compatible",
+      auth: { type: "api_key", credential: "cred-vllm" },
+      baseUrl: "http://localhost:8080/v1",
+      models,
+    });
+    expect(result.ok).toBe(true);
+    if (result.ok) {
+      expect(result.connection.baseUrl).toBe("http://localhost:8080/v1");
+      expect(result.connection.models).toEqual(models);
+    }
+
+    const fetched = getConnection(db, "my-vllm");
+    expect(fetched).not.toBeNull();
+    expect(fetched!.baseUrl).toBe("http://localhost:8080/v1");
+    expect(fetched!.models).toEqual(models);
+  });
+
+  test("non-openai-compatible providers leave baseUrl/models null", () => {
+    const db = bootDb();
+    const result = createConnection(db, {
+      name: "my-anthropic",
+      provider: "anthropic",
+      auth: { type: "api_key", credential: "cred-anthropic" },
+    });
+    expect(result.ok).toBe(true);
+    if (result.ok) {
+      expect(result.connection.baseUrl).toBeNull();
+      expect(result.connection.models).toBeNull();
+    }
+  });
+
+  test("updateConnection can change models without re-supplying baseUrl", () => {
+    const db = bootDb();
+    createConnection(db, {
+      name: "my-vllm",
+      provider: "openai-compatible",
+      auth: { type: "api_key", credential: "cred-vllm" },
+      baseUrl: "http://localhost:8080/v1",
+      models: [{ id: "llama-3-70b" }],
+    });
+
+    const result = updateConnection(db, "my-vllm", {
+      auth: { type: "api_key", credential: "cred-vllm" },
+      models: [{ id: "llama-3-70b" }, { id: "mistral-7b" }],
+    });
+    expect(result.ok).toBe(true);
+    if (result.ok) {
+      expect(result.connection.baseUrl).toBe("http://localhost:8080/v1");
+      expect(result.connection.models).toEqual([
+        { id: "llama-3-70b" },
+        { id: "mistral-7b" },
+      ]);
+    }
+  });
+
+  test("updateConnection rejects clearing models on openai-compatible", () => {
+    const db = bootDb();
+    createConnection(db, {
+      name: "my-vllm",
+      provider: "openai-compatible",
+      auth: { type: "api_key", credential: "cred-vllm" },
+      baseUrl: "http://localhost:8080/v1",
+      models: [{ id: "llama-3-70b" }],
+    });
+
+    const result = updateConnection(db, "my-vllm", {
+      auth: { type: "api_key", credential: "cred-vllm" },
+      models: [],
+    });
+    expect(result.ok).toBe(false);
+    if (!result.ok) {
+      expect(result.error.code).toBe("models_required");
+    }
+  });
+
+  test("listConnections includes baseUrl and models", () => {
+    const db = bootDb();
+    createConnection(db, {
+      name: "my-vllm",
+      provider: "openai-compatible",
+      auth: { type: "api_key", credential: "cred-vllm" },
+      baseUrl: "http://localhost:8080/v1",
+      models: [{ id: "llama-3-70b" }],
+    });
+
+    const connections = listConnections(db, {
+      provider: "openai-compatible",
+    });
+    expect(connections.length).toBe(1);
+    expect(connections[0]!.baseUrl).toBe("http://localhost:8080/v1");
+    expect(connections[0]!.models).toEqual([{ id: "llama-3-70b" }]);
+  });
+});

package/src/providers/inference/__tests__/connections-status-label.test.ts

@@ -5,6 +5,7 @@ import { drizzle } from "drizzle-orm/bun-sqlite";
 
 import { migrateCreateProviderConnections } from "../../../memory/migrations/243-provider-connections.js";
 import { migrateProviderConnectionStatusLabel } from "../../../memory/migrations/244-provider-connection-status-label.js";
+import { migrateProviderConnectionBaseUrlAndModels } from "../../../memory/migrations/250-provider-connection-base-url-and-models.js";
 import * as schema from "../../../memory/schema.js";
 import {
   createConnection,
@@ -25,6 +26,7 @@ function bootDb() {
   const db = createTestDb();
   migrateCreateProviderConnections(db);
   migrateProviderConnectionStatusLabel(db);
+  migrateProviderConnectionBaseUrlAndModels(db);
   return db;
 }
 
@@ -199,6 +201,19 @@ describe("disableManagedConnectionsForByokHatch", () => {
     expect(getConnection(db, "gemini-managed")?.status).toBe("disabled");
   });
 
+  test("leaves an excluded managed connection active", () => {
+    const db = bootDb();
+    seedCanonicalConnections(db);
+
+    disableManagedConnectionsForByokHatch(db, {
+      excludeConnection: "anthropic-managed",
+    });
+
+    expect(getConnection(db, "anthropic-managed")?.status).toBe("active");
+    expect(getConnection(db, "openai-managed")?.status).toBe("disabled");
+    expect(getConnection(db, "gemini-managed")?.status).toBe("disabled");
+  });
+
   test("subsequent seedCanonicalConnections call does NOT re-flip a user-re-enabled connection", () => {
     // Models the post-hatch lifecycle: at hatch we disable; the user
     // later flips one back to active (e.g. after Vellum login). Every

package/src/providers/inference/adapter-factory.ts

@@ -41,6 +41,8 @@ export interface AdapterCreateOpts {
   baseURL?: string;
   /** Forwarded to providers that wire native provider-side web search. */
   useNativeWebSearch: boolean;
+  /** When true, the OpenAI adapter targets the Codex subscription endpoint. */
+  codexSubscription?: boolean;
 }
 
 type AdapterFactory = (opts: AdapterCreateOpts) => Provider;
@@ -65,10 +67,18 @@ const ADAPTER_FACTORIES: Record<string, AdapterFactory> = {
       streamTimeoutMs,
       ...(baseURL ? { baseURL } : {}),
     }),
-  openai: ({ apiKey, model, streamTimeoutMs, baseURL, useNativeWebSearch }) =>
+  openai: ({
+    apiKey,
+    model,
+    streamTimeoutMs,
+    baseURL,
+    useNativeWebSearch,
+    codexSubscription,
+  }) =>
     new OpenAIResponsesProvider(apiKey, model, {
       useNativeWebSearch,
       streamTimeoutMs,
+      codexSubscription,
       ...(baseURL ? { baseURL } : {}),
     }),
   gemini: ({ apiKey, model, streamTimeoutMs, baseURL }) =>
@@ -84,33 +94,22 @@ const ADAPTER_FACTORIES: Record<string, AdapterFactory> = {
       apiKey: apiKey || undefined,
       streamTimeoutMs,
     }),
-  fireworks: ({ apiKey, model, streamTimeoutMs }) =>
-    new FireworksProvider(apiKey, model, { streamTimeoutMs }),
+  fireworks: ({ apiKey, model, streamTimeoutMs, baseURL }) =>
+    new FireworksProvider(apiKey, model, {
+      streamTimeoutMs,
+      ...(baseURL ? { baseURL } : {}),
+    }),
   openrouter: ({ apiKey, model, streamTimeoutMs, useNativeWebSearch }) =>
     new OpenRouterProvider(apiKey, model, {
       useNativeWebSearch,
       streamTimeoutMs,
     }),
-  zai: ({ apiKey, model, streamTimeoutMs }) =>
+  "openai-compatible": ({ apiKey, model, streamTimeoutMs, baseURL }) =>
     new OpenAIChatCompletionsProvider(apiKey, model, {
-      providerName: "zai",
-      providerLabel: "z.ai",
-      baseURL: "https://api.z.ai/api/paas/v4/",
-      streamTimeoutMs,
-    }),
-  deepseek: ({ apiKey, model, streamTimeoutMs }) =>
-    new OpenAIChatCompletionsProvider(apiKey, model, {
-      providerName: "deepseek",
-      providerLabel: "DeepSeek",
-      baseURL: "https://api.deepseek.com",
-      streamTimeoutMs,
-    }),
-  minimax: ({ apiKey, model, streamTimeoutMs }) =>
-    new OpenAIChatCompletionsProvider(apiKey, model, {
-      providerName: "minimax",
-      providerLabel: "MiniMax",
-      baseURL: "https://api.minimax.io/v1",
+      providerName: "openai-compatible",
+      providerLabel: "OpenAI-compatible",
       streamTimeoutMs,
+      ...(baseURL ? { baseURL } : {}),
     }),
 };
 
@@ -187,12 +186,16 @@ export function createAdapterFromConnection(
   const baseURL =
     resolvedAuth.kind === "header" ? resolvedAuth.baseUrl : undefined;
 
+  const codexSubscription =
+    connection.auth.type === "oauth_subscription" && provider === "openai";
+
   const adapter = buildProviderAdapter(provider, {
     apiKey,
     model: opts.model,
     streamTimeoutMs: opts.streamTimeoutMs ?? 1_800_000,
     baseURL,
     useNativeWebSearch: opts.useNativeWebSearch ?? false,
+    codexSubscription,
   });
   if (!adapter) return null;
 

package/src/providers/inference/auth.ts

@@ -9,13 +9,13 @@ import { PROVIDER_CATALOG } from "../model-catalog.js";
 /**
  * Auth configuration stored in the `provider_connections` table.
  *
- * v1 runtime-supported variants:
+ * Runtime-supported variants:
  *   - api_key: look up `credential` in vault, inject as bearer/provider header.
  *   - platform: route via Vellum managed proxy; no client-side credential.
  *   - none: no auth (e.g. Ollama running locally).
+ *   - oauth_subscription: OAuth-based subscription auth (e.g. ChatGPT Codex).
  *
- * v2 schema-accepted variants (runtime rejects with a clear "not yet shipped" error):
- *   - oauth_subscription: OAuth-based subscription auth.
+ * Schema-accepted variants (runtime rejects with a clear "not yet shipped" error):
  *   - service_account: service-account credentials (Vertex AI, Bedrock).
  */
 export const AuthSchema = z.discriminatedUnion("type", [
@@ -86,6 +86,16 @@ export const ConnectionProviderSchema = z.enum(
 export const ConnectionStatusSchema = z.enum(["active", "disabled"]);
 export type ConnectionStatus = z.infer<typeof ConnectionStatusSchema>;
 
+// ---------------------------------------------------------------------------
+// Per-connection model entries (openai-compatible)
+// ---------------------------------------------------------------------------
+
+export const ConnectionModelSchema = z.object({
+  id: z.string().min(1),
+  displayName: z.string().min(1).optional(),
+});
+export type ConnectionModel = z.infer<typeof ConnectionModelSchema>;
+
 // ---------------------------------------------------------------------------
 // Full connection shape used by CRUD layer
 // ---------------------------------------------------------------------------
@@ -96,6 +106,8 @@ export const ProviderConnectionSchema = z.object({
   auth: AuthSchema,
   status: ConnectionStatusSchema,
   label: z.string().min(1).nullable(),
+  baseUrl: z.string().url().nullable(),
+  models: z.array(ConnectionModelSchema).nullable(),
   createdAt: z.number().int(),
   updatedAt: z.number().int(),
   /**

package/src/providers/inference/backfill.ts

@@ -21,7 +21,12 @@ import { loadRawConfig, saveRawConfig } from "../../config/loader.js";
 import type { DrizzleDb } from "../../memory/db-connection.js";
 import { credentialKey } from "../../security/credential-key.js";
 import { getLogger } from "../../util/logger.js";
-import { createConnection, getConnection, seedCanonicalConnections } from "./connections.js";
+import {
+  createConnection,
+  getConnection,
+  PROVIDERS_REQUIRING_BASE_URL_AND_MODELS,
+  seedCanonicalConnections,
+} from "./connections.js";
 
 const log = getLogger("provider-connections-backfill");
 
@@ -149,6 +154,14 @@ function ensureProviderConnection(
   const provider = entry.provider as string | undefined;
   if (!provider) return false;
 
+  if (PROVIDERS_REQUIRING_BASE_URL_AND_MODELS.has(provider)) {
+    log.warn(
+      { entry: entryLabel, provider },
+      "Skipping backfill for provider that requires per-connection base_url/models",
+    );
+    return false;
+  }
+
   let connectionName: string;
 
   if (globalMode === "managed" && MANAGED_PROVIDERS.has(provider)) {