diff --git a/packages/mcp/src/client.ts b/packages/mcp/src/client.ts index 8c421a21..092b89ec 100644 --- a/packages/mcp/src/client.ts +++ b/packages/mcp/src/client.ts @@ -167,6 +167,35 @@ function isUuid(value: string): boolean { return typeof value === "string" && UUID_RE.test(value); } +/** + * Collab-token cache TTL in milliseconds (issue #435). Read fresh from the + * environment on every mint — like collab-session.ts readConfig — so tests and a + * live rollback can change it without reloading the module. + * + * Why a cache at all: the live CollabSession registry (#400/#431) keys sessions + * on (wsUrl, pageId, collabToken) for identity isolation (invariant 4). But BOTH + * collab-token sources mint a FRESH token per mutation — the in-app provider + * re-signs a JWT whose iat/exp (seconds) changes every second, and the external + * MCP POSTs /auth/collab-token each call — so the token in the key changed on + * every op and the session was almost never reused (connect-storms, 25s + * timeouts, zombie sessions). Caching the token per-client keeps the key stable + * across a burst of mutations so ONE session is reused. + * + * Default 5 min: well under the 24h collab-token lifetime AND <= the collab + * session max-age (10 min, MCP_COLLAB_SESSION_MAX_AGE_MS), so the + * permission-staleness window is not widened beyond what #431 already accepted. + * The rollback knob is an EXPLICIT 0 (or a negative number): that DISABLES the + * cache — an exact fetch-per-call legacy path, mirroring how idleMs<=0 disables + * the session cache. Unset OR unparseable (e.g. a typo like "5min", "abc") falls + * back to the 5-min default with the cache ON — parseInt yields NaN, which is + * treated as "not configured", not as "disabled". So to turn the cache off you + * must set the value to exactly 0, not to garbage. + */ +function readCollabTokenTtlMs(): number { + const raw = parseInt(process.env.MCP_COLLAB_TOKEN_TTL_MS ?? "", 10); + return Number.isFinite(raw) ? Math.max(0, raw) : 5 * 60 * 1000; +} + export class DocmostClient { private client: AxiosInstance; private token: string | null = null; @@ -205,6 +234,15 @@ export class DocmostClient { // resolvePageId), so only slugId->uuid entries are stored/read here. private pageIdCache = new Map(); + // Collab-token cache (issue #435): the last minted collab token plus the + // wall-clock time it was minted, so a burst of content mutations reuses ONE + // token and therefore ONE live CollabSession (whose registry key includes the + // token — #400 invariant 4). Per-instance: a DocmostClient is built per + // user/per chat request, so a cached token can never leak across identities. + // Reset whenever the client's identity changes (login() / this.token cleared); + // bypassed on a forced refresh (the 401/403 reauth path). null = no token yet. + private collabTokenCache: { token: string; mintedAt: number } | null = null; + // Two construction forms: // - new DocmostClient(config) // discriminated union (current) // - new DocmostClient(baseURL, email, password) // legacy positional creds @@ -273,8 +311,11 @@ export class DocmostClient { if (config && isAuthError && !config._retry && !isLoginRequest) { config._retry = true; - // Drop the stale token + Authorization header before re-login. + // Drop the stale token + Authorization header before re-login. Also + // clear the collab-token cache (#435): a new identity/login must not + // keep serving a collab token minted under the old one. this.token = null; + this.collabTokenCache = null; delete this.client.defaults.headers.common["Authorization"]; try { await this.login(); @@ -323,6 +364,9 @@ export class DocmostClient { throw new Error("getToken returned an empty token"); } this.token = token; + // Identity (re)established: drop any collab token minted under a + // previous identity so the #435 cache can never outlive it. + this.collabTokenCache = null; this.client.defaults.headers.common["Authorization"] = `Bearer ${token}`; }) @@ -345,8 +389,34 @@ export class DocmostClient { * by this.client's response interceptor; this helper replicates that * behaviour for collab-token requests: ensure a token, try once, and on an * expired-token auth error perform a fresh login and retry exactly once. + * + * Collab-token cache (issue #435): both sources — the getCollabToken provider + * (in-app agent) AND the REST /auth/collab-token endpoint (external MCP) — mint + * a FRESH token per call, whose string therefore changes every op. Since the + * live CollabSession registry keys on the token string (#400/#431 invariant 4), + * that churned the key and defeated session reuse. So we cache the last minted + * token per-client for readCollabTokenTtlMs() and hand it back for a burst of + * mutations, keeping the session key stable. `forceRefresh` bypasses the cache + * (the 401/403 reauth retry uses it, so the retry cannot be handed the same + * stale token that just failed — otherwise reauth would be a no-op). TTL 0 + * disables the cache: exact fetch-per-call legacy behaviour. */ - private async getCollabTokenWithReauth(): Promise { + private async getCollabTokenWithReauth( + forceRefresh = false, + ): Promise { + const ttl = readCollabTokenTtlMs(); + // Serve the cached collab token while it is still fresh (identity isolation + // is preserved: the cache is a per-instance field on a client built per + // user/per chat request, and it is cleared on every identity change). + if ( + !forceRefresh && + ttl > 0 && + this.collabTokenCache && + Date.now() - this.collabTokenCache.mintedAt < ttl + ) { + return this.collabTokenCache.token; + } + // Collab-token PROVIDER path: when a getCollabToken provider was supplied // (the internal agent's provenance collab token), use it instead of the // REST /auth/collab-token endpoint. Re-invoke it once on a 401/403 (e.g. the @@ -357,23 +427,13 @@ export class DocmostClient { if (typeof token !== "string" || token.length === 0) { throw new Error("getCollabToken returned an empty token"); } - return token; + return this.rememberCollabToken(token, ttl); } catch (e) { - const axiosStatus = axios.isAxiosError(e) - ? e.response?.status - : undefined; - const attachedStatus = (e as any)?.status; - const isAuthError = - axiosStatus === 401 || - axiosStatus === 403 || - attachedStatus === 401 || - attachedStatus === 403; - if (isAuthError) { - const token = await this.getCollabTokenFn(); - if (typeof token !== "string" || token.length === 0) { - throw new Error("getCollabToken returned an empty token"); - } - return token; + // On an auth error retry EXACTLY once, forcing a refresh so the retry + // re-invokes the provider (bypassing the cache) for a genuinely fresh + // token. `!forceRefresh` bounds it to a single retry (no loop). + if (this.isCollabAuthError(e) && !forceRefresh) { + return this.getCollabTokenWithReauth(true); } throw e; } @@ -381,28 +441,51 @@ export class DocmostClient { await this.ensureAuthenticated(); try { - return await getCollabToken(this.apiUrl, this.token!); + const token = await getCollabToken(this.apiUrl, this.token!); + return this.rememberCollabToken(token, ttl); } catch (e) { // getCollabToken wraps the AxiosError in a plain Error but attaches the - // HTTP status as `.status`, so detect an auth failure via either the raw - // AxiosError shape OR the attached status. - const axiosStatus = axios.isAxiosError(e) - ? e.response?.status - : undefined; - const attachedStatus = (e as any)?.status; - const isAuthError = - axiosStatus === 401 || - axiosStatus === 403 || - attachedStatus === 401 || - attachedStatus === 403; - if (isAuthError) { + // HTTP status as `.status`, so isCollabAuthError detects an auth failure + // via either the raw AxiosError shape OR the attached status. + if (this.isCollabAuthError(e) && !forceRefresh) { + // Fresh login (which clears this.token AND the collab-token cache), then + // retry exactly once with the cache bypassed via forceRefresh. await this.login(); - return await getCollabToken(this.apiUrl, this.token!); + return this.getCollabTokenWithReauth(true); } throw e; } } + /** + * Store a freshly minted collab token in the per-client cache (issue #435) and + * return it unchanged. No-op write when the cache is disabled (ttl<=0) or the + * token is empty, so a disabled cache is exact fetch-per-call legacy behaviour + * and a bad token is never cached. + */ + private rememberCollabToken(token: string, ttl: number): string { + if (ttl > 0 && typeof token === "string" && token.length > 0) { + this.collabTokenCache = { token, mintedAt: Date.now() }; + } + return token; + } + + /** + * True when an error carries a 401/403 — either as a raw AxiosError + * (`error.response.status`) or as the plain-Error `.status` that + * lib/auth-utils.getCollabToken attaches after wrapping the AxiosError. + */ + private isCollabAuthError(e: unknown): boolean { + const axiosStatus = axios.isAxiosError(e) ? e.response?.status : undefined; + const attachedStatus = (e as any)?.status; + return ( + axiosStatus === 401 || + axiosStatus === 403 || + attachedStatus === 401 || + attachedStatus === 403 + ); + } + /** * Connect to the collaboration websocket, read the live doc, apply * `transform`, write the result, and wait for the server to persist it — diff --git a/packages/mcp/test/mock/collab-token-cache.test.mjs b/packages/mcp/test/mock/collab-token-cache.test.mjs new file mode 100644 index 00000000..366a48a0 --- /dev/null +++ b/packages/mcp/test/mock/collab-token-cache.test.mjs @@ -0,0 +1,282 @@ +// Unit tests for the collab-token cache (issue #435). The live CollabSession +// registry (#400/#431) keys sessions on (wsUrl, pageId, collabToken), so a token +// string that changes every op defeats reuse. This cache holds the last minted +// token per DocmostClient for MCP_COLLAB_TOKEN_TTL_MS so a burst of mutations +// reuses ONE token -> ONE session. These tests exercise both mint sources: +// - the getCollabToken PROVIDER path (in-app agent), via a counting provider fn; +// - the REST /auth/collab-token path (external MCP), via a mock http server. +// getCollabTokenWithReauth is private in TS but a plain method on the compiled +// build, so the tests call it directly (same convention as reauth.test.mjs). +import { test, afterEach, after } from "node:test"; +import assert from "node:assert/strict"; +import http from "node:http"; +import { DocmostClient } from "../../build/client.js"; + +// Restore the env knob after each test so cases do not leak into one another. +const ENV_KEY = "MCP_COLLAB_TOKEN_TTL_MS"; +afterEach(() => { + delete process.env[ENV_KEY]; +}); + +// --------------------------------------------------------------------------- +// Small mock server for the REST /auth/collab-token path. Counts collab-token +// mints and can be told to 401 the first N of them (to drive the reauth retry). +// --------------------------------------------------------------------------- +function readBody(req) { + return new Promise((resolve) => { + let raw = ""; + req.on("data", (c) => (raw += c)); + req.on("end", () => resolve(raw)); + }); +} +function sendJson(res, status, obj, extra = {}) { + res.writeHead(status, { "Content-Type": "application/json", ...extra }); + res.end(JSON.stringify(obj)); +} + +const openServers = []; +after(async () => { + await Promise.all( + openServers.map((s) => new Promise((r) => s.close(r))), + ); +}); + +// state: { collabCalls, loginCalls, unauthorizedCollabHits } +function spawnCollabServer(state, { collabAuthFailsFor = 0 } = {}) { + return new Promise((resolve) => { + const server = http.createServer(async (req, res) => { + await readBody(req); + if (req.url === "/api/auth/login") { + state.loginCalls++; + // A fresh authToken per login so an identity change is observable. + sendJson(res, 200, { success: true }, { + "Set-Cookie": `authToken=login-${state.loginCalls}; Path=/; HttpOnly`, + }); + return; + } + if (req.url === "/api/auth/collab-token") { + state.collabCalls++; + if (state.collabCalls <= collabAuthFailsFor) { + sendJson(res, 401, { message: "Unauthorized" }); + return; + } + // Unique token per mint so a stale cached value is distinguishable. + sendJson(res, 200, { data: { token: `collab-${state.collabCalls}` } }); + return; + } + sendJson(res, 404, { message: "not found" }); + }); + server.listen(0, "127.0.0.1", () => { + openServers.push(server); + resolve(`http://127.0.0.1:${server.address().port}/api`); + }); + }); +} + +// =========================================================================== +// PROVIDER path (in-app agent getCollabToken fn) +// =========================================================================== + +// A counting provider that returns a distinct token each call so a cached +// (reused) token is visibly the SAME string while a fresh mint is different. +function countingProvider() { + let n = 0; + const fn = async () => { + n++; + return `provider-token-${n}`; + }; + return { + fn, + get calls() { + return n; + }, + }; +} + +test("within TTL, repeated calls return the SAME token and mint ONCE (provider path)", async () => { + process.env[ENV_KEY] = "300000"; // 5 min + const p = countingProvider(); + const client = new DocmostClient({ + apiUrl: "http://127.0.0.1:1/api", + getToken: async () => "access", + getCollabToken: p.fn, + }); + + const a = await client.getCollabTokenWithReauth(); + const b = await client.getCollabTokenWithReauth(); + const c = await client.getCollabTokenWithReauth(); + + assert.equal(a, "provider-token-1"); + assert.equal(b, a, "second call reuses the cached token"); + assert.equal(c, a, "third call reuses the cached token"); + assert.equal(p.calls, 1, "the provider is invoked exactly once within the TTL"); +}); + +test("after TTL expiry a new token is minted (provider path)", async () => { + process.env[ENV_KEY] = "20"; // 20ms TTL + const p = countingProvider(); + const client = new DocmostClient({ + apiUrl: "http://127.0.0.1:1/api", + getToken: async () => "access", + getCollabToken: p.fn, + }); + + const a = await client.getCollabTokenWithReauth(); + await new Promise((r) => setTimeout(r, 40)); // let the TTL lapse + const b = await client.getCollabTokenWithReauth(); + + assert.equal(a, "provider-token-1"); + assert.equal(b, "provider-token-2", "a fresh token is minted after expiry"); + assert.equal(p.calls, 2); +}); + +test("MCP_COLLAB_TOKEN_TTL_MS=0 disables the cache: mint on EVERY call (provider path)", async () => { + process.env[ENV_KEY] = "0"; + const p = countingProvider(); + const client = new DocmostClient({ + apiUrl: "http://127.0.0.1:1/api", + getToken: async () => "access", + getCollabToken: p.fn, + }); + + await client.getCollabTokenWithReauth(); + await client.getCollabTokenWithReauth(); + await client.getCollabTokenWithReauth(); + + assert.equal(p.calls, 3, "cache disabled -> exact fetch-per-call legacy path"); +}); + +test("a 401 triggers the internal reauth retry, which bypasses the cache and mints fresh (provider path)", async () => { + process.env[ENV_KEY] = "300000"; + let n = 0; + const provider = async () => { + n++; + if (n === 1) { + // The FIRST mint fails with an auth error; the internal reauth retry must + // re-invoke the provider (bypassing the empty cache) for a fresh token. + const err = new Error("collab token expired"); + err.status = 401; + throw err; + } + return `provider-token-${n}`; + }; + const client = new DocmostClient({ + apiUrl: "http://127.0.0.1:1/api", + getToken: async () => "access", + getCollabToken: provider, + }); + + // Cache is empty: mint #1 401s -> the reauth retry mints #2 and caches it. + const tok = await client.getCollabTokenWithReauth(); + assert.equal(tok, "provider-token-2", "the post-401 retry token wins"); + assert.equal(n, 2, "exactly one failed mint + one retry, no loop"); + + // The retried token is what got cached (no extra mint on a cache hit). + const cached = await client.getCollabTokenWithReauth(); + assert.equal(cached, "provider-token-2"); + assert.equal(n, 2, "served from cache, provider not re-invoked"); +}); + +test("forceRefresh=true bypasses a warm cache and mints a fresh token (provider path)", async () => { + process.env[ENV_KEY] = "300000"; + const p = countingProvider(); + const client = new DocmostClient({ + apiUrl: "http://127.0.0.1:1/api", + getToken: async () => "access", + getCollabToken: p.fn, + }); + + const first = await client.getCollabTokenWithReauth(); // caches token-1 + assert.equal(first, "provider-token-1"); + + // A forced refresh (what the reauth path passes) must NOT return the cached + // token-1; it mints a fresh token-2 and replaces the cache. + const forced = await client.getCollabTokenWithReauth(true); + assert.equal(forced, "provider-token-2", "cache bypassed on forceRefresh"); + assert.equal(p.calls, 2); + + const cached = await client.getCollabTokenWithReauth(); + assert.equal(cached, "provider-token-2", "the fresh token replaced the cache"); + assert.equal(p.calls, 2); +}); + +test("two consecutive mutations keep the SAME token, so the session key is stable (provider path)", async () => { + // The whole point of #435: acquireCollabSession keys on the token, so two + // acquire calls in a burst must be handed the identical token string. + process.env[ENV_KEY] = "300000"; + const p = countingProvider(); + const client = new DocmostClient({ + apiUrl: "http://127.0.0.1:1/api", + getToken: async () => "access", + getCollabToken: p.fn, + }); + + const t1 = await client.getCollabTokenWithReauth(); + const t2 = await client.getCollabTokenWithReauth(); + assert.equal(t1, t2, "identical token across two mutations -> one session key"); + assert.equal(p.calls, 1); +}); + +// =========================================================================== +// REST /auth/collab-token path (external MCP) +// =========================================================================== + +test("within TTL, the REST /auth/collab-token endpoint is hit ONCE", async () => { + process.env[ENV_KEY] = "300000"; + const state = { collabCalls: 0, loginCalls: 0 }; + const baseURL = await spawnCollabServer(state); + const client = new DocmostClient(baseURL, "user@example.com", "pw"); + + const a = await client.getCollabTokenWithReauth(); + const b = await client.getCollabTokenWithReauth(); + + assert.equal(a, "collab-1"); + assert.equal(b, a, "cached token reused"); + assert.equal(state.collabCalls, 1, "POST /auth/collab-token called once"); +}); + +test("TTL=0 hits the REST endpoint on every call", async () => { + process.env[ENV_KEY] = "0"; + const state = { collabCalls: 0, loginCalls: 0 }; + const baseURL = await spawnCollabServer(state); + const client = new DocmostClient(baseURL, "user@example.com", "pw"); + + await client.getCollabTokenWithReauth(); + await client.getCollabTokenWithReauth(); + + assert.equal(state.collabCalls, 2, "cache disabled -> fetch each call"); +}); + +test("401 on REST collab-token re-logs-in and refetches (cache bypassed)", async () => { + process.env[ENV_KEY] = "300000"; + const state = { collabCalls: 0, loginCalls: 0 }; + // The first collab-token mint 401s; the reauth path logs in and retries. + const baseURL = await spawnCollabServer(state, { collabAuthFailsFor: 1 }); + const client = new DocmostClient(baseURL, "user@example.com", "pw"); + // Pre-seed a token so the initial call does not perform an initial login. + client.token = "seed"; + client.client.defaults.headers.common["Authorization"] = "Bearer seed"; + + const tok = await client.getCollabTokenWithReauth(); + assert.equal(tok, "collab-2", "the post-reauth mint wins, not the failed one"); + assert.equal(state.loginCalls, 1, "re-login happened exactly once"); + assert.equal(state.collabCalls, 2, "one failed mint + one successful retry"); +}); + +test("a fresh login clears the cache so a collab token cannot outlive the identity", async () => { + process.env[ENV_KEY] = "300000"; + const state = { collabCalls: 0, loginCalls: 0 }; + const baseURL = await spawnCollabServer(state); + const client = new DocmostClient(baseURL, "user@example.com", "pw"); + + const before = await client.getCollabTokenWithReauth(); + assert.equal(before, "collab-1"); + + // Simulate an identity change (the 401 interceptor / re-login path calls + // login(), which must drop the cached collab token). + await client.login(); + + const after = await client.getCollabTokenWithReauth(); + assert.equal(after, "collab-2", "cache was invalidated by login(); refetched"); + assert.equal(state.collabCalls, 2); +});