diff --git a/packages/mcp/src/client/context.ts b/packages/mcp/src/client/context.ts index 498c3b79..b65a7b5a 100644 --- a/packages/mcp/src/client/context.ts +++ b/packages/mcp/src/client/context.ts @@ -19,7 +19,10 @@ import { assertYjsEncodable, MutationResult, } from "../lib/collaboration.js"; -import { acquireCollabSession } from "../lib/collab-session.js"; +import { + acquireCollabSession, + isCollabAuthFailedError, +} from "../lib/collab-session.js"; import { withPageLock, isUuid } from "../lib/page-lock.js"; import { getCollabToken, performLogin } from "../lib/auth-utils.js"; import { formatDocmostAxiosError } from "./errors.js"; @@ -484,6 +487,37 @@ export abstract class DocmostClientContext { ); } + /** + * Run a collab write and, on a Hocuspocus HANDSHAKE auth failure, self-heal + * once (#486). Symmetric to the HTTP-401 path in getCollabTokenWithReauth: the + * REST interceptor and login() already drop the cached collab token on a 401/ + * 403, but a rejected WEBSOCKET handshake left the stale token in the cache, so + * every subsequent mutation kept re-presenting the same bad token for up to the + * collab-token TTL (minutes) with no self-heal. Here, when the write rejects + * with the tagged collab-auth error, we invalidate the cached token and retry + * the write EXACTLY once with a force-refreshed token. Not a loop: a second + * failure (or any non-auth error) propagates unchanged. + * + * `write` receives the token to use, so the retry can hand it a genuinely fresh + * one rather than re-running with the same stale string. + */ + protected async writeWithCollabAuthRetry( + collabToken: string, + write: (token: string) => Promise, + ): Promise { + try { + return await write(collabToken); + } catch (e) { + if (!isCollabAuthFailedError(e)) throw e; + // The WS handshake rejected our token: drop it from the cache so it can't + // be reused for the rest of the TTL, mint a fresh one (forceRefresh bypasses + // the cache and re-invokes the provider/login), and retry the write once. + this.collabTokenCache = null; + const fresh = await this.getCollabTokenWithReauth(true); + return await write(fresh); + } + } + /** * Connect to the collaboration websocket, read the live doc, apply * `transform`, write the result, and wait for the server to persist it — @@ -518,19 +552,24 @@ export abstract class DocmostClientContext { // unsyncedChanges/connectionLost ack logic live in CollabSession.mutate, // preserved verbatim from the old inline machine (incl. the #152 structural // diff that keeps a live editor's cursor anchored). - const session = await acquireCollabSession(pageId, collabToken, this.apiUrl, { - // Only the actual 25s collab connect timeout emits this — the connect-vs- - // unload signal; the other failure paths must NOT emit it. - onConnectTimeout: () => - this.onMetricFn?.("collab_connect_timeouts_total", 1), + // Wrap in the collab-auth self-heal (#486): a rejected WS handshake drops the + // cached collab token and retries once with a fresh one (the retry passes the + // refreshed token down to acquireCollabSession via `token`). + return this.writeWithCollabAuthRetry(collabToken, async (token) => { + const session = await acquireCollabSession(pageId, token, this.apiUrl, { + // Only the actual 25s collab connect timeout emits this — the connect-vs- + // unload signal; the other failure paths must NOT emit it. + onConnectTimeout: () => + this.onMetricFn?.("collab_connect_timeouts_total", 1), + }); + try { + return await session.mutate(transform); + } catch (e) { + // Drop the session on any failure so the next call reconnects fresh. + session.destroy("mutate failed"); + throw e; + } }); - try { - return await session.mutate(transform); - } catch (e) { - // Drop the session on any failure so the next call reconnects fresh. - session.destroy("mutate failed"); - throw e; - } } /** @@ -659,7 +698,11 @@ export abstract class DocmostClientContext { transform: (doc: any) => any, ): Promise<{ doc?: any; verify?: any }> { const pageUuid = await this.resolvePageId(pageId); - return mutatePageContent(pageUuid, collabToken, apiUrl, transform); + // #486: on a rejected collab-WS handshake, invalidate + refresh the token and + // retry the write once (symmetric to the HTTP-401 reauth path). + return this.writeWithCollabAuthRetry(collabToken, (token) => + mutatePageContent(pageUuid, token, apiUrl, transform), + ); } /** @@ -679,7 +722,11 @@ export abstract class DocmostClientContext { apiUrl: string, ): Promise<{ doc?: any; verify?: any }> { const pageUuid = await this.resolvePageId(pageId); - return replacePageContent(pageUuid, doc, collabToken, apiUrl); + // #486: on a rejected collab-WS handshake, invalidate + refresh the token and + // retry the write once (symmetric to the HTTP-401 reauth path). + return this.writeWithCollabAuthRetry(collabToken, (token) => + replacePageContent(pageUuid, doc, token, apiUrl), + ); } /** diff --git a/packages/mcp/src/lib/collab-session.ts b/packages/mcp/src/lib/collab-session.ts index 3dc16161..5e1e40cc 100644 --- a/packages/mcp/src/lib/collab-session.ts +++ b/packages/mcp/src/lib/collab-session.ts @@ -37,6 +37,25 @@ const CONNECT_TIMEOUT_MS = 25000; /** Time we wait for the server to acknowledge our write before giving up. */ const PERSIST_TIMEOUT_MS = 20000; +/** + * Marker property set on the Error thrown when the Hocuspocus handshake REJECTS + * our collab token (onAuthenticationFailed). The client wraps content writes so + * that on this specific failure it invalidates its cached collab token and + * retries once with a fresh one — symmetric to the HTTP-401 reauth path (#486). + * A plain message-match would be brittle; a tagged property is unambiguous and + * survives teardown (which rejects pending ops with this SAME error object). + */ +const COLLAB_AUTH_FAILED_MARKER = "collabAuthFailed"; + +/** True when `e` is the tagged collab-WS auth-failure error (see marker above). */ +export function isCollabAuthFailedError(e: unknown): boolean { + return !!( + e && + typeof e === "object" && + (e as Record)[COLLAB_AUTH_FAILED_MARKER] === true + ); +} + /** * Tunables, read fresh from the environment on every acquire so tests (and a * live rollback) can change them without reloading the module. Mirrors how @@ -302,10 +321,13 @@ export class CollabSession { this.openResolve?.(); }, onAuthenticationFailed: () => { - this.teardown( - new Error("Authentication failed for collaboration connection"), - true, - ); + // Tag the error so the client can tell a REJECTED collab token apart + // from a generic disconnect and invalidate + refresh it (#486). + const err = new Error( + "Authentication failed for collaboration connection", + ) as Error & { [COLLAB_AUTH_FAILED_MARKER]?: boolean }; + err[COLLAB_AUTH_FAILED_MARKER] = true; + this.teardown(err, true); }, }); }); diff --git a/packages/mcp/test/mock/collab-auth-failure-reset.test.mjs b/packages/mcp/test/mock/collab-auth-failure-reset.test.mjs new file mode 100644 index 00000000..573098ce --- /dev/null +++ b/packages/mcp/test/mock/collab-auth-failure-reset.test.mjs @@ -0,0 +1,137 @@ +// Unit tests for the collab-token reset on a Hocuspocus WS auth failure (#486). +// +// Before this fix the cached collab token (#435) was dropped ONLY on an HTTP +// 401/403 (the REST interceptor + login()); a rejected collab-WEBSOCKET handshake +// left the stale token in the cache, so every subsequent mutation re-presented +// the SAME bad token for up to the collab-token TTL (minutes) with no self-heal. +// +// The fix wraps collab writes in `writeWithCollabAuthRetry`: when the write +// rejects with the tagged collab-auth error (collab-session.ts's +// onAuthenticationFailed), it invalidates the cached token and retries the write +// ONCE with a force-refreshed token — symmetric to the HTTP-401 path. +// +// writeWithCollabAuthRetry / getCollabTokenWithReauth are protected in TS but +// plain methods on the compiled build, so the tests call them directly (same +// convention as collab-token-cache.test.mjs). +import { test, afterEach } from "node:test"; +import assert from "node:assert/strict"; +import { DocmostClient } from "../../build/client.js"; + +const ENV_KEY = "MCP_COLLAB_TOKEN_TTL_MS"; +afterEach(() => { + delete process.env[ENV_KEY]; +}); + +// 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; + }, + }; +} + +// The tagged error collab-session.ts throws on a rejected WS handshake. +function collabAuthError() { + const err = new Error("Authentication failed for collaboration connection"); + err.collabAuthFailed = true; + return err; +} + +test("a WS auth failure clears the cached token and retries the write with a FRESH one (#486)", async () => { + process.env[ENV_KEY] = "300000"; // 5 min: the cache is warm across the burst. + const p = countingProvider(); + const client = new DocmostClient({ + apiUrl: "http://127.0.0.1:1/api", + getToken: async () => "access", + getCollabToken: p.fn, + }); + + // Warm the cache the way a real write would (mints provider-token-1). + const initial = await client.getCollabTokenWithReauth(); + assert.equal(initial, "provider-token-1"); + assert.equal(p.calls, 1); + + const tokensSeen = []; + const write = async (token) => { + tokensSeen.push(token); + // The FIRST attempt (with the stale cached token) fails the WS handshake; + // the retry (with a fresh token) succeeds. + if (tokensSeen.length === 1) throw collabAuthError(); + return `written-with:${token}`; + }; + + const result = await client.writeWithCollabAuthRetry(initial, write); + + assert.equal(tokensSeen.length, 2, "write attempted exactly twice (one retry)"); + assert.equal(tokensSeen[0], "provider-token-1", "first attempt used the stale token"); + assert.equal( + tokensSeen[1], + "provider-token-2", + "retry used a FRESH force-refreshed token, not the stale cached one", + ); + assert.equal(result, "written-with:provider-token-2", "the retry's result wins"); + assert.equal(p.calls, 2, "exactly one extra mint for the retry — no loop"); + + // The cache now holds the fresh token, so a subsequent op reuses it (proving + // the stale token was evicted and the fresh one cached, not re-minted). + const next = await client.getCollabTokenWithReauth(); + assert.equal(next, "provider-token-2", "the fresh token replaced the stale cache"); + assert.equal(p.calls, 2, "served from cache — provider not re-invoked"); +}); + +test("a successful write is NOT retried and mints nothing extra", 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 initial = await client.getCollabTokenWithReauth(); // provider-token-1 + let attempts = 0; + const result = await client.writeWithCollabAuthRetry(initial, async (token) => { + attempts++; + return `ok:${token}`; + }); + + assert.equal(attempts, 1, "no retry on success"); + assert.equal(result, "ok:provider-token-1"); + assert.equal(p.calls, 1, "no extra mint"); +}); + +test("a NON-auth write error propagates unchanged (no reset, no retry)", 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 initial = await client.getCollabTokenWithReauth(); // provider-token-1 + let attempts = 0; + await assert.rejects( + client.writeWithCollabAuthRetry(initial, async () => { + attempts++; + throw new Error("collab connection closed before persist"); // NOT tagged. + }), + /closed before persist/, + ); + + assert.equal(attempts, 1, "a non-auth error is not retried"); + assert.equal(p.calls, 1, "the cache is untouched -> no fresh mint"); + + // Cache still holds the original token (was never invalidated). + const still = await client.getCollabTokenWithReauth(); + assert.equal(still, "provider-token-1"); + assert.equal(p.calls, 1); +});