@vellumai/assistant 0.11.2 → 0.11.3-dev.202608080427.92f668a

package/src/mcp/__tests__/mcp-oauth-client-registration.test.ts

@@ -0,0 +1,196 @@
+/**
+ * Reuse and invalidation of the stored dynamic client registration.
+ *
+ * Dynamic registration writes a record on the authorization server, so
+ * registering per attempt leaves records behind that nobody cleans up. The
+ * stored registration is therefore reused, and only withheld when it
+ * provably no longer applies. Both halves matter: reusing one bound to a
+ * stale redirect URI produces authorization requests the server rejects,
+ * and re-registering unnecessarily is the accumulation this exists to stop.
+ */
+
+import { beforeEach, describe, expect, mock, test } from "bun:test";
+
+const CALLBACK_URL =
+  "https://platform.example/v1/gateway/callbacks/abc/webhooks/oauth/callback/";
+
+let resolvedCallbackUrl = CALLBACK_URL;
+
+mock.module("../../inbound/platform-callback-registration.js", () => ({
+  resolveCallbackUrl: async () => resolvedCallbackUrl,
+}));
+
+mock.module("../../inbound/public-ingress-urls.js", () => ({
+  getOAuthCallbackUrl: () => resolvedCallbackUrl,
+}));
+
+mock.module("../../config/loader.js", () => ({
+  loadConfig: () => ({}),
+}));
+
+/** In-memory stand-in for the credential store. */
+const store = new Map<string, string>();
+
+mock.module("../../security/secure-keys.js", () => ({
+  getSecureKeyAsync: async (k: string) => store.get(k) ?? null,
+  setSecureKeyAsync: async (k: string, v: string) => {
+    store.set(k, v);
+    return true;
+  },
+  deleteSecureKeyAsync: async (k: string) =>
+    store.delete(k) ? "deleted" : "not-found",
+}));
+
+mock.module("../../daemon/identity-helpers.js", () => ({
+  getAssistantName: () => "Jarvis",
+}));
+
+mock.module("../../config/env.js", () => ({
+  getPlatformAssistantId: () => "019d6d4f-6dbd-779f-91d3-cb273b9429a5",
+}));
+
+const { McpOAuthProvider } = await import("../mcp-oauth-provider.js");
+
+const REGISTRATION = { client_id: "generated-client-id" };
+
+function newProvider() {
+  return new McpOAuthProvider(
+    "unabyss",
+    "https://mcp.unabyss.com",
+    /* interactive */ false,
+  );
+}
+
+/** Seed discovery state the way the SDK would after running discovery. */
+async function seedDiscovery(
+  provider: InstanceType<typeof McpOAuthProvider>,
+  issuer: string,
+): Promise<void> {
+  await provider.saveDiscoveryState({
+    authorizationServerUrl: issuer,
+    authorizationServerMetadata: { issuer },
+  } as never);
+}
+
+/** Run one registration the way a completed flow would. */
+async function register(issuer = "https://mcp.unabyss.com"): Promise<void> {
+  const provider = newProvider();
+  await provider.startCallbackServer();
+  await seedDiscovery(provider, issuer);
+  await provider.saveClientInformation(REGISTRATION as never);
+  provider.stopCallbackServer();
+}
+
+beforeEach(() => {
+  store.clear();
+  resolvedCallbackUrl = CALLBACK_URL;
+});
+
+describe("McpOAuthProvider client registration reuse", () => {
+  test("a stored registration is reused on the next flow", async () => {
+    await register();
+
+    const next = newProvider();
+    await next.startCallbackServer();
+    await seedDiscovery(next, "https://mcp.unabyss.com");
+
+    // Returning the stored value is what keeps the server from accruing a
+    // fresh client record on every `assistant mcp auth`.
+    expect(await next.clientInformation()).toEqual(REGISTRATION as never);
+  });
+
+  test("no stored registration means the SDK registers", async () => {
+    const provider = newProvider();
+    await provider.startCallbackServer();
+    expect(await provider.clientInformation()).toBeUndefined();
+  });
+
+  test("a changed redirect URI withholds the registration", async () => {
+    await register();
+
+    // The assistant moved, so its callback URL moved with it. The stored
+    // client is registered for the old one and the server would reject it.
+    resolvedCallbackUrl = "https://other.example/webhooks/oauth/callback";
+    const next = newProvider();
+    await next.startCallbackServer();
+    await seedDiscovery(next, "https://mcp.unabyss.com");
+
+    expect(await next.clientInformation()).toBeUndefined();
+  });
+
+  test("a changed authorization server withholds the registration", async () => {
+    await register("https://auth.unabyss.com");
+
+    const next = newProvider();
+    await next.startCallbackServer();
+    await seedDiscovery(next, "https://different-auth.example");
+
+    expect(await next.clientInformation()).toBeUndefined();
+  });
+
+  test("a silent reconnect reuses the registration without a redirect URI", async () => {
+    await register();
+
+    // `McpClient.connect` never prepares a callback, so `_redirectUrl` is
+    // unset. Treating that as a mismatch would turn every background
+    // reconnect into a registration.
+    const silent = newProvider();
+    await seedDiscovery(silent, "https://mcp.unabyss.com");
+
+    expect(await silent.clientInformation()).toEqual(REGISTRATION as never);
+  });
+
+  test("a registration stored before bindings existed is kept", async () => {
+    // Upgrading must not invalidate every existing registration at once.
+    store.set(
+      "mcp:unabyss:client_info",
+      JSON.stringify({ client_id: "legacy" }),
+    );
+
+    const provider = newProvider();
+    await provider.startCallbackServer();
+    await seedDiscovery(provider, "https://mcp.unabyss.com");
+
+    expect(await provider.clientInformation()).toEqual({
+      client_id: "legacy",
+    } as never);
+  });
+
+  test("invalidating the client drops its binding too", async () => {
+    await register();
+    expect(store.has("mcp:unabyss:client_binding")).toBe(true);
+
+    const provider = newProvider();
+    await provider.invalidateCredentials("client");
+
+    expect(store.has("mcp:unabyss:client_info")).toBe(false);
+    // A surviving binding would let the next registration inherit the
+    // previous one's issuer and redirect URI.
+    expect(store.has("mcp:unabyss:client_binding")).toBe(false);
+  });
+});
+
+describe("McpOAuthProvider client metadata", () => {
+  test("registers under the assistant's own name", async () => {
+    const provider = newProvider();
+    expect(provider.clientMetadata.client_name).toEqual("Jarvis");
+  });
+
+  test("carries the assistant id as software_id", async () => {
+    const provider = newProvider();
+    expect(provider.clientMetadata.software_id).toEqual(
+      "019d6d4f-6dbd-779f-91d3-cb273b9429a5",
+    );
+  });
+
+  test("omits logo_uri, which has no anonymously fetchable URL", async () => {
+    const provider = newProvider();
+    expect(provider.clientMetadata.logo_uri).toBeUndefined();
+  });
+
+  test("registers the resolved callback URL as the redirect URI", async () => {
+    const provider = newProvider();
+    await provider.startCallbackServer();
+    expect(provider.clientMetadata.redirect_uris).toEqual([CALLBACK_URL]);
+  });
+});

