@vellumai/assistant 0.5.1 → 0.5.2

package/src/oauth/connection-resolver.test.ts

@@ -0,0 +1,191 @@
+import { beforeEach, describe, expect, mock, test } from "bun:test";
+
+// ---------------------------------------------------------------------------
+// Mutable mock state
+// ---------------------------------------------------------------------------
+
+let mockProvider: Record<string, unknown> | undefined;
+let mockConnection: Record<string, unknown> | undefined;
+let mockAccessToken: string | undefined;
+let mockConfig: Record<string, unknown> = {};
+let mockManagedProxyCtx = {
+  enabled: false,
+  platformBaseUrl: "",
+  assistantApiKey: "",
+};
+let mockAssistantId = "";
+
+// ---------------------------------------------------------------------------
+// Module mocks (must precede imports of the module under test)
+// ---------------------------------------------------------------------------
+
+mock.module("../util/logger.js", () => ({
+  getLogger: () =>
+    new Proxy({} as Record<string, unknown>, {
+      get: () => () => {},
+    }),
+}));
+
+mock.module("./oauth-store.js", () => ({
+  getProvider: () => mockProvider,
+  getActiveConnection: (
+    _pk: string,
+    opts?: { clientId?: string; account?: string },
+  ) => {
+    if (opts?.clientId && mockConnection?.clientId !== opts.clientId)
+      return undefined;
+    if (opts?.account && mockConnection?.accountInfo !== opts.account)
+      return undefined;
+    return mockConnection;
+  },
+}));
+
+mock.module("../security/secure-keys.js", () => ({
+  getSecureKeyAsync: async () => mockAccessToken,
+}));
+
+mock.module("../config/loader.js", () => ({
+  getConfig: () => mockConfig,
+}));
+
+mock.module("../config/env.js", () => ({
+  getPlatformAssistantId: () => mockAssistantId,
+}));
+
+mock.module("../providers/managed-proxy/context.js", () => ({
+  resolveManagedProxyContext: async () => mockManagedProxyCtx,
+}));
+
+// ---------------------------------------------------------------------------
+// Import the module under test (after all mocks are registered)
+// ---------------------------------------------------------------------------
+
+import { BYOOAuthConnection } from "./byo-connection.js";
+import { resolveOAuthConnection } from "./connection-resolver.js";
+import { PlatformOAuthConnection } from "./platform-connection.js";
+
+// ---------------------------------------------------------------------------
+// Helpers
+// ---------------------------------------------------------------------------
+
+function setupDefaults(): void {
+  mockProvider = {
+    providerKey: "integration:google",
+    baseUrl: "https://gmail.googleapis.com/gmail/v1/users/me",
+    managedServiceConfigKey: null,
+  };
+  mockConnection = {
+    id: "conn-1",
+    providerKey: "integration:google",
+    oauthAppId: "app-1",
+    accountInfo: "user@example.com",
+    grantedScopes: JSON.stringify(["scope-a", "scope-b"]),
+    status: "active",
+    clientId: "client-1",
+  };
+  mockAccessToken = "tok-valid";
+  mockConfig = {
+    services: {
+      inference: {
+        mode: "your-own",
+        provider: "anthropic",
+        model: "claude-opus-4-6",
+      },
+      "image-generation": {
+        mode: "your-own",
+        provider: "gemini",
+        model: "gemini-3.1-flash-image-preview",
+      },
+      "web-search": { mode: "your-own", provider: "inference-provider-native" },
+      "google-oauth": { mode: "managed" },
+    },
+  };
+  mockManagedProxyCtx = {
+    enabled: true,
+    platformBaseUrl: "https://platform.example.com",
+    assistantApiKey: "sk-test-key",
+  };
+  mockAssistantId = "asst-123";
+}
+
+// ---------------------------------------------------------------------------
+// Tests
+// ---------------------------------------------------------------------------
+
+describe("resolveOAuthConnection", () => {
+  beforeEach(() => {
+    setupDefaults();
+  });
+
+  test("returns BYOOAuthConnection when provider has no managedServiceConfigKey", async () => {
+    const result = await resolveOAuthConnection("integration:google");
+    expect(result).toBeInstanceOf(BYOOAuthConnection);
+    expect(result.id).toBe("conn-1");
+    expect(result.providerKey).toBe("integration:google");
+  });
+
+  test("returns PlatformOAuthConnection when managed mode is active", async () => {
+    mockProvider!.managedServiceConfigKey = "google-oauth";
+
+    const result = await resolveOAuthConnection("integration:google");
+    expect(result).toBeInstanceOf(PlatformOAuthConnection);
+    expect(result.id).toBe("integration:google");
+    expect(result.providerKey).toBe("integration:google");
+    expect(result.accountInfo).toBeNull();
+  });
+
+  test("passes account through to PlatformOAuthConnection", async () => {
+    mockProvider!.managedServiceConfigKey = "google-oauth";
+
+    const result = await resolveOAuthConnection("integration:google", {
+      account: "user@example.com",
+    });
+    expect(result).toBeInstanceOf(PlatformOAuthConnection);
+    expect(result.accountInfo).toBe("user@example.com");
+  });
+
+  test("returns BYOOAuthConnection when service config mode is your-own", async () => {
+    mockProvider!.managedServiceConfigKey = "google-oauth";
+    (mockConfig.services as Record<string, unknown>)["google-oauth"] = {
+      mode: "your-own",
+    };
+
+    const result = await resolveOAuthConnection("integration:google");
+    expect(result).toBeInstanceOf(BYOOAuthConnection);
+    expect(result.id).toBe("conn-1");
+  });
+
+  test("managed path does not require a local connection row", async () => {
+    mockProvider!.managedServiceConfigKey = "google-oauth";
+    mockConnection = undefined;
+    mockAccessToken = undefined;
+
+    const result = await resolveOAuthConnection("integration:google");
+    expect(result).toBeInstanceOf(PlatformOAuthConnection);
+  });
+
+  test("managed path ignores clientId option", async () => {
+    mockProvider!.managedServiceConfigKey = "google-oauth";
+
+    const result = await resolveOAuthConnection("integration:google", {
+      clientId: "some-client-id",
+    });
+    expect(result).toBeInstanceOf(PlatformOAuthConnection);
+  });
+
+  test("BYO path narrows by clientId when provided", async () => {
+    const result = await resolveOAuthConnection("integration:google", {
+      clientId: "client-1",
+    });
+    expect(result).toBeInstanceOf(BYOOAuthConnection);
+    expect(result.id).toBe("conn-1");
+  });
+
+  test("BYO path returns no credential when clientId does not match", async () => {
+    await expect(
+      resolveOAuthConnection("integration:google", {
+        clientId: "wrong-client",
+      }),
+    ).rejects.toThrow(/No active OAuth connection found/);
+  });
+});

