@vellumai/assistant 0.8.3 → 0.8.4

package/src/providers/inference/__tests__/base-url-security.test.ts

@@ -0,0 +1,189 @@
+import { Database } from "bun:sqlite";
+import { describe, expect, mock, 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 { migrateStripBaseUrlNonOpenaiCompatible } from "../../../memory/migrations/257-strip-base-url-non-openai-compatible.js";
+import * as schema from "../../../memory/schema.js";
+import { providerConnections } from "../../../memory/schema/inference.js";
+import { getConnection } from "../connections.js";
+import { resolveAuth } from "../resolve-auth.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;
+}
+
+// ---------------------------------------------------------------------------
+// Migration: strip base_url from non-openai-compatible connections
+// ---------------------------------------------------------------------------
+
+describe("migration 257: strip base_url from non-openai-compatible connections", () => {
+  test("clears base_url on anthropic connection", () => {
+    const db = bootDb();
+
+    // Manually insert a row with a base_url on an anthropic provider (simulating
+    // a row created before the validation was added).
+    const now = Date.now();
+    db.insert(providerConnections)
+      .values({
+        name: "bad-anthropic",
+        provider: "anthropic",
+        auth: JSON.stringify({ type: "api_key", credential: "cred-abc" }),
+        status: "active",
+        baseUrl: "https://evil.example.com/v1",
+        createdAt: now,
+        updatedAt: now,
+      })
+      .run();
+
+    migrateStripBaseUrlNonOpenaiCompatible(db);
+
+    const conn = getConnection(db, "bad-anthropic");
+    expect(conn).not.toBeNull();
+    expect(conn!.baseUrl).toBeNull();
+  });
+
+  test("preserves base_url on openai-compatible connection", () => {
+    const db = bootDb();
+
+    const now = Date.now();
+    db.insert(providerConnections)
+      .values({
+        name: "good-vllm",
+        provider: "openai-compatible",
+        auth: JSON.stringify({ type: "api_key", credential: "cred-vllm" }),
+        status: "active",
+        baseUrl: "https://my-vllm.example.com/v1",
+        models: JSON.stringify([{ id: "my-model" }]),
+        createdAt: now,
+        updatedAt: now,
+      })
+      .run();
+
+    migrateStripBaseUrlNonOpenaiCompatible(db);
+
+    const conn = getConnection(db, "good-vllm");
+    expect(conn).not.toBeNull();
+    expect(conn!.baseUrl).toBe("https://my-vllm.example.com/v1");
+  });
+
+  test("is idempotent", () => {
+    const db = bootDb();
+
+    const now = Date.now();
+    db.insert(providerConnections)
+      .values({
+        name: "bad-openai",
+        provider: "openai",
+        auth: JSON.stringify({ type: "api_key", credential: "cred-abc" }),
+        status: "active",
+        baseUrl: "https://evil.example.com/v1",
+        createdAt: now,
+        updatedAt: now,
+      })
+      .run();
+
+    migrateStripBaseUrlNonOpenaiCompatible(db);
+    migrateStripBaseUrlNonOpenaiCompatible(db);
+
+    const conn = getConnection(db, "bad-openai");
+    expect(conn).not.toBeNull();
+    expect(conn!.baseUrl).toBeNull();
+  });
+});
+
+// ---------------------------------------------------------------------------
+// resolveAuth defense-in-depth: strip baseUrl for non-openai-compatible
+// ---------------------------------------------------------------------------
+
+describe("resolveAuth defense-in-depth", () => {
+  // Mock the secure key store to return a predictable value.
+  mock.module("../../../security/secure-keys.js", () => ({
+    getSecureKeyAsync: async (credential: string) =>
+      credential === "cred-test" ? "sk-test-key" : null,
+  }));
+
+  // Mock the platform proxy context to avoid real platform calls.
+  mock.module("../../../providers/platform-proxy/context.js", () => ({
+    buildManagedBaseUrl: async () => null,
+    resolveManagedProxyContext: async () => ({
+      assistantApiKey: "platform-key",
+    }),
+  }));
+
+  test("strips baseUrl for anthropic provider", async () => {
+    const result = await resolveAuth(
+      { type: "api_key", credential: "cred-test" },
+      "anthropic",
+      { baseUrl: "https://evil.example.com/v1" },
+    );
+    expect(result.ok).toBe(true);
+    if (result.ok) {
+      expect(result.resolved.kind).toBe("header");
+      if (result.resolved.kind === "header") {
+        expect(result.resolved.baseUrl).toBeUndefined();
+      }
+    }
+  });
+
+  test("strips baseUrl for openai provider", async () => {
+    const result = await resolveAuth(
+      { type: "api_key", credential: "cred-test" },
+      "openai",
+      { baseUrl: "https://evil.example.com/v1" },
+    );
+    expect(result.ok).toBe(true);
+    if (result.ok && result.resolved.kind === "header") {
+      expect(result.resolved.baseUrl).toBeUndefined();
+    }
+  });
+
+  test("strips baseUrl for gemini provider", async () => {
+    const result = await resolveAuth(
+      { type: "api_key", credential: "cred-test" },
+      "gemini",
+      { baseUrl: "https://evil.example.com/v1" },
+    );
+    expect(result.ok).toBe(true);
+    if (result.ok && result.resolved.kind === "header") {
+      expect(result.resolved.baseUrl).toBeUndefined();
+    }
+  });
+
+  test("preserves baseUrl for openai-compatible provider", async () => {
+    const result = await resolveAuth(
+      { type: "api_key", credential: "cred-test" },
+      "openai-compatible",
+      { baseUrl: "https://my-vllm.example.com/v1" },
+    );
+    expect(result.ok).toBe(true);
+    if (result.ok && result.resolved.kind === "header") {
+      expect(result.resolved.baseUrl).toBe("https://my-vllm.example.com/v1");
+    }
+  });
+
+  test("handles null baseUrl gracefully for any provider", async () => {
+    const result = await resolveAuth(
+      { type: "api_key", credential: "cred-test" },
+      "anthropic",
+      { baseUrl: null },
+    );
+    expect(result.ok).toBe(true);
+    if (result.ok && result.resolved.kind === "header") {
+      expect(result.resolved.baseUrl).toBeUndefined();
+    }
+  });
+});

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/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 }) =>
@@ -176,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", [

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;
+}