@vellumai/assistant 0.8.1 → 0.8.3

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
 // ---------------------------------------------------------------------------
@@ -23,7 +43,11 @@ export function listConnections(
   filter?: { provider?: string },
 ): ProviderConnection[] {
   const rows = filter?.provider
-    ? db.select().from(providerConnections).where(eq(providerConnections.provider, filter.provider)).all()
+    ? db
+        .select()
+        .from(providerConnections)
+        .where(eq(providerConnections.provider, filter.provider))
+        .all()
     : db.select().from(providerConnections).all();
 
   return rows.flatMap((row) => {
@@ -32,15 +56,21 @@ export function listConnections(
     const provider = ConnectionProviderSchema.safeParse(row.provider);
     if (!provider.success) return [];
     const statusResult = ConnectionStatusSchema.safeParse(row.status);
-    const status: ConnectionStatus = statusResult.success ? statusResult.data : "active";
-    return [{
-      ...row,
-      auth: auth.data,
-      provider: provider.data,
-      status,
-      label: row.label ?? null,
-      isManaged: MANAGED_CONNECTION_NAMES.has(row.name),
-    }];
+    const status: ConnectionStatus = statusResult.success
+      ? statusResult.data
+      : "active";
+    return [
+      {
+        ...row,
+        auth: auth.data,
+        provider: provider.data,
+        status,
+        label: row.label ?? null,
+        baseUrl: row.baseUrl ?? null,
+        models: parseModelsColumn(row.models),
+        isManaged: MANAGED_CONNECTION_NAMES.has(row.name),
+      },
+    ];
   });
 }
 
@@ -60,13 +90,17 @@ export function getConnection(
   const provider = ConnectionProviderSchema.safeParse(row.provider);
   if (!provider.success) return null;
   const statusResult = ConnectionStatusSchema.safeParse(row.status);
-  const status: ConnectionStatus = statusResult.success ? statusResult.data : "active";
+  const status: ConnectionStatus = statusResult.success
+    ? statusResult.data
+    : "active";
   return {
     ...row,
     auth: auth.data,
     provider: provider.data,
     status,
     label: row.label ?? null,
+    baseUrl: row.baseUrl ?? null,
+    models: parseModelsColumn(row.models),
     isManaged: MANAGED_CONNECTION_NAMES.has(row.name),
   };
 }
@@ -81,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" }
@@ -105,9 +147,14 @@ export type ConnectionDeleteError =
 export function createConnection(
   db: DrizzleDb,
   input: CreateConnectionInput,
-): { ok: true; connection: ProviderConnection } | { ok: false; error: ConnectionCreateError } {
+):
+  | { ok: true; connection: ProviderConnection }
+  | { ok: false; error: ConnectionCreateError } {
   if (!VALID_CONNECTION_PROVIDERS.includes(input.provider as never)) {
-    return { ok: false, error: { code: "invalid_provider", provider: input.provider } };
+    return {
+      ok: false,
+      error: { code: "invalid_provider", provider: input.provider },
+    };
   }
   // Safe cast: VALID_CONNECTION_PROVIDERS.includes() guards above.
   const provider = input.provider as ConnectionProvider;
@@ -128,17 +175,30 @@ 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).values({
-    name: input.name,
-    provider,
-    auth: JSON.stringify(authResult.data),
-    status,
-    label,
-    createdAt: now,
-    updatedAt: now,
-  }).run();
+  db.insert(providerConnections)
+    .values({
+      name: input.name,
+      provider,
+      auth: JSON.stringify(authResult.data),
+      status,
+      label,
+      baseUrl,
+      models: models === null ? null : JSON.stringify(models),
+      createdAt: now,
+      updatedAt: now,
+    })
+    .run();
 
   // Invalidate per-connection adapter cache so subsequent dispatch
   // resolves the freshly-inserted row's auth.
@@ -152,6 +212,8 @@ export function createConnection(
       auth: authResult.data,
       status,
       label,
+      baseUrl,
+      models,
       createdAt: now,
       updatedAt: now,
       isManaged: MANAGED_CONNECTION_NAMES.has(input.name),
@@ -163,7 +225,9 @@ export function updateConnection(
   db: DrizzleDb,
   name: string,
   input: UpdateConnectionInput,
-): { ok: true; connection: ProviderConnection } | { ok: false; error: ConnectionUpdateError } {
+):
+  | { ok: true; connection: ProviderConnection }
+  | { ok: false; error: ConnectionUpdateError } {
   const existing = getConnection(db, name);
   if (!existing) {
     return { ok: false, error: { code: "not_found" } };
@@ -174,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)
@@ -199,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,
     },
   };
@@ -226,14 +311,20 @@ export function deleteConnection(
     return { ok: false, error: { code: "not_found" } };
   }
 
-  if (!opts.force && opts.referencingProfiles && opts.referencingProfiles.length > 0) {
+  if (
+    !opts.force &&
+    opts.referencingProfiles &&
+    opts.referencingProfiles.length > 0
+  ) {
     return {
       ok: false,
       error: { code: "has_references", count: opts.referencingProfiles.length },
     };
   }
 
-  db.delete(providerConnections).where(eq(providerConnections.name, name)).run();
+  db.delete(providerConnections)
+    .where(eq(providerConnections.name, name))
+    .run();
 
   // Evict cached adapter for the deleted connection name.
   clearConnectionProviderCache();
@@ -251,9 +342,30 @@ const CANONICAL_CONNECTIONS: Array<{
   auth: Auth;
   label: string;
 }> = [
-  { name: "anthropic-managed", provider: "anthropic", auth: { type: "platform" }, label: "Anthropic" },
-  { name: "openai-managed",    provider: "openai",    auth: { type: "platform" }, label: "OpenAI" },
-  { name: "gemini-managed",    provider: "gemini",    auth: { type: "platform" }, label: "Google Gemini" },
+  {
+    name: "anthropic-managed",
+    provider: "anthropic",
+    auth: { type: "platform" },
+    label: "Anthropic",
+  },
+  {
+    name: "openai-managed",
+    provider: "openai",
+    auth: { type: "platform" },
+    label: "OpenAI",
+  },
+  {
+    name: "gemini-managed",
+    provider: "gemini",
+    auth: { type: "platform" },
+    label: "Google Gemini",
+  },
+  {
+    name: "fireworks-managed",
+    provider: "fireworks",
+    auth: { type: "platform" },
+    label: "Fireworks",
+  },
 ];
 
 /**
@@ -290,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
@@ -323,13 +435,18 @@ export function seedCanonicalConnections(db: DrizzleDb): void {
     // user-set label.
     db.update(providerConnections)
       .set({ label, updatedAt: now })
-      .where(and(eq(providerConnections.name, name), isNull(providerConnections.label)))
+      .where(
+        and(
+          eq(providerConnections.name, name),
+          isNull(providerConnections.label),
+        ),
+      )
       .run();
   }
 }
 
 /**
- * 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
@@ -344,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

@@ -11,7 +11,7 @@
 import {
   buildManagedBaseUrl,
   resolveManagedProxyContext,
-} from "../../providers/managed-proxy/context.js";
+} from "../../providers/platform-proxy/context.js";
 import { getSecureKeyAsync } from "../../security/secure-keys.js";
 import type { Auth, ResolvedAuth } from "./auth.js";
 
@@ -23,16 +23,26 @@ export type ResolveAuthError =
 export async function resolveAuth(
   auth: Auth,
   provider: string,
-): Promise<{ ok: true; resolved: ResolvedAuth } | { ok: false; error: ResolveAuthError }> {
+  opts: { baseUrl?: string | null } = {},
+): Promise<
+  { ok: true; resolved: ResolvedAuth } | { ok: false; error: ResolveAuthError }
+> {
   switch (auth.type) {
     case "api_key": {
       const value = await getSecureKeyAsync(auth.credential);
       if (!value) {
-        return { ok: false, error: { code: "credential_not_found", credential: auth.credential } };
+        return {
+          ok: false,
+          error: { code: "credential_not_found", credential: auth.credential },
+        };
       }
       return {
         ok: true,
-        resolved: { kind: "header", headers: { Authorization: `Bearer ${value}` } },
+        resolved: {
+          kind: "header",
+          headers: { Authorization: `Bearer ${value}` },
+          ...(opts.baseUrl ? { baseUrl: opts.baseUrl } : {}),
+        },
       };
     }