package/src/mcp/__tests__/mcp-oauth-provider.test.ts

@@ -16,22 +16,24 @@ mock.module("../../inbound/public-ingress-urls.js", () => ({
 
 mock.module("../../config/loader.js", () => ({
   loadConfig: () => ({}),
+  // `clientMetadata` reads the assistant's name, which pulls the persona
+  // resolver into the provider's import graph, and that reads config.
+  getConfig: () => ({}),
 }));
 
 // ── Import SUT after mocks ────────────────────────────────────────────────────
 
 const { McpOAuthProvider } = await import("../mcp-oauth-provider.js");
 
-function newGatewayProvider() {
+function newProvider() {
   return new McpOAuthProvider(
     "comms",
     "https://comms.example/mcp",
     /* interactive */ false,
-    "gateway",
   );
 }
 
-describe("McpOAuthProvider gateway callback", () => {
+describe("McpOAuthProvider callback", () => {
   test("stopCallbackServer() before a consumer attaches does not emit an unhandled rejection", async () => {
     // stopCallbackServer() rejects the deferred code promise. When no consumer
     // is attached, that rejection must stay observed so it does not surface as
@@ -40,7 +42,7 @@ describe("McpOAuthProvider gateway callback", () => {
     const onUnhandled = (reason: unknown) => unhandled.push(reason);
     process.on("unhandledRejection", onUnhandled);
     try {
-      const provider = newGatewayProvider();
+      const provider = newProvider();
 
       // Intentionally do NOT consume the returned codePromise — this mirrors the
       // early-exit paths where connect() throws before the OAuth tail is wired up.
@@ -53,7 +55,7 @@ describe("McpOAuthProvider gateway callback", () => {
       const cancelRejections = unhandled.filter(
         (r) =>
           r instanceof Error &&
-          r.message.includes("MCP OAuth gateway callback cancelled"),
+          r.message.includes("MCP OAuth callback cancelled"),
       );
       expect(cancelRejections).toHaveLength(0);
     } finally {
@@ -62,14 +64,12 @@ describe("McpOAuthProvider gateway callback", () => {
   });
 
   test("stopCallbackServer() still rejects a consumer that is awaiting the code", async () => {
-    const provider = newGatewayProvider();
+    const provider = newProvider();
     await provider.startCallbackServer();
 
     const codePromise = provider.waitForCode();
     provider.stopCallbackServer();
 
-    await expect(codePromise).rejects.toThrow(
-      "MCP OAuth gateway callback cancelled",
-    );
+    await expect(codePromise).rejects.toThrow("MCP OAuth callback cancelled");
   });
 });

package/src/mcp/client.ts

@@ -3,8 +3,8 @@ import { Client } from "@modelcontextprotocol/sdk/client/index.js";
 import { SSEClientTransport } from "@modelcontextprotocol/sdk/client/sse.js";
 import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js";
 import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
+import type { ToolAnnotations } from "@modelcontextprotocol/sdk/types.js";
 
-import { getIsPlatform } from "../config/env-registry.js";
 import type { McpTransport } from "../config/schemas/mcp.js";
 import { getSecureKeyAsync } from "../security/secure-keys.js";
 import { getLogger } from "../util/logger.js";
@@ -33,10 +33,18 @@ function normalizeInputSchema(
   return { type: "object", properties: {} };
 }
 
+/**
+ * Behavioral hints a server may attach to a tool (MCP `ToolAnnotations`).
+ * Self-reported by the server, so consumers treat them as hints rather than
+ * guarantees.
+ */
+export type McpToolAnnotations = ToolAnnotations;
+
 export interface McpToolInfo {
   name: string;
   description: string;
   inputSchema: Record<string, unknown>;
+  annotations?: McpToolAnnotations;
 }
 
 export interface McpCallResult {
@@ -92,21 +100,19 @@ export class McpClient {
       transportConfig.type === "sse" ||
       transportConfig.type === "streamable-http";
 
-    // For HTTP transports, only attach an OAuth provider if cached tokens exist.
-    // This avoids triggering client registration (which binds to a random port)
-    // during daemon startup. If no tokens, try without auth — if the server
-    // requires it, skip silently.
+    // For HTTP transports, only attach an OAuth provider if cached tokens
+    // exist. This avoids triggering client registration during daemon
+    // startup. If no tokens, try without auth: if the server requires it,
+    // skip silently.
     if (isHttpTransport) {
       const cachedTokens = await getSecureKeyAsync(
         `mcp:${this.serverId}:tokens`,
       );
       if (cachedTokens) {
-        const callbackTransport = getIsPlatform() ? "gateway" : "loopback";
         this.oauthProvider = new McpOAuthProvider(
           this.serverId,
           transportConfig.url,
           /* interactive */ false,
-          callbackTransport,
         );
       }
     }
@@ -197,6 +203,7 @@ export class McpClient {
       name: tool.name,
       description: tool.description ?? "",
       inputSchema: normalizeInputSchema(tool.inputSchema, tool.name),
+      annotations: tool.annotations,
     }));
   }
 

package/src/mcp/mcp-auth-orchestrator.ts

@@ -13,7 +13,6 @@ import { Client } from "@modelcontextprotocol/sdk/client/index.js";
 import { SSEClientTransport } from "@modelcontextprotocol/sdk/client/sse.js";
 import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
 
-import { getIsContainerized } from "../config/env-registry.js";
 import { reloadMcpServers } from "../daemon/mcp-reload-service.js";
 import { getLogger } from "../util/logger.js";
 import {
@@ -22,10 +21,7 @@ import {
   setMcpAuthPending,
 } from "./mcp-auth-state.js";
 import { getMcpHeaders } from "./mcp-header-store.js";
-import {
-  type McpOAuthCallbackTransport,
-  McpOAuthProvider,
-} from "./mcp-oauth-provider.js";
+import { McpOAuthProvider } from "./mcp-oauth-provider.js";
 
 const log = getLogger("mcp-auth-orchestrator");
 
@@ -53,20 +49,11 @@ export async function orchestrateMcpOAuthConnect(args: {
 }): Promise<OrchestrateMcpOAuthConnectResult> {
   const { serverId, transport } = args;
 
-  // Containerized deployments (platform-managed AND self-hosted Docker) must
-  // use the gateway transport: the browser is outside the daemon container,
-  // so a loopback callback inside the container is unreachable. Bare-metal
-  // daemons can use loopback as before.
-  const callbackTransport: McpOAuthCallbackTransport = getIsContainerized()
-    ? "gateway"
-    : "loopback";
-
   let capturedAuthUrl: string | undefined;
   const provider = new McpOAuthProvider(
     serverId,
     transport.url,
     /* interactive */ false,
-    callbackTransport,
     {
       onAuthorizationUrl: (url) => {
         capturedAuthUrl = url;
@@ -74,8 +61,11 @@ export async function orchestrateMcpOAuthConnect(args: {
     },
   );
 
-  // Clear stale credentials so the flow starts fresh
-  await provider.invalidateCredentials("client");
+  // Re-discover, but keep the client registration. Discovery is a read and
+  // costs a request; registration is a write on the authorization server,
+  // and remaking it per attempt leaves a record behind every time. The
+  // provider drops a stored registration on its own when the redirect URI
+  // or the authorization server it was made against has changed.
   await provider.invalidateCredentials("discovery");
 
   // Register the pending callback in the daemon heap
@@ -146,10 +136,9 @@ export async function orchestrateMcpOAuthConnect(args: {
   // so the daemon can reconnect on the next MCP reload without re-connecting here.
   void (async () => {
     try {
-      // Apply an explicit timeout: the loopback callback path has a built-in
-      // 2-minute timer, but the gateway transport's deferred promise relies on
-      // the caller for time-boxing. Without this race, gateway-mode tails leak
-      // forever if the user never completes the OAuth handshake.
+      // Apply an explicit timeout: the deferred code promise relies on the
+      // caller for time-boxing. Without this race, a tail leaks forever if
+      // the user never completes the OAuth handshake.
       const code = await Promise.race([
         codePromise,
         new Promise<never>((_, reject) => {
@@ -215,6 +204,8 @@ export async function orchestrateMcpOAuthConnect(args: {
   return { auth_url: capturedAuthUrl };
 }
 
-// Matches the loopback callback timeout; keep both in lockstep so loopback
-// and gateway transports time-box OAuth identically.
+// How long a user has to complete the browser handshake before the tail
+// gives up. Shorter than the callback registry's own 5-minute TTL, so the
+// failure surfaces here with a clear message rather than as a registry
+// expiry the caller cannot attribute.
 const CALLBACK_TIMEOUT_MS = 2 * 60 * 1000;