@vellumai/assistant 0.7.3 → 0.8.1

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

@@ -0,0 +1,250 @@
+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 * as schema from "../../../memory/schema.js";
+import {
+  createConnection,
+  disableManagedConnectionsForByokHatch,
+  getConnection,
+  listConnections,
+  seedCanonicalConnections,
+  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);
+  return db;
+}
+
+describe("connection CRUD status + label defaults", () => {
+  test("new connection without status/label gets status=active and label=null", () => {
+    const db = bootDb();
+    const result = createConnection(db, {
+      name: "my-conn",
+      provider: "anthropic",
+      auth: { type: "platform" },
+    });
+    expect(result.ok).toBe(true);
+    if (result.ok) {
+      expect(result.connection.status).toBe("active");
+      expect(result.connection.label).toBeNull();
+    }
+  });
+
+  test("createConnection passes explicit status and label", () => {
+    const db = bootDb();
+    const result = createConnection(db, {
+      name: "disabled-conn",
+      provider: "openai",
+      auth: { type: "platform" },
+      status: "disabled",
+      label: "My OpenAI",
+    });
+    expect(result.ok).toBe(true);
+    if (result.ok) {
+      expect(result.connection.status).toBe("disabled");
+      expect(result.connection.label).toBe("My OpenAI");
+    }
+  });
+
+  test("getConnection returns status and label from DB", () => {
+    const db = bootDb();
+    createConnection(db, {
+      name: "get-me",
+      provider: "gemini",
+      auth: { type: "platform" },
+      status: "disabled",
+      label: "Gemini Pro",
+    });
+
+    const conn = getConnection(db, "get-me");
+    expect(conn).not.toBeNull();
+    expect(conn!.status).toBe("disabled");
+    expect(conn!.label).toBe("Gemini Pro");
+  });
+
+  test("updateConnection updates status", () => {
+    const db = bootDb();
+    createConnection(db, {
+      name: "toggle-me",
+      provider: "anthropic",
+      auth: { type: "platform" },
+    });
+
+    const result = updateConnection(db, "toggle-me", {
+      auth: { type: "platform" },
+      status: "disabled",
+    });
+    expect(result.ok).toBe(true);
+    if (result.ok) {
+      expect(result.connection.status).toBe("disabled");
+    }
+  });
+
+  test("updateConnection clears label when set to null", () => {
+    const db = bootDb();
+    createConnection(db, {
+      name: "clear-label",
+      provider: "openai",
+      auth: { type: "platform" },
+      label: "Old Label",
+    });
+
+    const result = updateConnection(db, "clear-label", {
+      auth: { type: "platform" },
+      label: null,
+    });
+    expect(result.ok).toBe(true);
+    if (result.ok) {
+      expect(result.connection.label).toBeNull();
+    }
+  });
+});
+
+describe("seedCanonicalConnections labels", () => {
+  test("first boot seeds default labels on all three managed connections", () => {
+    const db = bootDb();
+    seedCanonicalConnections(db);
+
+    const conns = listConnections(db);
+    const byName = Object.fromEntries(conns.map((c) => [c.name, c]));
+
+    expect(byName["anthropic-managed"]?.label).toBe("Anthropic");
+    expect(byName["openai-managed"]?.label).toBe("OpenAI");
+    expect(byName["gemini-managed"]?.label).toBe("Google Gemini");
+  });
+
+  test("second boot preserves user-customized label", () => {
+    const db = bootDb();
+    seedCanonicalConnections(db);
+
+    // User customizes the label.
+    updateConnection(db, "anthropic-managed", {
+      auth: { type: "platform" },
+      label: "Work Anthropic",
+    });
+
+    // Reboot.
+    seedCanonicalConnections(db);
+
+    const conn = getConnection(db, "anthropic-managed");
+    expect(conn?.label).toBe("Work Anthropic");
+  });
+
+  test("second boot backfills default label when existing row has null label", () => {
+    const db = bootDb();
+
+    // `bootDb()` runs migration 243 which already inserted the three
+    // canonical rows with `label=null` (the label column was added by 244
+    // and defaults NULL for pre-existing rows). This matches the state
+    // every pre-label install carries forward into the boot that ships
+    // the label seed.
+    const before = getConnection(db, "anthropic-managed");
+    expect(before?.label).toBeNull();
+
+    seedCanonicalConnections(db);
+
+    const after = getConnection(db, "anthropic-managed");
+    expect(after?.label).toBe("Anthropic");
+  });
+
+  test("backfill leaves explicit empty-overwrite null untouched on subsequent boot", () => {
+    const db = bootDb();
+    seedCanonicalConnections(db);
+
+    // User clears the label (PATCH label: null).
+    updateConnection(db, "openai-managed", {
+      auth: { type: "platform" },
+      label: null,
+    });
+
+    // Subsequent boots refill it — there's no distinction between "user
+    // explicitly cleared" and "pre-seed row that never had one". Treating
+    // both as "fill with default" is intentional; users who want a blank
+    // label aren't a real cohort and we'd rather guarantee the default is
+    // present for everyone.
+    seedCanonicalConnections(db);
+
+    const conn = getConnection(db, "openai-managed");
+    expect(conn?.label).toBe("OpenAI");
+  });
+});
+
+describe("disableManagedConnectionsForByokHatch", () => {
+  test("flips all three canonical managed connections to status='disabled'", () => {
+    const db = bootDb();
+    seedCanonicalConnections(db);
+
+    // Sanity: seeded rows default to active.
+    expect(getConnection(db, "anthropic-managed")?.status).toBe("active");
+    expect(getConnection(db, "openai-managed")?.status).toBe("active");
+    expect(getConnection(db, "gemini-managed")?.status).toBe("active");
+
+    disableManagedConnectionsForByokHatch(db);
+
+    expect(getConnection(db, "anthropic-managed")?.status).toBe("disabled");
+    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
+    // subsequent daemon boot runs seedCanonicalConnections — and that
+    // boot must NOT revert the user's choice. The hatch-disable helper
+    // is only ever called from the seedInferenceProfiles hatch branch,
+    // so it does not run on a non-hatch boot; this test confirms the
+    // ambient seed pass leaves status alone.
+    const db = bootDb();
+    seedCanonicalConnections(db);
+    disableManagedConnectionsForByokHatch(db);
+
+    // User re-enables anthropic post-hatch.
+    updateConnection(db, "anthropic-managed", {
+      auth: { type: "platform" },
+      status: "active",
+    });
+    expect(getConnection(db, "anthropic-managed")?.status).toBe("active");
+
+    // Simulate a normal restart: seedCanonicalConnections runs every boot,
+    // disableManagedConnectionsForByokHatch does NOT.
+    seedCanonicalConnections(db);
+
+    expect(getConnection(db, "anthropic-managed")?.status).toBe("active");
+    // The two the user didn't touch stay disabled.
+    expect(getConnection(db, "openai-managed")?.status).toBe("disabled");
+    expect(getConnection(db, "gemini-managed")?.status).toBe("disabled");
+  });
+
+  test("idempotent re-hatch leaves all three at disabled", () => {
+    // Workspace reset / re-hatch scenario: helper runs again and any
+    // user re-enable from before the reset is intentionally undone —
+    // re-hatch means re-onboard.
+    const db = bootDb();
+    seedCanonicalConnections(db);
+    disableManagedConnectionsForByokHatch(db);
+
+    updateConnection(db, "anthropic-managed", {
+      auth: { type: "platform" },
+      status: "active",
+    });
+
+    disableManagedConnectionsForByokHatch(db);
+
+    expect(getConnection(db, "anthropic-managed")?.status).toBe("disabled");
+    expect(getConnection(db, "openai-managed")?.status).toBe("disabled");
+    expect(getConnection(db, "gemini-managed")?.status).toBe("disabled");
+  });
+});

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