package/src/oauth/connection-resolver.ts

@@ -1,72 +1,100 @@
+import { getPlatformAssistantId } from "../config/env.js";
+import { getConfig } from "../config/loader.js";
+import { type Services, ServicesSchema } from "../config/schemas/services.js";
+import { resolveManagedProxyContext } from "../providers/managed-proxy/context.js";
 import { getSecureKeyAsync } from "../security/secure-keys.js";
 import { BYOOAuthConnection } from "./byo-connection.js";
 import type { OAuthConnection } from "./connection.js";
-import {
-  getApp,
-  getConnectionByProvider,
-  getConnectionByProviderAndAccount,
-  getProvider,
-} from "./oauth-store.js";
+import { getActiveConnection, getProvider } from "./oauth-store.js";
+import { PlatformOAuthConnection } from "./platform-connection.js";
+
+export interface ResolveOAuthConnectionOptions {
+  /** OAuth app client ID — narrows to a specific app when multiple BYO apps
+   *  exist for the same provider. */
+  clientId?: string;
+  /** Account identifier (e.g. email, username) — disambiguates when multiple
+   *  accounts are connected for the same provider. Best-effort: not guaranteed
+   *  to be present on all connections. */
+  account?: string;
+}
 
 /**
- * Resolve an OAuthConnection for a given credential service.
+ * Resolve an OAuthConnection for a given provider.
+ *
+ * Managed providers (where the service config `mode` is `"managed"`) are
+ * routed through the platform proxy with no local state required.
  *
- * When `accountInfo` is provided, resolves the connection for that specific
- * account (e.g. "user@gmail.com"). Otherwise falls back to the most recent
- * active connection.
+ * BYO providers resolve from the local SQLite oauth-store and require an
+ * active connection row and a stored access token.
  *
- * Reads exclusively from the SQLite oauth-store. Throws if no connection
- * exists (authorization required).
+ * @param providerKey - Provider identifier (e.g. "integration:google").
+ *   Maps to the `provider_key` primary key in the `oauth_providers` table.
+ * @param options.clientId - Optional OAuth app client ID. When multiple BYO
+ *   apps exist for the same provider, narrows the connection lookup to the
+ *   app matching this client ID. Ignored for managed providers.
+ * @param options.account - Optional account identifier to disambiguate
+ *   multi-account connections.
  */
 export async function resolveOAuthConnection(
-  credentialService: string,
-  accountInfo?: string,
+  providerKey: string,
+  options?: ResolveOAuthConnectionOptions,
 ): Promise<OAuthConnection> {
-  const conn = accountInfo
-    ? getConnectionByProviderAndAccount(credentialService, accountInfo)
-    : getConnectionByProvider(credentialService);
+  const { clientId, account } = options ?? {};
+  const provider = getProvider(providerKey);
+  const managedKey = provider?.managedServiceConfigKey;
+
+  if (managedKey && managedKey in ServicesSchema.shape) {
+    const services: Services = getConfig().services;
+    if (services[managedKey as keyof Services].mode === "managed") {
+      const ctx = await resolveManagedProxyContext();
+      const assistantId = getPlatformAssistantId();
+      return new PlatformOAuthConnection({
+        id: providerKey,
+        providerKey,
+        externalId: providerKey,
+        accountInfo: account ?? null,
+        assistantId,
+        platformBaseUrl: ctx.platformBaseUrl,
+        apiKey: ctx.assistantApiKey,
+      });
+    }
+  }
+
+  // BYO path — requires a local connection row, access token, and base URL.
+  const conn = getActiveConnection(providerKey, { clientId, account });
   if (!conn) {
+    const filters = [
+      account && `account "${account}"`,
+      clientId && `client ID "${clientId}"`,
+    ].filter(Boolean);
+    const qualifier = filters.length
+      ? ` matching ${filters.join(" and ")}`
+      : "";
     throw new Error(
-      `No credential found for "${credentialService}". Authorization required.`,
+      `No active OAuth connection found for "${providerKey}"${qualifier}. Connect the service first with oauth2_connect.`,
     );
   }
 
   const accessToken = await getSecureKeyAsync(
     `oauth_connection/${conn.id}/access_token`,
   );
-
   if (!accessToken) {
     throw new Error(
-      `No access token found for "${credentialService}". Authorization required.`,
+      `OAuth connection for "${providerKey}" exists but has no access token. Re-authorize with oauth2_connect.`,
     );
   }
 
-  // Look up the provider by credentialService first; fall back to the
-  // connection's app's canonical providerKey so custom credential_service
-  // overrides (e.g. "integration:github-work") still resolve to the well-known
-  // provider's base URL. We traverse conn -> oauthApp -> providerKey because
-  // conn.providerKey equals credentialService (getConnectionByProvider queries
-  // WHERE providerKey = credentialService), whereas the app's providerKey is a
-  // foreign key to the oauthProviders table.
-  const provider =
-    getProvider(credentialService) ??
-    getProvider(getApp(conn.oauthAppId)?.providerKey ?? "");
   const baseUrl = provider?.baseUrl;
-
   if (!baseUrl) {
-    throw new Error(`No base URL configured for "${credentialService}".`);
+    throw new Error(
+      `OAuth provider "${providerKey}" has no base URL configured. Check provider setup.`,
+    );
   }
 
-  const grantedScopes: string[] = conn.grantedScopes
-    ? JSON.parse(conn.grantedScopes)
-    : [];
-
   return new BYOOAuthConnection({
     id: conn.id,
     providerKey: conn.providerKey,
     baseUrl,
     accountInfo: conn.accountInfo,
-    grantedScopes,
-    credentialService,
   });
 }

