@vellumai/assistant 0.6.3 → 0.6.5

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

@@ -77,6 +77,7 @@ const mockConnections = new Map<
     expiresAt: number | null;
     grantedScopes?: string;
     accountInfo?: string | null;
+    status?: string;
   }
 >();
 const mockApps = new Map<
@@ -92,14 +93,26 @@ const mockProviders = new Map<
   string,
   {
     key: string;
+    provider: string;
     tokenExchangeUrl: string;
     tokenEndpointAuthMethod?: string;
+    tokenExchangeBodyFormat?: string;
     baseUrl?: string;
+    managedServiceConfigKey?: string | null;
   }
 >();
 
 mock.module("./oauth-store.js", () => ({
   getConnectionByProvider: (service: string) => mockConnections.get(service),
+  getActiveConnection: (
+    service: string,
+    opts?: { clientId?: string; account?: string },
+  ) => {
+    const conn = mockConnections.get(service);
+    if (!conn) return undefined;
+    if (opts?.account && conn.accountInfo !== opts.account) return undefined;
+    return conn;
+  },
   getConnection: (id: string) => {
     for (const conn of mockConnections.values()) {
       if (conn.id === id) return conn;
@@ -192,12 +205,15 @@ async function setupCredential(
   const connId = `conn-${service}`;
   mockProviders.set(service, {
     key: service,
+    provider: service,
     tokenExchangeUrl: "https://oauth2.googleapis.com/token",
+    tokenExchangeBodyFormat: "form",
     // Only well-known providers (gmail) have a baseUrl; custom services don't
     baseUrl:
       service === "google"
         ? "https://gmail.googleapis.com/gmail/v1/users/me"
         : undefined,
+    managedServiceConfigKey: null,
   });
   mockApps.set(appId, {
     id: appId,
@@ -212,6 +228,7 @@ async function setupCredential(
     expiresAt: opts?.expiresAt ?? Date.now() + 3600 * 1000,
     grantedScopes: JSON.stringify(opts?.grantedScopes ?? ["read", "write"]),
     accountInfo: null,
+    status: "active",
   });
   // Store access token in oauth-store key format
   await setSecureKeyAsync(
@@ -509,7 +526,7 @@ describe("resolveOAuthConnection", () => {
   test("throws when no base URL configured", async () => {
     await setupCredential("custom-service");
     await expect(resolveOAuthConnection("custom-service")).rejects.toThrow(
-      /No base URL configured for "custom-service"/,
+      /OAuth provider "custom-service" has no base URL configured/,
     );
   });
 });

package/src/oauth/byo-connection.ts

@@ -82,7 +82,9 @@ export class BYOOAuthConnection implements OAuthConnection {
           method: req.method,
           headers,
           body: req.body ? JSON.stringify(req.body) : undefined,
-          signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS),
+          signal: req.signal
+            ? AbortSignal.any([req.signal, AbortSignal.timeout(REQUEST_TIMEOUT_MS)])
+            : AbortSignal.timeout(REQUEST_TIMEOUT_MS),
         });
 
         if (resp.status === 401) {

package/src/oauth/connect-orchestrator.ts

@@ -209,6 +209,8 @@ export async function orchestrateOAuthConnect(
     authorizeParams,
     userinfoUrl,
     tokenEndpointAuthMethod,
+    tokenExchangeBodyFormat:
+      (providerRow.tokenExchangeBodyFormat as "form" | "json") ?? "form",
   };
 
   const storageParams = {

package/src/oauth/connection-resolver.ts

@@ -1,5 +1,9 @@
 import { getConfig } from "../config/loader.js";
-import { type Services, ServicesSchema } from "../config/schemas/services.js";
+import {
+  getServiceMode,
+  type Services,
+  ServicesSchema,
+} from "../config/schemas/services.js";
 import { VellumPlatformClient } from "../platform/client.js";
 import { getSecureKeyAsync } from "../security/secure-keys.js";
 import { getLogger } from "../util/logger.js";
@@ -47,7 +51,7 @@ export async function resolveOAuthConnection(
 
   if (managedKey && managedKey in ServicesSchema.shape) {
     const services: Services = getConfig().services;
-    if (services[managedKey as keyof Services].mode === "managed") {
+    if (getServiceMode(services, managedKey as keyof Services) === "managed") {
       const client = await VellumPlatformClient.create();
       if (!client || !client.platformAssistantId) {
         const detail = !client

package/src/oauth/connection.ts

@@ -11,6 +11,8 @@ export interface OAuthConnectionRequest {
    * use the same credential but different base URLs).
    */
   baseUrl?: string;
+  /** Optional abort signal to cancel the request. */
+  signal?: AbortSignal;
 }
 
 export interface OAuthConnectionResponse {

package/src/oauth/oauth-store.ts

@@ -72,6 +72,7 @@ export function seedProviders(
     tokenExchangeUrl: string;
     refreshUrl?: string;
     tokenEndpointAuthMethod?: string;
+    tokenExchangeBodyFormat?: string;
     userinfoUrl?: string;
     pingUrl?: string;
     pingMethod?: string;
@@ -122,6 +123,7 @@ export function seedProviders(
     // token endpoint auth method.
     const tokenEndpointAuthMethod =
       p.tokenEndpointAuthMethod || "client_secret_post";
+    const tokenExchangeBodyFormat = p.tokenExchangeBodyFormat || "form";
     const userinfoUrl = p.userinfoUrl ?? null;
     const pingUrl = p.pingUrl ?? null;
     const pingMethod = p.pingMethod ?? null;
@@ -176,6 +178,7 @@ export function seedProviders(
         tokenExchangeUrl,
         refreshUrl,
         tokenEndpointAuthMethod,
+        tokenExchangeBodyFormat,
         userinfoUrl,
         baseUrl,
         defaultScopes,
@@ -217,8 +220,10 @@ export function seedProviders(
           tokenExchangeUrl,
           refreshUrl,
           tokenEndpointAuthMethod,
+          tokenExchangeBodyFormat,
           userinfoUrl,
           baseUrl: sql`COALESCE(${oauthProviders.baseUrl}, ${baseUrl})`,
+          defaultScopes,
           scopeSeparator,
           authorizeParams,
           pingUrl,
@@ -279,6 +284,7 @@ export function registerProvider(params: {
   tokenExchangeUrl: string;
   refreshUrl?: string;
   tokenEndpointAuthMethod?: string;
+  tokenExchangeBodyFormat?: string;
   userinfoUrl?: string;
   pingUrl?: string;
   pingMethod?: string;
@@ -331,6 +337,7 @@ export function registerProvider(params: {
     refreshUrl: params.refreshUrl ?? null,
     tokenEndpointAuthMethod:
       params.tokenEndpointAuthMethod || "client_secret_post",
+    tokenExchangeBodyFormat: params.tokenExchangeBodyFormat || "form",
     userinfoUrl: params.userinfoUrl ?? null,
     baseUrl: params.baseUrl ?? null,
     defaultScopes: JSON.stringify(params.defaultScopes),
@@ -402,6 +409,7 @@ export function updateProvider(
     tokenExchangeUrl: string;
     refreshUrl: string;
     tokenEndpointAuthMethod: string;
+    tokenExchangeBodyFormat: string;
     userinfoUrl: string;
     pingUrl: string;
     pingMethod: string;
@@ -452,6 +460,8 @@ export function updateProvider(
   if (params.tokenEndpointAuthMethod !== undefined)
     set.tokenEndpointAuthMethod =
       params.tokenEndpointAuthMethod || "client_secret_post";
+  if (params.tokenExchangeBodyFormat !== undefined)
+    set.tokenExchangeBodyFormat = params.tokenExchangeBodyFormat || "form";
   if (params.userinfoUrl !== undefined) set.userinfoUrl = params.userinfoUrl;
   if (params.pingUrl !== undefined) set.pingUrl = params.pingUrl;
   if (params.pingMethod !== undefined) set.pingMethod = params.pingMethod;

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

@@ -1,9 +1,18 @@
 import { describe, expect, mock, test } from "bun:test";
 
+import * as actualRetry from "../util/retry.js";
+
+// Stub out sleep so retry tests don't wait for real delays.
+mock.module("../util/retry.js", () => ({
+  ...actualRetry,
+  sleep: () => Promise.resolve(),
+}));
+
 import type { VellumPlatformClient } from "../platform/client.js";
 import { BackendError, VellumError } from "../util/errors.js";
 import {
   CredentialRequiredError,
+  InsufficientBalanceError,
   PlatformOAuthConnection,
   ProviderUnreachableError,
 } from "./platform-connection.js";
@@ -191,6 +200,52 @@ describe("PlatformOAuthConnection", () => {
     const provErr = new ProviderUnreachableError();
     expect(provErr).toBeInstanceOf(BackendError);
     expect(provErr).toBeInstanceOf(VellumError);
+
+    const balErr = new InsufficientBalanceError();
+    expect(balErr).toBeInstanceOf(BackendError);
+    expect(balErr).toBeInstanceOf(VellumError);
+  });
+
+  test("402 response throws InsufficientBalanceError", async () => {
+    const client = makeMockClient(
+      mock(
+        async () => new Response("", { status: 402 }),
+      ) as unknown as typeof globalThis.fetch,
+    );
+
+    const conn = new PlatformOAuthConnection({ ...DEFAULT_OPTIONS, client });
+    await expect(
+      conn.request({ method: "GET", path: "/test" }),
+    ).rejects.toThrow(InsufficientBalanceError);
+  });
+
+  test("402 response includes actionable billing message", async () => {
+    const client = makeMockClient(
+      mock(
+        async () => new Response("", { status: 402 }),
+      ) as unknown as typeof globalThis.fetch,
+    );
+
+    const conn = new PlatformOAuthConnection({ ...DEFAULT_OPTIONS, client });
+    await expect(
+      conn.request({ method: "GET", path: "/test" }),
+    ).rejects.toThrow(/add funds/i);
+  });
+
+  test("does not retry on 402", async () => {
+    let callCount = 0;
+    const client = makeMockClient(
+      mock(async () => {
+        callCount++;
+        return new Response("", { status: 402 });
+      }) as unknown as typeof globalThis.fetch,
+    );
+
+    const conn = new PlatformOAuthConnection({ ...DEFAULT_OPTIONS, client });
+    await expect(
+      conn.request({ method: "GET", path: "/test" }),
+    ).rejects.toThrow(InsufficientBalanceError);
+    expect(callCount).toBe(1);
   });
 
   test("424 response throws CredentialRequiredError", async () => {
@@ -226,6 +281,96 @@ describe("PlatformOAuthConnection", () => {
     );
   });
 
+  test("retries on 429 and succeeds", async () => {
+    let callCount = 0;
+    const client = makeMockClient(
+      mock(async () => {
+        callCount++;
+        if (callCount <= 2) {
+          return new Response("", { status: 429 });
+        }
+        return new Response(
+          JSON.stringify({ status: 200, headers: {}, body: { ok: true } }),
+          { status: 200 },
+        );
+      }) as unknown as typeof globalThis.fetch,
+    );
+
+    const conn = new PlatformOAuthConnection({ ...DEFAULT_OPTIONS, client });
+    const result = await conn.request({ method: "GET", path: "/test" });
+
+    expect(result.status).toBe(200);
+    expect(result.body).toEqual({ ok: true });
+    expect(callCount).toBe(3);
+  });
+
+  test("throws after exhausting retries on 429", async () => {
+    const client = makeMockClient(
+      mock(
+        async () => new Response("", { status: 429 }),
+      ) as unknown as typeof globalThis.fetch,
+    );
+
+    const conn = new PlatformOAuthConnection({ ...DEFAULT_OPTIONS, client });
+    await expect(
+      conn.request({ method: "GET", path: "/test" }),
+    ).rejects.toThrow("Platform proxy returned unexpected status 429");
+  });
+
+  test("retries on 500 and succeeds", async () => {
+    let callCount = 0;
+    const client = makeMockClient(
+      mock(async () => {
+        callCount++;
+        if (callCount === 1) {
+          return new Response("", { status: 500 });
+        }
+        return new Response(
+          JSON.stringify({ status: 200, headers: {}, body: null }),
+          { status: 200 },
+        );
+      }) as unknown as typeof globalThis.fetch,
+    );
+
+    const conn = new PlatformOAuthConnection({ ...DEFAULT_OPTIONS, client });
+    const result = await conn.request({ method: "GET", path: "/test" });
+
+    expect(result.status).toBe(200);
+    expect(callCount).toBe(2);
+  });
+
+  test("does not retry on 424", async () => {
+    let callCount = 0;
+    const client = makeMockClient(
+      mock(async () => {
+        callCount++;
+        return new Response("", { status: 424 });
+      }) as unknown as typeof globalThis.fetch,
+    );
+
+    const conn = new PlatformOAuthConnection({ ...DEFAULT_OPTIONS, client });
+    await expect(
+      conn.request({ method: "GET", path: "/test" }),
+    ).rejects.toThrow(CredentialRequiredError);
+    expect(callCount).toBe(1);
+  });
+
+  test("does not retry on 403", async () => {
+    let callCount = 0;
+    const client = makeMockClient(
+      mock(async () => {
+        callCount++;
+        return new Response("", { status: 403 });
+      }) as unknown as typeof globalThis.fetch,
+    );
+
+    const conn = new PlatformOAuthConnection({ ...DEFAULT_OPTIONS, client });
+    await expect(
+      conn.request({ method: "GET", path: "/test" }),
+    ).rejects.toThrow("Platform proxy returned unexpected status 403");
+    expect(callCount).toBe(1);
+  });
+
   test("uses connectionId in proxy URL regardless of provider format", async () => {
     const client = makeMockClient(
       mock(async (url: string | URL | Request) => {

package/src/oauth/platform-connection.ts

@@ -1,11 +1,14 @@
 import type { VellumPlatformClient } from "../platform/client.js";
 import { BackendError } from "../util/errors.js";
+import { getHttpRetryDelay, isRetryableStatus, sleep } from "../util/retry.js";
 import type {
   OAuthConnection,
   OAuthConnectionRequest,
   OAuthConnectionResponse,
 } from "./connection.js";
 
+const MAX_RETRIES = 3;
+
 export class CredentialRequiredError extends BackendError {
   constructor(message = "Connection not set up on platform") {
     super(message);
@@ -20,6 +23,16 @@ export class ProviderUnreachableError extends BackendError {
   }
 }
 
+export class InsufficientBalanceError extends BackendError {
+  constructor(
+    message = "Your Vellum account balance is too low to use this managed OAuth connection. " +
+      "You can add funds or switch to using your own OAuth app.",
+  ) {
+    super(message);
+    this.name = "InsufficientBalanceError";
+  }
+}
+
 export interface PlatformOAuthConnectionOptions {
   id: string;
   provider: string;
@@ -76,39 +89,57 @@ export class PlatformOAuthConnection implements OAuthConnection {
       },
     };
 
-    const response = await this.client.fetch(proxyPath, {
-      method: "POST",
-      headers: {
-        "Content-Type": "application/json",
-      },
-      body: JSON.stringify(body),
-    });
-
-    if (response.status === 424) {
-      throw new CredentialRequiredError();
-    }
-
-    if (response.status === 502) {
-      throw new ProviderUnreachableError();
-    }
-
-    if (!response.ok) {
-      throw new BackendError(
-        `Platform proxy returned unexpected status ${response.status}`,
-      );
+    for (let attempt = 0; attempt <= MAX_RETRIES; attempt++) {
+      const response = await this.client.fetch(proxyPath, {
+        method: "POST",
+        headers: {
+          "Content-Type": "application/json",
+        },
+        body: JSON.stringify(body),
+        signal: req.signal,
+      });
+
+      if (response.status === 402) {
+        throw new InsufficientBalanceError();
+      }
+
+      if (response.status === 424) {
+        throw new CredentialRequiredError();
+      }
+
+      if (response.status === 502) {
+        throw new ProviderUnreachableError();
+      }
+
+      if (
+        !response.ok &&
+        isRetryableStatus(response.status) &&
+        attempt < MAX_RETRIES
+      ) {
+        await sleep(getHttpRetryDelay(response, attempt));
+        continue;
+      }
+
+      if (!response.ok) {
+        throw new BackendError(
+          `Platform proxy returned unexpected status ${response.status}`,
+        );
+      }
+
+      const json = (await response.json()) as {
+        status: number;
+        headers: Record<string, string>;
+        body: unknown;
+      };
+
+      return {
+        status: json.status,
+        headers: json.headers,
+        body: json.body,
+      };
     }
 
-    const json = (await response.json()) as {
-      status: number;
-      headers: Record<string, string>;
-      body: unknown;
-    };
-
-    return {
-      status: json.status,
-      headers: json.headers,
-      body: json.body,
-    };
+    throw new BackendError("Platform proxy request failed after retries");
   }
 
   async withToken<T>(_fn: (token: string) => Promise<T>): Promise<T> {

package/src/oauth/seed-providers.ts

@@ -16,8 +16,11 @@ import { seedProviders } from "./oauth-store.js";
  * description, dashboardUrl, clientIdPlaceholder, requiresClientSecret,
  * logoUrl)
  * are overwritten on subsequent startups — user-customizable
- * fields (defaultScopes, scopePolicy) are only
+ * fields (scopePolicy) are only
  * written on initial insert and preserved across restarts.
+ * defaultScopes are also overwritten on subsequent startups so that
+ * upstream scope additions (e.g. new Gmail API scopes) propagate to
+ * existing installations.
  */
 export const PROVIDER_SEED_DATA: Record<
   string,
@@ -27,6 +30,7 @@ export const PROVIDER_SEED_DATA: Record<
     tokenExchangeUrl: string;
     refreshUrl?: string;
     tokenEndpointAuthMethod?: string;
+    tokenExchangeBodyFormat?: string;
     userinfoUrl?: string;
     pingUrl?: string;
     pingMethod?: string;
@@ -85,6 +89,7 @@ export const PROVIDER_SEED_DATA: Record<
       "https://www.googleapis.com/auth/gmail.readonly",
       "https://www.googleapis.com/auth/gmail.modify",
       "https://www.googleapis.com/auth/gmail.send",
+      "https://www.googleapis.com/auth/gmail.settings.basic",
       "https://www.googleapis.com/auth/calendar.readonly",
       "https://www.googleapis.com/auth/calendar.events",
       "https://www.googleapis.com/auth/userinfo.email",
@@ -140,6 +145,7 @@ export const PROVIDER_SEED_DATA: Record<
     clientIdPlaceholder: null,
     logoUrl: "https://cdn.simpleicons.org/slack",
     defaultScopes: [
+      "channels:join",
       "channels:read",
       "channels:history",
       "groups:read",
@@ -199,6 +205,8 @@ export const PROVIDER_SEED_DATA: Record<
     },
     authorizeParams: { owner: "user" },
     tokenEndpointAuthMethod: "client_secret_basic",
+    tokenExchangeBodyFormat: "json",
+    managedServiceConfigKey: "notion-oauth",
     loopbackPort: 17323,
     injectionTemplates: [
       {
@@ -238,6 +246,7 @@ export const PROVIDER_SEED_DATA: Record<
     },
     tokenEndpointAuthMethod: "client_secret_basic",
     loopbackPort: 17335,
+    managedServiceConfigKey: "twitter-oauth",
     injectionTemplates: [
       {
         hostPattern: "api.x.com",