@@ -0,0 +1,173 @@
+/**
+ * Provider adapter construction.
+ *
+ * One catalog-keyed factory table feeds two construction paths:
+ *
+ *   - `buildProviderAdapter` returns a raw `Provider` instance for a given
+ *     provider id + options. The caller wraps with `RetryProvider` /
+ *     `UsageTrackingProvider` to match the boot-time vs per-connection
+ *     wrapping conventions in `registry.ts`.
+ *   - `createAdapterFromConnection` is the per-call dispatcher entry point.
+ *     It resolves a `ResolvedAuth` into `AdapterCreateOpts`, validates
+ *     keyless/keyed compatibility, and returns a fully-wrapped
+ *     `Provider | null`.
+ *
+ * Adding a new provider:
+ *   1. Add an entry to `PROVIDER_CATALOG` in `model-catalog.ts`.
+ *   2. Implement the client in `src/providers/<id>/client.ts`.
+ *   3. Register the client in `ADAPTER_FACTORIES` below.
+ */
+
+import { AnthropicProvider } from "../anthropic/client.js";
+import { FireworksProvider } from "../fireworks/client.js";
+import { GeminiProvider } from "../gemini/client.js";
+import { PROVIDER_CATALOG } from "../model-catalog.js";
+import { OllamaProvider } from "../ollama/client.js";
+import { OpenAIResponsesProvider } from "../openai/responses-provider.js";
+import { OpenRouterProvider } from "../openrouter/client.js";
+import { RetryProvider } from "../retry.js";
+import type { Provider } from "../types.js";
+import { UsageTrackingProvider } from "../usage-tracking.js";
+import type { ResolvedAuth } from "./auth.js";
+import type { ProviderConnection } from "./auth.js";
+
+/** Unified construction opts. Adapters ignore fields they don't consume. */
+export interface AdapterCreateOpts {
+  apiKey: string;
+  model: string;
+  streamTimeoutMs: number;
+  /** Set when an explicit base URL override or managed proxy is in play. */
+  baseURL?: string;
+  /** Forwarded to providers that wire native provider-side web search. */
+  useNativeWebSearch: boolean;
+}
+
+type AdapterFactory = (opts: AdapterCreateOpts) => Provider;
+
+/**
+ * Catalog-keyed factory table. Each entry takes a unified
+ * `AdapterCreateOpts` and constructs the underlying provider client. The
+ * `id` field must match the corresponding `ProviderCatalogEntry.id` in
+ * `PROVIDER_CATALOG` — `PROVIDER_CATALOG_FACTORY_PARITY` enforces this at
+ * module-load time.
+ */
+const ADAPTER_FACTORIES: Record<string, AdapterFactory> = {
+  anthropic: ({ apiKey, model, streamTimeoutMs, baseURL, useNativeWebSearch }) =>
+    new AnthropicProvider(apiKey, model, {
+      useNativeWebSearch,
+      streamTimeoutMs,
+      ...(baseURL ? { baseURL } : {}),
+    }),
+  openai: ({ apiKey, model, streamTimeoutMs, baseURL, useNativeWebSearch }) =>
+    new OpenAIResponsesProvider(apiKey, model, {
+      useNativeWebSearch,
+      streamTimeoutMs,
+      ...(baseURL ? { baseURL } : {}),
+    }),
+  gemini: ({ apiKey, model, streamTimeoutMs, baseURL }) =>
+    new GeminiProvider(apiKey, model, {
+      streamTimeoutMs,
+      // Gemini routes managed proxies through `managedBaseUrl`, not `baseURL`.
+      ...(baseURL ? { managedBaseUrl: baseURL } : {}),
+    }),
+  ollama: ({ apiKey, model, streamTimeoutMs }) =>
+    new OllamaProvider(model, {
+      // Empty string means keyless — Ollama's client treats undefined as
+      // "no key provided" and defaults its internal placeholder.
+      apiKey: apiKey || undefined,
+      streamTimeoutMs,
+    }),
+  fireworks: ({ apiKey, model, streamTimeoutMs }) =>
+    new FireworksProvider(apiKey, model, { streamTimeoutMs }),
+  openrouter: ({ apiKey, model, streamTimeoutMs, useNativeWebSearch }) =>
+    new OpenRouterProvider(apiKey, model, {
+      useNativeWebSearch,
+      streamTimeoutMs,
+    }),
+};
+
+/**
+ * Module-load parity guard. Surfaces a clear startup error if someone adds
+ * a catalog entry without a matching factory (or vice versa).
+ */
+const PROVIDER_CATALOG_FACTORY_PARITY = (() => {
+  const catalogIds = new Set(PROVIDER_CATALOG.map((entry) => entry.id));
+  const factoryIds = new Set(Object.keys(ADAPTER_FACTORIES));
+  const missingFactories = [...catalogIds].filter((id) => !factoryIds.has(id));
+  const orphanFactories = [...factoryIds].filter((id) => !catalogIds.has(id));
+  if (missingFactories.length > 0 || orphanFactories.length > 0) {
+    const parts: string[] = [];
+    if (missingFactories.length > 0) {
+      parts.push(`missing adapter factories: ${missingFactories.join(", ")}`);
+    }
+    if (orphanFactories.length > 0) {
+      parts.push(`orphan adapter factories: ${orphanFactories.join(", ")}`);
+    }
+    throw new Error(
+      `PROVIDER_CATALOG / ADAPTER_FACTORIES drift: ${parts.join("; ")}`,
+    );
+  }
+  return true;
+})();
+
+// Reference the parity guard so unused-variable lint doesn't strip it.
+void PROVIDER_CATALOG_FACTORY_PARITY;
+
+/**
+ * Build a raw `Provider` instance from a provider id and unified opts.
+ *
+ * Returns null when no factory exists for the given provider id. The
+ * caller is responsible for wrapping (RetryProvider, UsageTrackingProvider).
+ */
+export function buildProviderAdapter(
+  providerId: string,
+  opts: AdapterCreateOpts,
+): Provider | null {
+  const factory = ADAPTER_FACTORIES[providerId];
+  if (!factory) return null;
+  return factory(opts);
+}
+
+/**
+ * Build a Provider instance for a given connection + resolved auth.
+ *
+ * Returns null when the provider/auth combination is not usable
+ * (e.g. `none` auth on a keyed provider). The caller decides whether to
+ * log a warning or fall back to the global registry.
+ */
+export function createAdapterFromConnection(
+  connection: ProviderConnection,
+  resolvedAuth: ResolvedAuth,
+  opts: { model: string; streamTimeoutMs?: number; useNativeWebSearch?: boolean },
+): Provider | null {
+  const { provider } = connection;
+  const entry = PROVIDER_CATALOG.find((e) => e.id === provider);
+  if (!entry) return null;
+  const isKeyless = entry.setupMode === "keyless";
+
+  // Keyed providers can't operate without a credential.
+  if (!isKeyless && resolvedAuth.kind === "none") return null;
+
+  const apiKey =
+    resolvedAuth.kind === "header"
+      ? (resolvedAuth.headers["Authorization"] ?? "").replace(/^Bearer /, "")
+      : "";
+  const baseURL =
+    resolvedAuth.kind === "header" ? resolvedAuth.baseUrl : undefined;
+
+  const adapter = buildProviderAdapter(provider, {
+    apiKey,
+    model: opts.model,
+    streamTimeoutMs: opts.streamTimeoutMs ?? 1_800_000,
+    baseURL,
+    useNativeWebSearch: opts.useNativeWebSearch ?? false,
+  });
+  if (!adapter) return null;
+
+  const isProxy = baseURL !== undefined;
+  return new UsageTrackingProvider(
+    new RetryProvider(adapter, {
+      forwardUsageAttributionHeaders: isProxy,
+    }),
+  );
+}