package/src/oauth/connection.ts

@@ -34,5 +34,4 @@ export interface OAuthConnection {
   readonly id: string;
   readonly providerKey: string;
   readonly accountInfo: string | null;
-  readonly grantedScopes: string[];
 }

package/src/oauth/oauth-store.ts

@@ -45,7 +45,9 @@ export type OAuthConnectionRow = typeof oauthConnections.$inferSelect;
  * Seed well-known provider profiles into the database. Uses INSERT … ON
  * CONFLICT DO UPDATE so that implementation fields (authUrl, tokenUrl,
  * tokenEndpointAuthMethod, userinfoUrl, extraParams, callbackTransport,
- * pingUrl) propagate to existing installations on every startup, while
+ * pingUrl, managedServiceConfigKey) and display metadata (displayName,
+ * description, dashboardUrl, clientIdPlaceholder, requiresClientSecret)
+ * propagate to existing installations on every startup, while
  * user-customizable fields (defaultScopes, scopePolicy, baseUrl) are
  * only written on the initial insert.
  */
@@ -62,6 +64,12 @@ export function seedProviders(
     scopePolicy: Record<string, unknown>;
     extraParams?: Record<string, string>;
     callbackTransport?: string;
+    managedServiceConfigKey?: string;
+    displayName?: string;
+    description?: string;
+    dashboardUrl?: string | null;
+    clientIdPlaceholder?: string | null;
+    requiresClientSecret?: boolean;
   }>,
 ): void {
   const db = getDb();
@@ -77,6 +85,12 @@ export function seedProviders(
     const scopePolicy = JSON.stringify(p.scopePolicy);
     const extraParams = p.extraParams ? JSON.stringify(p.extraParams) : null;
     const callbackTransport = p.callbackTransport ?? null;
+    const managedServiceConfigKey = p.managedServiceConfigKey ?? null;
+    const displayName = p.displayName ?? null;
+    const description = p.description ?? null;
+    const dashboardUrl = p.dashboardUrl ?? null;
+    const clientIdPlaceholder = p.clientIdPlaceholder ?? null;
+    const requiresClientSecret = p.requiresClientSecret !== false ? 1 : 0;
 
     db.insert(oauthProviders)
       .values({
@@ -91,6 +105,12 @@ export function seedProviders(
         extraParams,
         callbackTransport,
         pingUrl,
+        managedServiceConfigKey,
+        displayName,
+        description,
+        dashboardUrl,
+        clientIdPlaceholder,
+        requiresClientSecret,
         createdAt: now,
         updatedAt: now,
       })
@@ -104,6 +124,12 @@ export function seedProviders(
           extraParams,
           callbackTransport,
           pingUrl,
+          managedServiceConfigKey,
+          displayName,
+          description,
+          dashboardUrl,
+          clientIdPlaceholder,
+          requiresClientSecret,
           updatedAt: now,
         },
       })
@@ -143,6 +169,12 @@ export function registerProvider(params: {
   scopePolicy: Record<string, unknown>;
   extraParams?: Record<string, string>;
   callbackTransport?: string;
+  managedServiceConfigKey?: string;
+  displayName?: string;
+  description?: string;
+  dashboardUrl?: string;
+  clientIdPlaceholder?: string;
+  requiresClientSecret?: number;
 }): OAuthProviderRow {
   const db = getDb();
   const now = Date.now();
@@ -164,6 +196,12 @@ export function registerProvider(params: {
     extraParams: params.extraParams ? JSON.stringify(params.extraParams) : null,
     callbackTransport: params.callbackTransport ?? null,
     pingUrl: params.pingUrl ?? null,
+    managedServiceConfigKey: params.managedServiceConfigKey ?? null,
+    displayName: params.displayName ?? null,
+    description: params.description ?? null,
+    dashboardUrl: params.dashboardUrl ?? null,
+    clientIdPlaceholder: params.clientIdPlaceholder ?? null,
+    requiresClientSecret: params.requiresClientSecret ?? 1,
     createdAt: now,
     updatedAt: now,
   };
@@ -232,6 +270,14 @@ export async function upsertApp(
       if (!stored) {
         throw new Error("Failed to store client_secret in secure storage");
       }
+      // Bump updatedAt so the rollback guard in the new-row insertion path
+      // can detect that a concurrent caller has claimed this row. Without
+      // this, a concurrent inserter's rollback DELETE would still match on
+      // the original updatedAt and delete the row we just validated.
+      db.update(oauthApps)
+        .set({ updatedAt: Date.now() })
+        .where(eq(oauthApps.id, existingRow.id))
+        .run();
     }
     if (clientSecretCredentialPath) {
       db.update(oauthApps)
@@ -272,7 +318,14 @@ export async function upsertApp(
     if (!stored) {
       // Roll back the just-inserted row to avoid an orphaned app pointing
       // at a non-existent client_secret in secure storage.
-      db.delete(oauthApps).where(eq(oauthApps.id, id)).run();
+      //
+      // Guard: only delete if updatedAt still matches our insertion timestamp.
+      // A concurrent upsertApp call may have observed this row, successfully
+      // stored the secret, and updated the row — deleting it would orphan that
+      // caller's valid reference.
+      db.delete(oauthApps)
+        .where(and(eq(oauthApps.id, id), eq(oauthApps.updatedAt, now)))
+        .run();
       throw new Error("Failed to store client_secret in secure storage");
     }
   }
@@ -286,6 +339,15 @@ export function getApp(id: string): OAuthAppRow | undefined {
   return db.select().from(oauthApps).where(eq(oauthApps.id, id)).get();
 }
 
+/** Read an app client_secret from secure storage. */
+export async function getAppClientSecret(
+  appOrId: OAuthAppRow | string,
+): Promise<string | undefined> {
+  const app = typeof appOrId === "string" ? getApp(appOrId) : appOrId;
+  if (!app) return undefined;
+  return getSecureKeyAsync(app.clientSecretCredentialPath);
+}
+
 /** Look up an app by (provider_key, client_id). */
 export function getAppByProviderAndClientId(
   providerKey: string,
@@ -407,71 +469,59 @@ export function getConnection(id: string): OAuthConnectionRow | undefined {
 
 /**
  * Get the most recent active connection for a provider.
- * When `clientId` is provided, only connections linked to the matching app are considered.
- * Returns undefined if no active connection exists.
+ *
+ * Optional filters narrow the result:
+ * - `account` — match a specific account identifier (e.g. email).
+ * - `clientId` — restrict to connections linked to a specific OAuth app.
+ *
+ * Returns `undefined` when no matching active connection exists.
  */
-export function getConnectionByProvider(
+export function getActiveConnection(
   providerKey: string,
-  clientId?: string,
+  options?: { clientId?: string; account?: string },
 ): OAuthConnectionRow | undefined {
+  const { clientId, account } = options ?? {};
   const db = getDb();
 
+  const conditions = [
+    eq(oauthConnections.providerKey, providerKey),
+    eq(oauthConnections.status, "active"),
+  ];
+
+  if (account) {
+    conditions.push(eq(oauthConnections.accountInfo, account));
+  }
+
   if (clientId) {
     const app = getAppByProviderAndClientId(providerKey, clientId);
     if (!app) return undefined;
-    return db
-      .select()
-      .from(oauthConnections)
-      .where(
-        and(
-          eq(oauthConnections.providerKey, providerKey),
-          eq(oauthConnections.oauthAppId, app.id),
-          eq(oauthConnections.status, "active"),
-        ),
-      )
-      .orderBy(desc(oauthConnections.createdAt), sql`rowid DESC`)
-      .limit(1)
-      .get();
+    conditions.push(eq(oauthConnections.oauthAppId, app.id));
   }
 
   return db
     .select()
     .from(oauthConnections)
-    .where(
-      and(
-        eq(oauthConnections.providerKey, providerKey),
-        eq(oauthConnections.status, "active"),
-      ),
-    )
+    .where(and(...conditions))
     .orderBy(desc(oauthConnections.createdAt), sql`rowid DESC`)
     .limit(1)
     .get();
 }
 
-/**
- * Get the active connection for a provider matching a specific account.
- * Falls back to getConnectionByProvider when accountInfo is undefined.
- */
+/** @deprecated Use {@link getActiveConnection} instead. */
+export function getConnectionByProvider(
+  providerKey: string,
+  clientId?: string,
+): OAuthConnectionRow | undefined {
+  return getActiveConnection(providerKey, { clientId });
+}
+
+/** @deprecated Use {@link getActiveConnection} instead. */
 export function getConnectionByProviderAndAccount(
   providerKey: string,
   accountInfo?: string,
+  clientId?: string,
 ): OAuthConnectionRow | undefined {
-  if (!accountInfo) return getConnectionByProvider(providerKey);
-
-  const db = getDb();
-  return db
-    .select()
-    .from(oauthConnections)
-    .where(
-      and(
-        eq(oauthConnections.providerKey, providerKey),
-        eq(oauthConnections.accountInfo, accountInfo),
-        eq(oauthConnections.status, "active"),
-      ),
-    )
-    .orderBy(desc(oauthConnections.createdAt), sql`rowid DESC`)
-    .limit(1)
-    .get();
+  return getActiveConnection(providerKey, { clientId, account: accountInfo });
 }
 
 /**
@@ -505,7 +555,7 @@ export function listActiveConnectionsByProvider(
 export async function isProviderConnected(
   providerKey: string,
 ): Promise<boolean> {
-  const conn = getConnectionByProvider(providerKey);
+  const conn = getActiveConnection(providerKey);
   if (!conn || conn.status !== "active") return false;
   return (
     (await getSecureKeyAsync(oauthConnectionAccessTokenPath(conn.id))) !==
@@ -623,7 +673,7 @@ export async function disconnectOAuthProvider(
 ): Promise<"disconnected" | "not-found" | "error"> {
   const conn = connectionId
     ? getConnection(connectionId)
-    : getConnectionByProvider(providerKey, clientId);
+    : getActiveConnection(providerKey, { clientId });
   if (!conn) return "not-found";
 
   // Wrap the assistant's secure-key functions into the SecureKeyBackend

package/src/oauth/platform-connection.test.ts

@@ -12,7 +12,6 @@ const DEFAULT_OPTIONS = {
   providerKey: "integration:google",
   externalId: "ext-123",
   accountInfo: "user@example.com",
-  grantedScopes: ["https://www.googleapis.com/auth/gmail.readonly"],
   assistantId: "asst-abc",
   platformBaseUrl: "https://platform.example.com",
   apiKey: "test-api-key",

package/src/oauth/platform-connection.ts

@@ -24,7 +24,6 @@ export interface PlatformOAuthConnectionOptions {
   providerKey: string;
   externalId: string;
   accountInfo: string | null;
-  grantedScopes: string[];
   assistantId: string;
   platformBaseUrl: string;
   apiKey: string;
@@ -35,18 +34,27 @@ export class PlatformOAuthConnection implements OAuthConnection {
   readonly providerKey: string;
   readonly externalId: string;
   readonly accountInfo: string | null;
-  readonly grantedScopes: string[];
 
   private readonly assistantId: string;
   private readonly platformBaseUrl: string;
   private readonly apiKey: string;
 
   constructor(options: PlatformOAuthConnectionOptions) {
+    const missing: string[] = [];
+    if (!options.platformBaseUrl) missing.push("platform base URL");
+    if (!options.apiKey) missing.push("assistant API key");
+    if (!options.assistantId) missing.push("assistant ID");
+    if (missing.length > 0) {
+      throw new BackendError(
+        `Platform-managed connection for "${options.providerKey}" cannot be created: missing ${missing.join(", ")}. ` +
+          `Log in to the Vellum platform or switch to using your own OAuth app.`,
+      );
+    }
+
     this.id = options.id;
     this.providerKey = options.providerKey;
     this.externalId = options.externalId;
     this.accountInfo = options.accountInfo;
-    this.grantedScopes = options.grantedScopes;
     this.assistantId = options.assistantId;
     this.platformBaseUrl = options.platformBaseUrl.replace(/\/+$/, "");
     this.apiKey = options.apiKey;