@vellumai/assistant 0.4.45 → 0.4.48

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

@@ -0,0 +1,436 @@
+import { randomBytes } from "node:crypto";
+import { mkdirSync, rmSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import {
+  afterAll,
+  afterEach,
+  beforeAll,
+  beforeEach,
+  describe,
+  expect,
+  mock,
+  test,
+} from "bun:test";
+
+// ---------------------------------------------------------------------------
+// Mock logger
+// ---------------------------------------------------------------------------
+
+mock.module("../util/logger.js", () => ({
+  getLogger: () =>
+    new Proxy({} as Record<string, unknown>, {
+      get: () => () => {},
+    }),
+}));
+
+// ---------------------------------------------------------------------------
+// Use encrypted backend with a temp store path
+// ---------------------------------------------------------------------------
+
+import { _setStorePath } from "../security/encrypted-store.js";
+import { _resetBackend } from "../security/secure-keys.js";
+
+const TEST_DIR = join(
+  tmpdir(),
+  `vellum-byo-conn-test-${randomBytes(4).toString("hex")}`,
+);
+const STORE_PATH = join(TEST_DIR, "keys.enc");
+
+// ---------------------------------------------------------------------------
+// Mock OAuth2 token refresh
+// ---------------------------------------------------------------------------
+
+let mockRefreshOAuth2Token: ReturnType<
+  typeof mock<
+    () => Promise<{
+      accessToken: string;
+      expiresIn: number;
+      refreshToken?: string;
+    }>
+  >
+>;
+
+mock.module("../security/oauth2.js", () => {
+  mockRefreshOAuth2Token = mock(() =>
+    Promise.resolve({
+      accessToken: "refreshed-access-token",
+      expiresIn: 3600,
+    }),
+  );
+  return {
+    refreshOAuth2Token: mockRefreshOAuth2Token,
+  };
+});
+
+// ---------------------------------------------------------------------------
+// Imports (after mocks)
+// ---------------------------------------------------------------------------
+
+import {
+  _resetMigrationFlag,
+  credentialKey,
+} from "../security/credential-key.js";
+import { setSecureKey } from "../security/secure-keys.js";
+import {
+  _resetInflightRefreshes,
+  _resetRefreshBreakers,
+} from "../security/token-manager.js";
+import {
+  _setMetadataPath,
+  upsertCredentialMetadata,
+} from "../tools/credentials/metadata-store.js";
+import { BYOOAuthConnection } from "./byo-connection.js";
+import { resolveOAuthConnection } from "./connection-resolver.js";
+
+// ---------------------------------------------------------------------------
+// Mock fetch
+// ---------------------------------------------------------------------------
+
+const originalFetch = globalThis.fetch;
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+let mockFetch: ReturnType<typeof mock<any>>;
+
+// ---------------------------------------------------------------------------
+// Setup / teardown
+// ---------------------------------------------------------------------------
+
+beforeAll(() => {
+  mkdirSync(TEST_DIR, { recursive: true });
+});
+
+afterAll(() => {
+  rmSync(TEST_DIR, { recursive: true, force: true });
+  globalThis.fetch = originalFetch;
+});
+
+beforeEach(() => {
+  _setStorePath(STORE_PATH);
+  _setMetadataPath(join(TEST_DIR, "metadata.json"));
+  _resetBackend();
+  _resetRefreshBreakers();
+  _resetInflightRefreshes();
+  _resetMigrationFlag();
+
+  // Default mock fetch returning 200 JSON
+  mockFetch = mock(() =>
+    Promise.resolve(
+      new Response(JSON.stringify({ ok: true }), {
+        status: 200,
+        headers: { "content-type": "application/json" },
+      }),
+    ),
+  );
+  globalThis.fetch = mockFetch as unknown as typeof fetch;
+});
+
+afterEach(() => {
+  globalThis.fetch = originalFetch;
+  // Clean up store for next test
+  try {
+    rmSync(STORE_PATH, { force: true });
+    rmSync(join(TEST_DIR, "metadata.json"), { force: true });
+  } catch {
+    // ignore
+  }
+});
+
+function setupCredential(
+  service: string,
+  opts?: { expiresAt?: number; grantedScopes?: string[] },
+) {
+  setSecureKey(credentialKey(service, "access_token"), "test-access-token");
+  setSecureKey(credentialKey(service, "refresh_token"), "test-refresh-token");
+  setSecureKey(credentialKey(service, "client_secret"), "test-client-secret");
+  upsertCredentialMetadata(service, "access_token", {
+    expiresAt: opts?.expiresAt ?? Date.now() + 3600 * 1000,
+    grantedScopes: opts?.grantedScopes ?? ["read", "write"],
+    oauth2TokenUrl: "https://oauth2.googleapis.com/token",
+    oauth2ClientId: "test-client-id",
+    hasRefreshToken: true,
+  });
+}
+
+function createConnection(service = "integration:gmail"): BYOOAuthConnection {
+  return new BYOOAuthConnection({
+    id: "test-cred-id",
+    providerKey: service,
+    baseUrl: "https://gmail.googleapis.com/gmail/v1/users/me",
+    accountInfo: null,
+    grantedScopes: ["read", "write"],
+    credentialService: service,
+  });
+}
+
+// ---------------------------------------------------------------------------
+// Tests
+// ---------------------------------------------------------------------------
+
+describe("BYOOAuthConnection", () => {
+  describe("request()", () => {
+    test("makes authenticated request with Bearer token", async () => {
+      setupCredential("integration:gmail");
+      const conn = createConnection();
+
+      const result = await conn.request({
+        method: "GET",
+        path: "/messages",
+      });
+
+      expect(result.status).toBe(200);
+      expect(result.body).toEqual({ ok: true });
+      expect(mockFetch).toHaveBeenCalledTimes(1);
+
+      const [url, init] = mockFetch.mock.calls[0];
+      expect(url).toBe(
+        "https://gmail.googleapis.com/gmail/v1/users/me/messages",
+      );
+      expect((init as RequestInit).headers).toMatchObject({
+        Authorization: "Bearer test-access-token",
+        "Content-Type": "application/json",
+      });
+      expect((init as RequestInit).method).toBe("GET");
+    });
+
+    test("appends query parameters", async () => {
+      setupCredential("integration:gmail");
+      const conn = createConnection();
+
+      await conn.request({
+        method: "GET",
+        path: "/messages",
+        query: { maxResults: "10", labelIds: "INBOX" },
+      });
+
+      const [url] = mockFetch.mock.calls[0];
+      const parsed = new URL(url as string);
+      expect(parsed.searchParams.get("maxResults")).toBe("10");
+      expect(parsed.searchParams.get("labelIds")).toBe("INBOX");
+    });
+
+    test("uses per-request baseUrl override", async () => {
+      setupCredential("integration:gmail");
+      const conn = createConnection();
+
+      await conn.request({
+        method: "GET",
+        path: "/calendars",
+        baseUrl: "https://www.googleapis.com/calendar/v3",
+      });
+
+      const [url] = mockFetch.mock.calls[0];
+      expect(url).toBe("https://www.googleapis.com/calendar/v3/calendars");
+    });
+
+    test("sends JSON body for POST requests", async () => {
+      setupCredential("integration:gmail");
+      const conn = createConnection();
+
+      await conn.request({
+        method: "POST",
+        path: "/messages/send",
+        body: { raw: "base64-encoded-email" },
+      });
+
+      const [, init] = mockFetch.mock.calls[0];
+      expect((init as RequestInit).body).toBe(
+        JSON.stringify({ raw: "base64-encoded-email" }),
+      );
+      expect((init as RequestInit).method).toBe("POST");
+    });
+
+    test("retries once on 401 response", async () => {
+      setupCredential("integration:gmail");
+      const conn = createConnection();
+
+      // First call returns 401, second returns 200
+      let callCount = 0;
+      globalThis.fetch = mock(() => {
+        callCount++;
+        if (callCount === 1) {
+          return Promise.resolve(new Response("Unauthorized", { status: 401 }));
+        }
+        return Promise.resolve(
+          new Response(JSON.stringify({ ok: true }), { status: 200 }),
+        );
+      }) as unknown as typeof fetch;
+
+      const result = await conn.request({
+        method: "GET",
+        path: "/messages",
+      });
+
+      expect(result.status).toBe(200);
+      expect(result.body).toEqual({ ok: true });
+      expect(callCount).toBe(2);
+      // Verify refresh was called
+      expect(mockRefreshOAuth2Token).toHaveBeenCalled();
+    });
+
+    test("handles empty response body", async () => {
+      setupCredential("integration:gmail");
+      const conn = createConnection();
+
+      globalThis.fetch = mock(() =>
+        Promise.resolve(new Response("", { status: 204 })),
+      ) as unknown as typeof fetch;
+
+      const result = await conn.request({
+        method: "DELETE",
+        path: "/messages/123",
+      });
+
+      expect(result.status).toBe(204);
+      expect(result.body).toBeNull();
+    });
+
+    test("handles non-JSON response body", async () => {
+      setupCredential("integration:gmail");
+      const conn = createConnection();
+
+      globalThis.fetch = mock(() =>
+        Promise.resolve(new Response("plain text response", { status: 200 })),
+      ) as unknown as typeof fetch;
+
+      const result = await conn.request({
+        method: "GET",
+        path: "/raw",
+      });
+
+      expect(result.status).toBe(200);
+      expect(result.body).toBe("plain text response");
+    });
+
+    test("returns response headers", async () => {
+      setupCredential("integration:gmail");
+      const conn = createConnection();
+
+      globalThis.fetch = mock(() =>
+        Promise.resolve(
+          new Response(JSON.stringify({}), {
+            status: 200,
+            headers: {
+              "x-ratelimit-remaining": "99",
+              "content-type": "application/json",
+            },
+          }),
+        ),
+      ) as unknown as typeof fetch;
+
+      const result = await conn.request({
+        method: "GET",
+        path: "/messages",
+      });
+
+      expect(result.headers["x-ratelimit-remaining"]).toBe("99");
+    });
+
+    test("includes custom request headers", async () => {
+      setupCredential("integration:gmail");
+      const conn = createConnection();
+
+      await conn.request({
+        method: "GET",
+        path: "/messages",
+        headers: { "X-Custom-Header": "custom-value" },
+      });
+
+      const [, init] = mockFetch.mock.calls[0];
+      expect((init as RequestInit).headers).toMatchObject({
+        "X-Custom-Header": "custom-value",
+        Authorization: "Bearer test-access-token",
+      });
+    });
+  });
+
+  describe("proactive token refresh", () => {
+    test("refreshes token when near expiry (within 5-minute buffer)", async () => {
+      // Set token to expire in 2 minutes (within 5-min buffer)
+      setupCredential("integration:gmail", {
+        expiresAt: Date.now() + 2 * 60 * 1000,
+      });
+      const conn = createConnection();
+
+      await conn.request({
+        method: "GET",
+        path: "/messages",
+      });
+
+      // Token should have been refreshed proactively
+      expect(mockRefreshOAuth2Token).toHaveBeenCalled();
+
+      // The request should use the refreshed token
+      const [, init] = mockFetch.mock.calls[0];
+      expect((init as RequestInit).headers).toMatchObject({
+        Authorization: "Bearer refreshed-access-token",
+      });
+    });
+  });
+
+  describe("withToken()", () => {
+    test("provides valid token to callback", async () => {
+      setupCredential("integration:gmail");
+      const conn = createConnection();
+
+      const result = await conn.withToken(async (token) => {
+        return `got-${token}`;
+      });
+
+      expect(result).toBe("got-test-access-token");
+    });
+
+    test("retries callback on 401 error", async () => {
+      setupCredential("integration:gmail");
+      const conn = createConnection();
+
+      let callCount = 0;
+      const result = await conn.withToken(async (token) => {
+        callCount++;
+        if (callCount === 1) {
+          const err = new Error("Unauthorized");
+          (err as Error & { status: number }).status = 401;
+          throw err;
+        }
+        return `got-${token}`;
+      });
+
+      expect(callCount).toBe(2);
+      expect(result).toBe("got-refreshed-access-token");
+      expect(mockRefreshOAuth2Token).toHaveBeenCalled();
+    });
+  });
+
+  describe("missing credential", () => {
+    test("throws when no access token exists", async () => {
+      const conn = createConnection();
+
+      await expect(
+        conn.request({ method: "GET", path: "/messages" }),
+      ).rejects.toThrow(/No access token found/);
+    });
+  });
+});
+
+describe("resolveOAuthConnection", () => {
+  test("returns a BYOOAuthConnection for valid credential", () => {
+    setupCredential("integration:gmail");
+    const conn = resolveOAuthConnection("integration:gmail");
+
+    expect(conn).toBeInstanceOf(BYOOAuthConnection);
+    expect(conn.providerKey).toBe("integration:gmail");
+    expect(conn.grantedScopes).toEqual(["read", "write"]);
+  });
+
+  test("throws when no credential metadata exists", () => {
+    expect(() => resolveOAuthConnection("integration:unknown")).toThrow(
+      /No credential found for "integration:unknown"/,
+    );
+  });
+
+  test("throws when no base URL configured", () => {
+    setupCredential("integration:custom-service");
+    expect(() => resolveOAuthConnection("integration:custom-service")).toThrow(
+      /No base URL configured for "integration:custom-service"/,
+    );
+  });
+});

package/src/oauth/byo-connection.ts

@@ -0,0 +1,112 @@
+/**
+ * BYO (Bring-Your-Own) OAuth connection implementation.
+ *
+ * Wraps the existing `withValidToken` token management infrastructure to
+ * provide the OAuthConnection interface. Delegates all token resolution,
+ * proactive refresh, circuit breaker, and retry-on-401 logic to
+ * `withValidToken` from `token-manager.ts`.
+ */
+
+import { withValidToken } from "../security/token-manager.js";
+import { getLogger } from "../util/logger.js";
+import type {
+  OAuthConnection,
+  OAuthConnectionRequest,
+  OAuthConnectionResponse,
+} from "./connection.js";
+
+const log = getLogger("byo-oauth-connection");
+
+/** Default per-request timeout to prevent hung requests from blocking indefinitely. */
+const REQUEST_TIMEOUT_MS = 30_000;
+
+export interface BYOOAuthConnectionOptions {
+  id: string;
+  providerKey: string;
+  baseUrl: string;
+  accountInfo: string | null;
+  grantedScopes: string[];
+  credentialService: string;
+}
+
+export class BYOOAuthConnection implements OAuthConnection {
+  readonly id: string;
+  readonly providerKey: string;
+  readonly accountInfo: string | null;
+  readonly grantedScopes: string[];
+
+  private readonly baseUrl: string;
+  private readonly credentialService: string;
+
+  constructor(opts: BYOOAuthConnectionOptions) {
+    this.id = opts.id;
+    this.providerKey = opts.providerKey;
+    this.baseUrl = opts.baseUrl;
+    this.accountInfo = opts.accountInfo;
+    this.grantedScopes = opts.grantedScopes;
+    this.credentialService = opts.credentialService;
+  }
+
+  async request(req: OAuthConnectionRequest): Promise<OAuthConnectionResponse> {
+    return withValidToken(this.credentialService, async (token) => {
+      const effectiveBaseUrl = req.baseUrl ?? this.baseUrl;
+      let fullUrl = `${effectiveBaseUrl}${req.path}`;
+
+      if (req.query && Object.keys(req.query).length > 0) {
+        const params = new URLSearchParams(req.query);
+        fullUrl += `?${params.toString()}`;
+      }
+
+      log.debug(
+        { method: req.method, url: fullUrl, provider: this.providerKey },
+        "Making authenticated request",
+      );
+
+      const resp = await fetch(fullUrl, {
+        method: req.method,
+        headers: {
+          ...(req.body ? { "Content-Type": "application/json" } : {}),
+          ...req.headers,
+          Authorization: `Bearer ${token}`,
+        },
+        body: req.body ? JSON.stringify(req.body) : undefined,
+        signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS),
+      });
+
+      if (resp.status === 401) {
+        // Throw with a status property so withValidToken detects the 401
+        // and triggers its refresh-and-retry logic.
+        const err = new Error(`HTTP 401 from ${this.providerKey}`);
+        (err as Error & { status: number }).status = 401;
+        throw err;
+      }
+
+      return buildResponse(resp);
+    });
+  }
+
+  async withToken<T>(fn: (token: string) => Promise<T>): Promise<T> {
+    return withValidToken(this.credentialService, fn);
+  }
+}
+
+async function buildResponse(resp: Response): Promise<OAuthConnectionResponse> {
+  const headers: Record<string, string> = {};
+  resp.headers.forEach((value, key) => {
+    headers[key] = value;
+  });
+
+  let body: unknown;
+  const text = await resp.text().catch(() => "");
+  if (text) {
+    try {
+      body = JSON.parse(text);
+    } catch {
+      body = text;
+    }
+  } else {
+    body = null;
+  }
+
+  return { status: resp.status, headers, body };
+}