package/src/providers/inference/auth.ts

@@ -0,0 +1,112 @@
+import { z } from "zod";
+
+import { PROVIDER_CATALOG } from "../model-catalog.js";
+
+// ---------------------------------------------------------------------------
+// Auth discriminated union (stored in provider_connections.auth as JSON)
+// ---------------------------------------------------------------------------
+
+/**
+ * Auth configuration stored in the `provider_connections` table.
+ *
+ * v1 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).
+ *
+ * v2 schema-accepted variants (runtime rejects with a clear "not yet shipped" error):
+ *   - oauth_subscription: OAuth-based subscription auth.
+ *   - service_account: service-account credentials (Vertex AI, Bedrock).
+ */
+export const AuthSchema = z.discriminatedUnion("type", [
+  z.object({
+    type: z.literal("api_key"),
+    credential: z.string().min(1),
+  }),
+  z.object({
+    type: z.literal("platform"),
+  }),
+  z.object({
+    type: z.literal("none"),
+  }),
+  z.object({
+    type: z.literal("oauth_subscription"),
+    credential: z.string().min(1),
+  }),
+  z.object({
+    type: z.literal("service_account"),
+    credential: z.string().min(1),
+  }),
+]);
+
+export type Auth = z.infer<typeof AuthSchema>;
+
+// ---------------------------------------------------------------------------
+// ResolvedAuth — what the dispatcher hands to each adapter
+// ---------------------------------------------------------------------------
+
+/**
+ * The resolved form of an Auth, produced by the dispatcher before calling
+ * an adapter. Adapters are pure functions of (ResolvedAuth, request) → response
+ * and never access the vault themselves.
+ */
+export type ResolvedAuth =
+  | { kind: "header"; headers: Record<string, string>; baseUrl?: string }
+  | { kind: "runtime_proxy"; route: string }
+  | { kind: "none" };
+
+// ---------------------------------------------------------------------------
+// Valid provider identifiers — derived from PROVIDER_CATALOG
+// ---------------------------------------------------------------------------
+//
+// PROVIDER_CATALOG (in `model-catalog.ts`) is the single source of truth for
+// the closed set of inference-provider identifiers. The list below is
+// derived at module load; adding a provider to the catalog automatically
+// extends `VALID_CONNECTION_PROVIDERS` and `ConnectionProviderSchema`.
+//
+// Trade-off: because `PROVIDER_CATALOG` is a runtime value, the
+// `ConnectionProvider` static type is `string` rather than a narrow
+// literal-string union. Callers that need a narrowed value should parse
+// through `ConnectionProviderSchema`, which still rejects unknown
+// providers at runtime.
+
+export const VALID_CONNECTION_PROVIDERS: readonly string[] =
+  PROVIDER_CATALOG.map((p) => p.id);
+
+export type ConnectionProvider = string;
+
+export const ConnectionProviderSchema = z.enum(
+  VALID_CONNECTION_PROVIDERS as readonly [string, ...string[]],
+);
+
+// ---------------------------------------------------------------------------
+// Connection status
+// ---------------------------------------------------------------------------
+
+export const ConnectionStatusSchema = z.enum(["active", "disabled"]);
+export type ConnectionStatus = z.infer<typeof ConnectionStatusSchema>;
+
+// ---------------------------------------------------------------------------
+// Full connection shape used by CRUD layer
+// ---------------------------------------------------------------------------
+
+export const ProviderConnectionSchema = z.object({
+  name: z.string().min(1),
+  provider: ConnectionProviderSchema,
+  auth: AuthSchema,
+  status: ConnectionStatusSchema,
+  label: z.string().min(1).nullable(),
+  createdAt: z.number().int(),
+  updatedAt: z.number().int(),
+  /**
+   * Whether this row is a Vellum-managed connection (`anthropic-managed`,
+   * `openai-managed`, `gemini-managed`). Derived from
+   * `MANAGED_CONNECTION_NAMES` in `connections.ts` at serialize time; the
+   * DB column does not exist. Clients use this to render the read-only
+   * "Vellum" badge + view-only editor and to disable the delete affordance
+   * without mirroring the canonical name list locally.
+   */
+  isManaged: z.boolean(),
+});
+
+export type ProviderConnection = z.infer<typeof ProviderConnectionSchema>;