From ebc2acd124cf250b7c420a661ca2380b5d3ae774 Mon Sep 17 00:00:00 2001 From: jonnyparris <6400000+jonnyparris@users.noreply.github.com> Date: Tue, 26 May 2026 13:58:55 +0100 Subject: [PATCH] feat(mcp): refresh-token auth type + set_refresh_token_mcp tool MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Lets Dodo connect to MCP servers whose OAuth authorize endpoint only accepts loopback redirect URIs — the cf-portal case, where DCR works from any host but the authorize step rejects non-loopback redirect URIs. The flow: 1. A local helper (e.g. OpenCode, or a small script) performs the OAuth authorization-code flow against the upstream provider, registering itself with a 127.0.0.1 redirect URI that the upstream accepts. 2. The helper calls Dodo's new `set_refresh_token_mcp` MCP tool with { name, url, tokenEndpoint, clientId, accessToken, refreshToken, expiresAt? }. 3. Dodo stores them encrypted in UserControl. The mcp_configs row gets auth_type = 'refresh_token' and two new non-secret columns (oauth_token_endpoint, oauth_client_id) for the refresh call. 4. Every session DO that needs the access token reads it via /mcp-configs/:id/access-token. UserControl owns the cache + refresh loop. Per-user DO single-threading collapses concurrent reads near expiry into one refresh — exactly the property that prevents the single-use refresh token race that Kenny Johnson's wiki page documents. 5. Session DOs reconnect-once-on-401 with ?force=1 to force a refresh if the cached token is stale despite a not-yet-elapsed expires_at (clock skew, server-side revoke). The OAuth tokens (access, refresh, expires_at) live in encrypted_secrets, bundled with the existing envelope encryption. The non-secret OAuth metadata (token endpoint URL, client_id) lives in two new nullable columns on mcp_configs that are idempotently created via ALTER TABLE on every onStart. Six new unit tests in test/refresh-token-mcp-unit.test.ts exercise: - create-on-new-URL vs update-in-place - cached read (no fetch made when token still valid) - refresh on expiry - force-refresh via ?force=1 - 502 surfacing when the token endpoint rejects the refresh The existing connectMcpServers static-headers path is unchanged. oauth (Agents SDK-managed) configs are still filtered out. The only new behaviour is when auth_type === 'refresh_token', where the Authorization header is sourced from UserControl. Tests: 823/823 pass. Typecheck clean. beep-boop-🤖 --- src/coding-agent.ts | 89 +++++++-- src/mcp-client.ts | 10 +- src/mcp.ts | 55 +++++ src/user-control.ts | 297 ++++++++++++++++++++++++++- test/refresh-token-mcp-unit.test.ts | 298 ++++++++++++++++++++++++++++ 5 files changed, 732 insertions(+), 17 deletions(-) create mode 100644 test/refresh-token-mcp-unit.test.ts diff --git a/src/coding-agent.ts b/src/coding-agent.ts index 8212936..6857756 100644 --- a/src/coding-agent.ts +++ b/src/coding-agent.ts @@ -6475,7 +6475,10 @@ export class CodingAgent extends Think { // even if the MCP config itself is enabled. const browserEnabled = this.readMetadata("browser_enabled") === "true"; - // Filter to enabled HTTP configs with URLs + // Filter to enabled HTTP configs with URLs. `oauth` (SDK-managed) is + // federated through the per-user hub DO and never connected from + // session DOs directly. `refresh_token` is connected here with a + // bearer header sourced from UserControl, which owns the refresh. const enabled = configs.filter((c) => { if (!c.enabled || c.type !== "http" || !c.url) return false; if (c.auth_type === "oauth") return false; @@ -6484,13 +6487,32 @@ export class CodingAgent extends Think { }); if (enabled.length === 0) return; + // Helper: ask UserControl for a current access token. UserControl + // refreshes if expired (serialised by per-user DO single-threading). + const fetchRefreshTokenBearer = async (configId: string): Promise => { + const tokenRes = await stub.fetch( + `https://user-control/mcp-configs/${encodeURIComponent(configId)}/access-token`, + { headers: { "x-owner-email": ownerEmail } }, + ); + if (!tokenRes.ok) return null; + const { accessToken } = (await tokenRes.json()) as { accessToken?: string }; + return accessToken ?? null; + }; + // Resolve encrypted headers and connect each gatekeeper const connected: McpClient[] = []; for (const config of enabled) { try { - // Resolve headers via internal secret endpoint + // Resolve auth headers depending on the config's auth_type. let headers: Record | undefined; - if (config.headerKeys?.length) { + + if (config.auth_type === "refresh_token") { + const accessToken = await fetchRefreshTokenBearer(config.id); + if (!accessToken) { + throw new Error("No refresh-token access token available; run set_refresh_token_mcp again"); + } + headers = { Authorization: `Bearer ${accessToken}` }; + } else if (config.headerKeys?.length) { headers = {}; for (const headerName of config.headerKeys) { const secretRes = await stub.fetch( @@ -6504,21 +6526,60 @@ export class CodingAgent extends Think { } } - const gk = new HttpMcpClient({ + let gk = new HttpMcpClient({ ...config, headers, }, this.mcpDepth); - await gk.connect(); - const tools = await gk.listTools(); // Pre-populate cache for synchronous getTools() - connected.push(gk); - this.mcpStatus.set(config.id, { - name: config.name, - url: config.url, - ok: true, - toolCount: tools.length, - lastCheckedAt: Date.now(), - }); + try { + await gk.connect(); + // Pre-populate cache for synchronous getTools() + const tools = await gk.listTools(); + connected.push(gk); + this.mcpStatus.set(config.id, { + name: config.name, + url: config.url, + ok: true, + toolCount: tools.length, + lastCheckedAt: Date.now(), + }); + } catch (innerErr) { + // For refresh-token configs, a 401/auth-style failure is + // recoverable: force a refresh and reconnect once. UserControl + // is the source of truth for the token, so we ask it to + // refresh rather than retry with the same (probably-expired) + // token we just used. + const msg = innerErr instanceof Error ? innerErr.message : String(innerErr); + const looksLikeAuthFail = /401|403|unauthor/i.test(msg); + if (config.auth_type === "refresh_token" && looksLikeAuthFail) { + try { gk.disconnect(); } catch { /* best effort */ } + const refreshRes = await stub.fetch( + `https://user-control/mcp-configs/${encodeURIComponent(config.id)}/access-token?force=1`, + { headers: { "x-owner-email": ownerEmail } }, + ); + if (refreshRes.ok) { + const { accessToken } = (await refreshRes.json()) as { accessToken?: string }; + if (accessToken) { + gk = new HttpMcpClient({ + ...config, + headers: { Authorization: `Bearer ${accessToken}` }, + }, this.mcpDepth); + await gk.connect(); + const tools = await gk.listTools(); + connected.push(gk); + this.mcpStatus.set(config.id, { + name: config.name, + url: config.url, + ok: true, + toolCount: tools.length, + lastCheckedAt: Date.now(), + }); + continue; + } + } + } + throw innerErr; + } } catch (error) { // Log but don't fail — one broken MCP server shouldn't block the session. // We also record the failure on `mcpStatus` so the UI can surface it diff --git a/src/mcp-client.ts b/src/mcp-client.ts index 01a5b27..9a34592 100644 --- a/src/mcp-client.ts +++ b/src/mcp-client.ts @@ -64,7 +64,15 @@ export interface McpClientConfig { id: string; name: string; type: "http" | "service-binding"; - auth_type: "oauth" | "static_headers"; + /** + * - `static_headers` — fixed bearer/API-key headers stored in encrypted_secrets + * - `oauth` — Agents-SDK-managed OAuth (per-user hub DO). Filtered out of + * the static MCP gatekeeper path in coding-agent.ts. + * - `refresh_token` — bearer token that auto-refreshes via OAuth refresh-token + * grant. Use when the OAuth provider only allows loopback redirect URIs + * (e.g. Cloudflare Portal) and DCR was performed by a local helper. + */ + auth_type: "oauth" | "static_headers" | "refresh_token"; url?: string; headers?: Record; /** Header key names (without values) for display purposes. */ diff --git a/src/mcp.ts b/src/mcp.ts index 6114d45..07f41c0 100644 --- a/src/mcp.ts +++ b/src/mcp.ts @@ -1474,6 +1474,61 @@ export function createDodoMcpServer(env: Env, userEmail: string, depth = 0): Mcp jsonFetch(env, "user", "/mcp-configs"), ); + // Push a refresh-token MCP config (from a local helper that already ran + // DCR + browser OAuth for an MCP server whose authorize endpoint only + // accepts loopback redirect URIs — e.g. portal.mcp.cfdata.org). Idempotent + // on `url` so re-running the local helper just rotates the tokens in place. + server.tool( + "set_refresh_token_mcp", + [ + "Register an MCP server that authenticates with an OAuth refresh token.", + "Use this when the upstream OAuth provider only accepts loopback", + "redirect URIs (so Dodo can't do the OAuth dance itself) and a local", + "helper has already completed the authorize+exchange flow. Dodo will", + "use the access token directly and refresh it via the OAuth token", + "endpoint as it expires. Idempotent on `url`: re-pushing for the same", + "MCP URL updates the stored tokens in place.", + ].join(" "), + { + name: z.string().describe("Integration display name"), + url: z.string().url().describe("MCP server endpoint URL"), + tokenEndpoint: z + .string() + .url() + .describe( + "OAuth token endpoint used to refresh the access token (e.g. https://cf-mcp.cloudflareaccess.com/cdn-cgi/access/oauth/token)", + ), + clientId: z + .string() + .min(1) + .describe("OAuth client_id from the local helper's DCR"), + accessToken: z.string().min(1).describe("Current OAuth access token"), + refreshToken: z + .string() + .min(1) + .describe("Current OAuth refresh token. Will rotate on each refresh."), + expiresAt: z + .number() + .int() + .nonnegative() + .optional() + .describe("Absolute expiry of the access token in unix seconds. When omitted, the access token is treated as expired and refreshed on first use."), + }, + async (input) => { + const res = await userControlFetch(env, "/refresh-token-mcp", { + body: JSON.stringify(input), + headers: { "content-type": "application/json" }, + method: "POST", + }); + if (!res.ok) { + const err = await res.json(); + return errorResult(err); + } + const payload = (await res.json()) as { id: string; name: string; url: string; updated: boolean }; + return textResult(payload); + }, + ); + server.tool("remove_mcp_config", "Remove an MCP integration by id", { id: z.string().describe("MCP config id to remove"), }, async ({ id }) => { diff --git a/src/user-control.ts b/src/user-control.ts index 567ed9a..8a1ff38 100644 --- a/src/user-control.ts +++ b/src/user-control.ts @@ -126,6 +126,42 @@ const mcpConfigUpdateSchema = z }) .strict(); +/** + * Schema for the loopback-OAuth piggyback path. A local helper performs the + * OAuth authorization-code flow (which the upstream's redirect-URI policy + * only allows for loopback) and then pushes the resulting tokens here. + * Tokens are stored encrypted; the access token auto-refreshes against the + * OAuth token endpoint as it expires. + */ +const refreshTokenMcpUpsertSchema = z + .object({ + name: z.string().min(1).describe("Integration display name"), + url: z.string().url().describe("MCP server endpoint URL"), + tokenEndpoint: z + .string() + .url() + .describe("OAuth token endpoint used to refresh the access token"), + clientId: z + .string() + .min(1) + .describe("OAuth client_id issued to the local helper that ran DCR"), + accessToken: z.string().min(1).describe("Current OAuth access token"), + refreshToken: z + .string() + .min(1) + .describe("Current OAuth refresh token. Will rotate on each refresh."), + expiresAt: z + .number() + .int() + .nonnegative() + .optional() + .describe( + "Absolute expiry of the access token in unix seconds. When omitted, the access token is treated as expired and refreshed on first use.", + ), + enabled: z.boolean().default(true), + }) + .strict(); + const workerRunStatusEnum = z.enum([ "session_created", "repo_ready", @@ -222,6 +258,24 @@ function toAgentMode(value: string | undefined): "inprocess" | "facet" { return value === "facet" ? "facet" : "inprocess"; } +/** + * Coerce a stored `auth_type` value into the union the rest of the worker + * knows about. Falls back to `static_headers` for unknown / legacy values so + * a typo in the DB can never produce an `undefined` field. + */ +function normalizeAuthType(value: unknown): "oauth" | "static_headers" | "refresh_token" { + if (value === "oauth") return "oauth"; + if (value === "refresh_token") return "refresh_token"; + return "static_headers"; +} + +/** + * Minimum seconds remaining before we proactively refresh an OAuth access + * token. Set high enough to absorb network round-trip + slight clock drift + * between Cloudflare's edge and the OAuth issuer. + */ +const REFRESH_TOKEN_SAFETY_MARGIN_SECONDS = 30; + /** * UserControl DO — one per user (`idFromName(email)`). * @@ -713,6 +767,43 @@ export class UserControl extends DurableObject { return Response.json({ deleted: true, id }); } + // ─── Refresh-Token MCP (loopback-OAuth piggyback) ─── + // Used for providers whose OAuth authorize endpoint only accepts + // loopback redirect URIs (e.g. cf-portal). A local helper performs + // the OAuth dance, then pushes the resulting tokens here. + + if (request.method === "POST" && url.pathname === "/refresh-token-mcp") { + const body = refreshTokenMcpUpsertSchema.parse(await request.json()); + const ownerEmail = request.headers.get("x-owner-email") ?? ""; + const result = await this.upsertRefreshTokenMcp(body, ownerEmail); + return Response.json(result, { status: 201 }); + } + + // Internal: session DOs read the current access token via this endpoint. + // UserControl serialises refresh-if-needed (single-threaded per user), + // which is what stops two parallel session DOs from racing on the + // single-use refresh token. + // + // Pass `?force=1` to skip the cache-expiry check and refresh + // unconditionally. Session DOs use that path when their connection + // fails with a 401 — the cached token must be stale even if its + // expires_at hasn't elapsed yet (clock skew, server-side revoke). + if (request.method === "GET" && url.pathname.match(/^\/mcp-configs\/[^/]+\/access-token$/)) { + const id = decodeURIComponent(url.pathname.split("/").at(-2) ?? ""); + const ownerEmail = request.headers.get("x-owner-email") ?? ""; + const force = url.searchParams.get("force") === "1"; + try { + const accessToken = force + ? await this.refreshMcpAccessToken(id, ownerEmail) + : await this.getMcpAccessToken(id, ownerEmail); + if (!accessToken) return Response.json({ error: "no token" }, { status: 404 }); + return Response.json({ accessToken }); + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + return Response.json({ error: message }, { status: 502 }); + } + } + if (request.method === "POST" && url.pathname === "/user-mcp-tokens") { const { email, label } = await request.json() as { email: string; label?: string }; const result = await this.createUserMcpToken(email, label); @@ -1308,6 +1399,8 @@ export class UserControl extends DurableObject { url TEXT, headers_json TEXT, enabled INTEGER NOT NULL DEFAULT 1, + oauth_token_endpoint TEXT, + oauth_client_id TEXT, created_at INTEGER NOT NULL, updated_at INTEGER NOT NULL ) @@ -1317,6 +1410,18 @@ export class UserControl extends DurableObject { } catch { // Column already exists. } + // Idempotent migrations for the refresh_token auth path. The columns are + // nullable because they only apply when `auth_type = 'refresh_token'`. + try { + this.ctx.storage.sql.exec("ALTER TABLE mcp_configs ADD COLUMN oauth_token_endpoint TEXT"); + } catch { + // Column already exists. + } + try { + this.ctx.storage.sql.exec("ALTER TABLE mcp_configs ADD COLUMN oauth_client_id TEXT"); + } catch { + // Column already exists. + } this.ctx.storage.sql.exec("UPDATE mcp_configs SET auth_type = 'static_headers' WHERE auth_type IS NULL OR auth_type = ''"); this.ctx.storage.sql.exec(` @@ -2130,6 +2235,194 @@ export class UserControl extends DurableObject { return Array.from(this.ctx.storage.sql.exec("SELECT id, name, type, auth_type, url, headers_json, enabled FROM mcp_configs ORDER BY name ASC")).map((row) => this.mapMcpConfigRowSafe(row)); } + // ─── Refresh-Token MCP ───────────────────────────────────────────────── + // + // Persist + refresh OAuth tokens for MCP servers whose authorize endpoint + // only accepts loopback redirect URIs (e.g. portal.mcp.cfdata.org). A local + // helper does the DCR + browser auth flow on the user's machine, then + // pushes the resulting tokens into this DO via `upsertRefreshTokenMcp`. + // + // From then on UserControl owns the refresh chain. Refresh tokens rotate + // on every refresh, so two parallel refreshes would race and the second + // would fail with `invalid_grant` (this is the same bug Kenny Johnson's + // wiki page describes for the portal). UserControl is single-threaded per + // user, so we get the necessary serialisation for free — every session + // DO that needs the access token reads via `/mcp-configs/:id/access-token` + // which proxies into `getMcpAccessToken` here. + + /** + * Find a refresh-token MCP config row keyed by URL. Used by the upsert + * path so re-pushing tokens for the same MCP URL updates the existing + * config instead of creating a duplicate. + */ + private findRefreshTokenMcpByUrl(url: string): { id: string; name: string; enabled: number } | null { + const row = Array.from(this.ctx.storage.sql.exec( + "SELECT id, name, enabled FROM mcp_configs WHERE auth_type = 'refresh_token' AND url = ?", + url, + ))[0] as SqlRow | null; + if (!row) return null; + return { + id: String(row.id), + name: String(row.name), + enabled: Number(row.enabled), + }; + } + + /** + * Look up the OAuth `token_endpoint` and `client_id` for a refresh-token MCP + * config. Returns null when either field is missing — callers treat that as + * an unconfigured / non-refresh-token config and skip the refresh path. + */ + private getRefreshTokenMcpOAuthMetadata(configId: string): { tokenEndpoint: string; clientId: string } | null { + const row = Array.from(this.ctx.storage.sql.exec( + "SELECT oauth_token_endpoint, oauth_client_id, auth_type FROM mcp_configs WHERE id = ?", + configId, + ))[0] as SqlRow | null; + if (!row) return null; + if (row.auth_type !== "refresh_token") return null; + const tokenEndpoint = row.oauth_token_endpoint ? String(row.oauth_token_endpoint) : ""; + const clientId = row.oauth_client_id ? String(row.oauth_client_id) : ""; + if (!tokenEndpoint || !clientId) return null; + return { tokenEndpoint, clientId }; + } + + /** + * Atomically create or update a refresh-token MCP config. Stores the + * non-secret fields (token endpoint, client id) in `mcp_configs` and the + * secrets (access token, refresh token, expiry) in `encrypted_secrets`. + */ + private async upsertRefreshTokenMcp( + input: z.infer, + ownerEmail: string, + ): Promise<{ id: string; name: string; url: string; updated: boolean }> { + if (!ownerEmail || !this.hasKeyEnvelope()) { + throw new Error( + "Refresh-token MCP storage requires the user passkey envelope. Run /api/passkey/init first.", + ); + } + + const now = nowEpoch(); + const existing = this.findRefreshTokenMcpByUrl(input.url); + const id = existing?.id ?? crypto.randomUUID(); + const enabled = input.enabled ? 1 : 0; + + if (existing) { + this.ctx.storage.sql.exec( + "UPDATE mcp_configs SET name = ?, type = 'http', auth_type = 'refresh_token', url = ?, headers_json = NULL, enabled = ?, oauth_token_endpoint = ?, oauth_client_id = ?, updated_at = ? WHERE id = ?", + input.name, + input.url, + enabled, + input.tokenEndpoint, + input.clientId, + now, + id, + ); + } else { + this.ctx.storage.sql.exec( + "INSERT INTO mcp_configs (id, name, type, auth_type, url, headers_json, enabled, oauth_token_endpoint, oauth_client_id, created_at, updated_at) VALUES (?, ?, 'http', 'refresh_token', ?, NULL, ?, ?, ?, ?, ?)", + id, + input.name, + input.url, + enabled, + input.tokenEndpoint, + input.clientId, + now, + now, + ); + } + + // Tokens go in encrypted_secrets so plaintext never lands in the + // mcp_configs row dump. expires_at is bundled with the secrets so all + // three load together — slight over-encryption is fine. + await this.setSecret(`mcp:${id}:access_token`, input.accessToken, ownerEmail); + await this.setSecret(`mcp:${id}:refresh_token`, input.refreshToken, ownerEmail); + await this.setSecret( + `mcp:${id}:expires_at`, + String(input.expiresAt ?? 0), + ownerEmail, + ); + + return { + id, + name: input.name, + url: input.url, + updated: !!existing, + }; + } + + /** + * Public RPC entry point: return a valid access token for the given config, + * refreshing first if the cached token is expired or close to expiry. + * + * Called by session DOs via `/mcp-configs/:id/access-token`. Single-threaded + * per-user DO semantics serialise concurrent reads — two sessions calling + * here near expiry collapse into one refresh call. + */ + public async getMcpAccessToken(configId: string, ownerEmail: string): Promise { + if (!ownerEmail || !this.hasKeyEnvelope()) return null; + const accessToken = await this.getSecret(`mcp:${configId}:access_token`, ownerEmail); + const expiresAtRaw = await this.getSecret(`mcp:${configId}:expires_at`, ownerEmail); + if (!accessToken) return null; + const expiresAt = Number(expiresAtRaw ?? "0"); + const nowSec = Math.floor(Date.now() / 1000); + if (expiresAt > 0 && expiresAt - nowSec > REFRESH_TOKEN_SAFETY_MARGIN_SECONDS) { + return accessToken; + } + // Expired or close to expiry — refresh and return the new access token. + return this.refreshMcpAccessToken(configId, ownerEmail); + } + + /** + * Force a token refresh against the OAuth token endpoint, then persist the + * rotated refresh token + new access token + new expiry. Returns the new + * access token, or null if the config is missing OAuth metadata. + * + * Errors are propagated to the caller so a session DO can choose between + * "show needs_reauth" and "treat as transient and retry". + */ + public async refreshMcpAccessToken(configId: string, ownerEmail: string): Promise { + if (!ownerEmail || !this.hasKeyEnvelope()) return null; + const meta = this.getRefreshTokenMcpOAuthMetadata(configId); + if (!meta) return null; + const refreshToken = await this.getSecret(`mcp:${configId}:refresh_token`, ownerEmail); + if (!refreshToken) return null; + + const body = new URLSearchParams({ + grant_type: "refresh_token", + refresh_token: refreshToken, + client_id: meta.clientId, + }); + const res = await fetch(meta.tokenEndpoint, { + method: "POST", + headers: { "content-type": "application/x-www-form-urlencoded" }, + body, + }); + if (!res.ok) { + const text = await res.text().catch(() => ""); + throw new Error(`Refresh failed (${res.status}): ${text.slice(0, 200)}`); + } + const payload = (await res.json()) as { + access_token?: string; + refresh_token?: string; + expires_in?: number; + }; + if (!payload.access_token) { + throw new Error("Refresh response missing access_token"); + } + + const nowSec = Math.floor(Date.now() / 1000); + const newExpiresAt = payload.expires_in ? nowSec + payload.expires_in : 0; + + await this.setSecret(`mcp:${configId}:access_token`, payload.access_token, ownerEmail); + // Some providers omit refresh_token on rotation; only update when present. + if (payload.refresh_token) { + await this.setSecret(`mcp:${configId}:refresh_token`, payload.refresh_token, ownerEmail); + } + await this.setSecret(`mcp:${configId}:expires_at`, String(newExpiresAt), ownerEmail); + + return payload.access_token; + } + public async createUserMcpToken(email: string, label?: string): Promise<{ token: string; created_at: number }> { const token = `dodo_${crypto.randomUUID().replace(/-/g, "")}`; const normalisedEmail = email.trim().toLowerCase(); @@ -2350,7 +2643,7 @@ export class UserControl extends DurableObject { id: String(row.id), name: String(row.name), type: String(row.type) as "http" | "service-binding", - auth_type: row.auth_type === "oauth" ? "oauth" : "static_headers", + auth_type: normalizeAuthType(row.auth_type), url: row.url === null ? undefined : String(row.url), headers: undefined, // Never expose header values in listing headerKeys, @@ -2378,7 +2671,7 @@ export class UserControl extends DurableObject { id: String(row.id), name: String(row.name), type: String(row.type) as "http" | "service-binding", - auth_type: row.auth_type === "oauth" ? "oauth" : "static_headers", + auth_type: normalizeAuthType(row.auth_type), url: row.url === null ? undefined : String(row.url), headers: undefined, headerKeys, diff --git a/test/refresh-token-mcp-unit.test.ts b/test/refresh-token-mcp-unit.test.ts new file mode 100644 index 0000000..c8ef9f8 --- /dev/null +++ b/test/refresh-token-mcp-unit.test.ts @@ -0,0 +1,298 @@ +/** + * Unit tests for the refresh-token MCP path. + * + * The refresh-token auth type is used for MCP servers whose OAuth authorize + * endpoint only accepts loopback redirect URIs (e.g. cf-portal). A local + * helper performs the DCR + browser auth flow and pushes the resulting + * tokens to Dodo via `set_refresh_token_mcp`. UserControl then refreshes + * the access token against the OAuth token endpoint as it expires. + * + * These tests exercise the storage + refresh paths directly against + * UserControl. They mock the outbound token-endpoint fetch so the test + * doesn't need network. The single-flight refresh property (no two + * parallel refreshes for the same user) falls out of UserControl being a + * single-threaded Durable Object — we don't bother re-testing that here. + */ + +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { env } from "cloudflare:workers"; +import type { Env } from "../src/types"; + +vi.mock("@cloudflare/codemode", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + DynamicWorkerExecutor: vi.fn(function () { + return { execute: vi.fn().mockResolvedValue({ logs: [], result: null }) }; + }) as unknown as typeof import("@cloudflare/codemode").DynamicWorkerExecutor, + }; +}); +vi.mock("../src/agentic", async () => await import("./helpers/agentic-mock")); +vi.mock("../src/notify", () => ({ dispatchNotification: vi.fn() })); + +const typedEnv = env as Env; + +const OWNER = "refresh-token-tester@dodo.test"; +const MCP_URL = "https://refresh-target.example.com/mcp"; +const TOKEN_ENDPOINT = "https://refresh-issuer.example.com/oauth/token"; + +function userControlStub() { + return typedEnv.USER_CONTROL.get(typedEnv.USER_CONTROL.idFromName(OWNER)); +} + +async function userControlFetch(path: string, init?: RequestInit): Promise { + const headers = new Headers(init?.headers); + headers.set("x-owner-email", OWNER); + return userControlStub().fetch(`https://user-control${path}`, { ...init, headers }); +} + +async function initPasskeyIfNeeded(): Promise { + // Refresh-token storage uses envelope encryption which requires the + // passkey envelope. Initialise once with a deterministic passkey so + // setSecret works. + await userControlFetch("/passkey/init", { + method: "POST", + body: JSON.stringify({ passkey: "refresh-token-tester-passkey" }), + headers: { "content-type": "application/json" }, + }).catch(() => undefined); +} + +describe("refresh-token MCP — storage", () => { + beforeEach(async () => { + await initPasskeyIfNeeded(); + }); + + it("upserting tokens for a new URL creates a refresh-token config", async () => { + const res = await userControlFetch("/refresh-token-mcp", { + method: "POST", + body: JSON.stringify({ + name: "Refresh Target", + url: MCP_URL, + tokenEndpoint: TOKEN_ENDPOINT, + clientId: "client-abc", + accessToken: "access-1", + refreshToken: "refresh-1", + expiresAt: Math.floor(Date.now() / 1000) + 3600, + }), + headers: { "content-type": "application/json" }, + }); + expect(res.status).toBe(201); + const body = (await res.json()) as { id: string; name: string; url: string; updated: boolean }; + expect(body.url).toBe(MCP_URL); + expect(body.updated).toBe(false); + + // The config appears in /mcp-configs with auth_type = refresh_token + const listRes = await userControlFetch("/mcp-configs"); + const list = (await listRes.json()) as { configs: Array<{ id: string; auth_type: string; url?: string }> }; + const entry = list.configs.find((c) => c.url === MCP_URL); + expect(entry).toBeDefined(); + expect(entry?.auth_type).toBe("refresh_token"); + }); + + it("upserting tokens for the same URL updates in place", async () => { + const first = await userControlFetch("/refresh-token-mcp", { + method: "POST", + body: JSON.stringify({ + name: "Refresh Target Updatable", + url: "https://refresh-target-update.example.com/mcp", + tokenEndpoint: TOKEN_ENDPOINT, + clientId: "client-xyz", + accessToken: "access-v1", + refreshToken: "refresh-v1", + expiresAt: Math.floor(Date.now() / 1000) + 60, + }), + headers: { "content-type": "application/json" }, + }); + const firstBody = (await first.json()) as { id: string; updated: boolean }; + expect(firstBody.updated).toBe(false); + + const second = await userControlFetch("/refresh-token-mcp", { + method: "POST", + body: JSON.stringify({ + name: "Refresh Target Updatable", + url: "https://refresh-target-update.example.com/mcp", + tokenEndpoint: TOKEN_ENDPOINT, + clientId: "client-xyz", + accessToken: "access-v2", + refreshToken: "refresh-v2", + expiresAt: Math.floor(Date.now() / 1000) + 3600, + }), + headers: { "content-type": "application/json" }, + }); + const secondBody = (await second.json()) as { id: string; updated: boolean }; + expect(secondBody.updated).toBe(true); + expect(secondBody.id).toBe(firstBody.id); + }); +}); + +describe("refresh-token MCP — access token retrieval", () => { + let originalFetch: typeof fetch; + + beforeEach(async () => { + await initPasskeyIfNeeded(); + originalFetch = globalThis.fetch; + }); + + afterEach(() => { + globalThis.fetch = originalFetch; + }); + + it("returns the cached access token when not yet expired", async () => { + const url = "https://refresh-target-cached.example.com/mcp"; + await userControlFetch("/refresh-token-mcp", { + method: "POST", + body: JSON.stringify({ + name: "Cached Target", + url, + tokenEndpoint: TOKEN_ENDPOINT, + clientId: "client-cached", + accessToken: "still-valid-token", + refreshToken: "still-valid-refresh", + expiresAt: Math.floor(Date.now() / 1000) + 3600, + }), + headers: { "content-type": "application/json" }, + }); + + const list = (await (await userControlFetch("/mcp-configs")).json()) as { + configs: Array<{ id: string; url?: string }>; + }; + const id = list.configs.find((c) => c.url === url)!.id; + + // No fetch should be made — the token is still valid. + const fetchSpy = vi.fn(); + globalThis.fetch = fetchSpy as unknown as typeof fetch; + + const tokenRes = await userControlFetch(`/mcp-configs/${encodeURIComponent(id)}/access-token`); + expect(tokenRes.status).toBe(200); + const body = (await tokenRes.json()) as { accessToken: string }; + expect(body.accessToken).toBe("still-valid-token"); + expect(fetchSpy).not.toHaveBeenCalled(); + }); + + it("refreshes against the token endpoint when the cached token is expired", async () => { + const url = "https://refresh-target-expired.example.com/mcp"; + await userControlFetch("/refresh-token-mcp", { + method: "POST", + body: JSON.stringify({ + name: "Expired Target", + url, + tokenEndpoint: TOKEN_ENDPOINT, + clientId: "client-expired", + accessToken: "expired-token", + refreshToken: "old-refresh", + // Expired an hour ago — guarantees a refresh on next read. + expiresAt: Math.floor(Date.now() / 1000) - 3600, + }), + headers: { "content-type": "application/json" }, + }); + + const list = (await (await userControlFetch("/mcp-configs")).json()) as { + configs: Array<{ id: string; url?: string }>; + }; + const id = list.configs.find((c) => c.url === url)!.id; + + // Mock the refresh response: rotated refresh token + new access token. + // Capture the request for post-call assertions — putting `expect`s + // inside the mock function turns assertion failures into rejected + // promises, which surface as opaque 502s from the DO endpoint. + let captured: { url: string; method?: string; body: string } | null = null; + const fetchSpy = vi.fn(async (input: RequestInfo | URL, init?: RequestInit) => { + const requestUrl = typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url; + let bodyStr = ""; + if (typeof init?.body === "string") bodyStr = init.body; + else if (init?.body instanceof URLSearchParams) bodyStr = init.body.toString(); + captured = { url: requestUrl, method: init?.method, body: bodyStr }; + return new Response( + JSON.stringify({ + access_token: "fresh-token", + refresh_token: "rotated-refresh", + expires_in: 900, + }), + { status: 200, headers: { "content-type": "application/json" } }, + ); + }); + globalThis.fetch = fetchSpy as unknown as typeof fetch; + + const tokenRes = await userControlFetch(`/mcp-configs/${encodeURIComponent(id)}/access-token`); + expect(tokenRes.status).toBe(200); + const body = (await tokenRes.json()) as { accessToken: string }; + expect(body.accessToken).toBe("fresh-token"); + expect(fetchSpy).toHaveBeenCalledTimes(1); + expect(captured?.url).toBe(TOKEN_ENDPOINT); + expect(captured?.method).toBe("POST"); + expect(captured?.body).toContain("grant_type=refresh_token"); + expect(captured?.body).toContain("refresh_token=old-refresh"); + expect(captured?.body).toContain("client_id=client-expired"); + }); + + it("force=1 refreshes even if the cached token has not yet expired", async () => { + const url = "https://refresh-target-force.example.com/mcp"; + await userControlFetch("/refresh-token-mcp", { + method: "POST", + body: JSON.stringify({ + name: "Force Target", + url, + tokenEndpoint: TOKEN_ENDPOINT, + clientId: "client-force", + accessToken: "cached-token", + refreshToken: "cached-refresh", + expiresAt: Math.floor(Date.now() / 1000) + 3600, // still valid + }), + headers: { "content-type": "application/json" }, + }); + + const list = (await (await userControlFetch("/mcp-configs")).json()) as { + configs: Array<{ id: string; url?: string }>; + }; + const id = list.configs.find((c) => c.url === url)!.id; + + const fetchSpy = vi.fn(async () => + new Response( + JSON.stringify({ + access_token: "force-refreshed-token", + refresh_token: "force-refreshed-refresh", + expires_in: 900, + }), + { status: 200, headers: { "content-type": "application/json" } }, + ), + ); + globalThis.fetch = fetchSpy as unknown as typeof fetch; + + const tokenRes = await userControlFetch(`/mcp-configs/${encodeURIComponent(id)}/access-token?force=1`); + expect(tokenRes.status).toBe(200); + const body = (await tokenRes.json()) as { accessToken: string }; + expect(body.accessToken).toBe("force-refreshed-token"); + expect(fetchSpy).toHaveBeenCalledTimes(1); + }); + + it("surfaces a 502 when the token endpoint rejects the refresh", async () => { + const url = "https://refresh-target-broken.example.com/mcp"; + await userControlFetch("/refresh-token-mcp", { + method: "POST", + body: JSON.stringify({ + name: "Broken Target", + url, + tokenEndpoint: TOKEN_ENDPOINT, + clientId: "client-broken", + accessToken: "irrelevant", + refreshToken: "dead-refresh", + expiresAt: Math.floor(Date.now() / 1000) - 1, // already expired + }), + headers: { "content-type": "application/json" }, + }); + + const list = (await (await userControlFetch("/mcp-configs")).json()) as { + configs: Array<{ id: string; url?: string }>; + }; + const id = list.configs.find((c) => c.url === url)!.id; + + globalThis.fetch = (async () => + new Response(JSON.stringify({ error: "invalid_grant" }), { + status: 400, + headers: { "content-type": "application/json" }, + })) as unknown as typeof fetch; + + const tokenRes = await userControlFetch(`/mcp-configs/${encodeURIComponent(id)}/access-token`); + expect(tokenRes.status).toBe(502); + }); +});