package/src/oauth/connect-orchestrator.ts

@@ -52,6 +52,18 @@ export interface OAuthConnectOptions {
   /** Tools allowed to use the resulting credential. */
   allowedTools?: string[];
 
+  /**
+   * Called when the deferred (non-interactive) flow completes — either
+   * successfully after tokens are stored, or on failure. Lets callers
+   * surface the outcome via SSE events, logs, etc.
+   */
+  onDeferredComplete?: (result: {
+    success: boolean;
+    service: string;
+    accountInfo?: string;
+    error?: string;
+  }) => void;
+
   // Optional overrides — when provided, these take precedence over the
   // provider profile. This lets callers connect custom / unknown providers.
   authUrl?: string;
@@ -214,11 +226,21 @@ export async function orchestrateOAuthConnect(
               },
               "Deferred OAuth2 flow completed — tokens stored",
             );
+            options.onDeferredComplete?.({
+              success: true,
+              service: resolvedService,
+              accountInfo: stored.accountInfo ?? accountInfo,
+            });
           } catch (err) {
             log.error(
               { err, service: resolvedService },
               "Failed to store tokens from deferred OAuth2 flow",
             );
+            options.onDeferredComplete?.({
+              success: false,
+              service: resolvedService,
+              error: err instanceof Error ? err.message : "Unknown error",
+            });
           }
         })
         .catch((err) => {
@@ -226,6 +248,11 @@ export async function orchestrateOAuthConnect(
             { err, service: resolvedService },
             "Deferred OAuth2 flow failed",
           );
+          options.onDeferredComplete?.({
+            success: false,
+            service: resolvedService,
+            error: err instanceof Error ? err.message : "Unknown error",
+          });
         });
 
       return {

package/src/oauth/connection-resolver.ts

@@ -0,0 +1,34 @@
+import { getCredentialMetadata } from "../tools/credentials/metadata-store.js";
+import { BYOOAuthConnection } from "./byo-connection.js";
+import type { OAuthConnection } from "./connection.js";
+import { getProviderBaseUrl } from "./provider-base-urls.js";
+
+/**
+ * Resolve an OAuthConnection for a given credential service.
+ * Currently always creates a BYO connection from existing credential store.
+ * Will be extended to create Platform connections when storage model supports it.
+ */
+export function resolveOAuthConnection(
+  credentialService: string,
+): OAuthConnection {
+  const meta = getCredentialMetadata(credentialService, "access_token");
+  if (!meta) {
+    throw new Error(
+      `No credential found for "${credentialService}". Authorization required.`,
+    );
+  }
+
+  const baseUrl = getProviderBaseUrl(credentialService);
+  if (!baseUrl) {
+    throw new Error(`No base URL configured for "${credentialService}".`);
+  }
+
+  return new BYOOAuthConnection({
+    id: meta.credentialId,
+    providerKey: credentialService,
+    baseUrl,
+    accountInfo: null,
+    grantedScopes: meta.grantedScopes ?? [],
+    credentialService,
+  });
+}

package/src/oauth/connection.ts

@@ -0,0 +1,38 @@
+export interface OAuthConnectionRequest {
+  method: string;
+  path: string; // relative, e.g. "/2/tweets"
+  query?: Record<string, string>;
+  headers?: Record<string, string>;
+  body?: unknown; // JSON-serializable
+  /**
+   * Override the connection's default base URL for this request.
+   * Required for providers that span multiple API hosts sharing
+   * one OAuth token (e.g. Google: Gmail, Calendar, People all
+   * use the same credential but different base URLs).
+   */
+  baseUrl?: string;
+}
+
+export interface OAuthConnectionResponse {
+  status: number;
+  headers: Record<string, string>;
+  body: unknown;
+}
+
+export interface OAuthConnection {
+  /** Make an authenticated HTTP request through this connection. */
+  request(req: OAuthConnectionRequest): Promise<OAuthConnectionResponse>;
+
+  /**
+   * Execute a callback with a valid raw access token. This is an escape hatch
+   * for provider-specific endpoints that don't fit the relative-path model
+   * (e.g. Gmail batch API on a different host). Throws for platform connections
+   * where raw tokens are not available locally.
+   */
+  withToken<T>(fn: (token: string) => Promise<T>): Promise<T>;
+
+  readonly id: string;
+  readonly providerKey: string;
+  readonly accountInfo: string | null;
+  readonly grantedScopes: string[];
+}