diff --git a/packages/mcp/src/client.ts b/packages/mcp/src/client.ts index d370a2b6..84541ff7 100644 --- a/packages/mcp/src/client.ts +++ b/packages/mcp/src/client.ts @@ -8,10 +8,6 @@ import { filterComment, filterSearchResult, } from "./lib/filters.js"; -import { HocuspocusProvider } from "@hocuspocus/provider"; -import { TiptapTransformer } from "@hocuspocus/transformer"; -import * as Y from "yjs"; -import WebSocket from "ws"; import { convertProseMirrorToMarkdown } from "./lib/markdown-converter.js"; import { collectInternalFileNodes, @@ -24,11 +20,10 @@ import { markdownToProseMirror, markdownToProseMirrorCanonical, mutatePageContent, - buildCollabWsUrl, assertYjsEncodable, - applyDocToFragment, MutationResult, } from "./lib/collaboration.js"; +import { acquireCollabSession } from "./lib/collab-session.js"; import { footnoteWarningsField } from "./lib/footnote-analyze.js"; import { buildPageTree } from "./lib/tree.js"; import { @@ -417,177 +412,32 @@ export class DocmostClient { * change report. The report is computed AFTER the atomic read->write and * never throws. */ - private mutateLiveContentUnlocked( + private async mutateLiveContentUnlocked( pageId: string, collabToken: string, transform: (liveDoc: any) => any | null, ): Promise { - const CONNECT_TIMEOUT_MS = 25000; - const PERSIST_TIMEOUT_MS = 20000; - const ydoc = new Y.Doc(); - const wsUrl = buildCollabWsUrl(this.apiUrl); - - return new Promise((resolve, reject) => { - let provider: HocuspocusProvider | undefined; - let applied = false; // onSynced may fire again on reconnect — apply once. - let settled = false; - let connectionLost = false; - let connectTimer: ReturnType | undefined; - let persistTimer: ReturnType | undefined; - let unsyncedHandler: ((data: { number: number }) => void) | undefined; - // The verifiable result resolved on every success/abort path. Set on abort - // (no-op report) and after a real write (computed change report). - let mutationResult: MutationResult; - - const cleanup = () => { - if (connectTimer) clearTimeout(connectTimer); - if (persistTimer) clearTimeout(persistTimer); - if (provider) { - if (unsyncedHandler) { - try { - provider.off("unsyncedChanges", unsyncedHandler); - } catch (err) {} - } - try { - provider.destroy(); - } catch (err) {} - } - }; - - const finish = (err: Error | null, value?: MutationResult) => { - if (settled) return; - settled = true; - cleanup(); - if (err) reject(err); - else resolve(value as MutationResult); - }; - - connectTimer = setTimeout(() => { - // Only the actual 25s collab connect timeout fires here — the agent's - // collab connection to the server never became ready. This is the - // connect-vs-unload signal; the other finish() paths must NOT emit it. - this.onMetricFn?.("collab_connect_timeouts_total", 1); - finish(new Error("Connection timeout to collaboration server")); - }, CONNECT_TIMEOUT_MS); - - const waitForPersistence = () => { - if (settled) return; - if (!provider) { - finish(new Error("collab provider gone before persistence")); - return; - } - if (provider.unsyncedChanges === 0) { - finish(null, mutationResult); - return; - } - persistTimer = setTimeout(() => { - finish( - new Error( - "Timeout waiting for collaboration server to persist the update", - ), - ); - }, PERSIST_TIMEOUT_MS); - unsyncedHandler = (data: { number: number }) => { - if (data.number === 0 && !connectionLost) { - finish(null, mutationResult); - } - }; - provider.on("unsyncedChanges", unsyncedHandler); - }; - - provider = new HocuspocusProvider({ - url: wsUrl, - name: `page.${pageId}`, - document: ydoc, - token: collabToken, - // @ts-ignore - Required for Node.js environment - WebSocketPolyfill: WebSocket, - onDisconnect: () => { - connectionLost = true; - finish( - new Error( - "Collaboration connection closed before the update was persisted/synced", - ), - ); - }, - onClose: () => { - connectionLost = true; - finish( - new Error( - "Collaboration connection closed before the update was persisted/synced", - ), - ); - }, - onSynced: () => { - if (applied || settled) return; - applied = true; - - // CRITICAL: keep everything between reading and writing the live doc - // synchronous (no await) so no remote update can interleave. - let newDoc: any; - let beforeDoc: any; - try { - let liveDoc = TiptapTransformer.fromYdoc(ydoc, "default"); - if ( - !liveDoc || - typeof liveDoc !== "object" || - !Array.isArray(liveDoc.content) - ) { - liveDoc = { type: "doc", content: [] }; - } - - // Snapshot the before-doc for the change report (safe deep clone). - beforeDoc = JSON.parse(JSON.stringify(liveDoc)); - - newDoc = transform(liveDoc); - - if (newDoc == null) { - // Transform aborted — write nothing, return the live doc with a - // no-op change report. - mutationResult = { - doc: liveDoc, - verify: { - changed: false, - textInserted: 0, - textDeleted: 0, - blocksChanged: 0, - marks: {}, - summary: "no changes (transform aborted)", - }, - }; - finish(null, mutationResult); - return; - } - - // Structural diff into the live fragment (issue #152), mirroring - // the main write path: preserves the Yjs ids of unchanged nodes so - // an open editor's cursor is not yanked to the end of the document. - // The previous destructive rewrite (delete-all + applyUpdate of a - // fresh Y.Doc) discarded every node id, so replaceImage — the only - // caller of this method — still reproduced the #152 cursor jump - // (#164). applyDocToFragment runs its own atomic `transact`. - applyDocToFragment(ydoc, newDoc); - } catch (e) { - finish(e instanceof Error ? e : new Error(String(e))); - return; - } - - // Compute the verifiable change report AFTER the transact write: it - // only needs the JSON before/after, so it cannot affect the atomic - // read->write window, and summarizeChange never throws. - mutationResult = { - doc: newDoc, - verify: summarizeChange(beforeDoc, newDoc), - }; - waitForPersistence(); - }, - onAuthenticationFailed: () => { - finish( - new Error("Authentication failed for collaboration connection"), - ); - }, - }); + // Reuse a live CollabSession for the page (issue #400) instead of opening a + // fresh provider per op. acquireCollabSession does NOT take the per-page + // lock — the caller (replaceImage) already holds ONE withPageLock across its + // scan -> upload -> write sequence, and the mutex is not reentrant, so + // taking it here would deadlock. The synchronous read->write section and the + // 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), }); + 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; + } } /** diff --git a/packages/mcp/src/index.ts b/packages/mcp/src/index.ts index 45bb397a..17cabe18 100644 --- a/packages/mcp/src/index.ts +++ b/packages/mcp/src/index.ts @@ -13,6 +13,11 @@ import { SHARED_TOOL_SPECS, SharedToolSpec } from "./tool-specs.js"; export { DocmostClient } from "./client.js"; export type { DocmostMcpConfig } from "./client.js"; +// Teardown for the live per-page CollabSession cache (issue #400). An embedding +// HTTP host (the gitmost NestJS server) should call this from its own shutdown +// hook so no cached collab provider outlives the process. +export { destroyAllSessions } from "./lib/collab-session.js"; + // Re-export the zod-agnostic shared tool-spec registry so the in-app AI-SDK // service can read it off the loaded module (it cannot import the ESM package's // internals directly; it goes through loadDocmostMcp()). diff --git a/packages/mcp/src/lib/collab-session.ts b/packages/mcp/src/lib/collab-session.ts new file mode 100644 index 00000000..1c1c8d81 --- /dev/null +++ b/packages/mcp/src/lib/collab-session.ts @@ -0,0 +1,668 @@ +import { HocuspocusProvider } from "@hocuspocus/provider"; +import { TiptapTransformer } from "@hocuspocus/transformer"; +import * as Y from "yjs"; +import WebSocket from "ws"; +import { + buildCollabWsUrl, + applyDocToFragment, + MutationResult, +} from "./collaboration.js"; +import { summarizeChange } from "./diff.js"; + +/** + * Live per-page collaboration session cache (issue #400). + * + * The one-shot write path (collaboration.mutatePageContent / + * client.mutateLiveContentUnlocked) used to open a NEW HocuspocusProvider, run + * the full connect -> auth -> onLoadDocument -> initial-sync handshake, apply a + * single edit, wait for persistence, and then `provider.destroy()` — for EVERY + * content mutation. Disconnecting after every edit means that once the pause + * between calls exceeds the server's write debounce, the server does a full + * store -> unload -> reload per cell, causing 25s connect timeouts and + * event-loop lag under a burst of edits on one page. + * + * This module keeps ONE live provider + ydoc per (wsUrl, pageId, token) alive + * across a SERIES of edits. While the provider stays connected the server never + * enters store -> unload -> reload, its debounce coalesces N writes into 1-2 + * stores, and the repeated auth/load/initial-sync disappears. + * + * The synchronous read -> transform -> write section and the per-edit + * persistence-ack logic are preserved VERBATIM from the one-shot machine — the + * only change is that they run on a persistent provider instead of a throwaway + * one. See CollabSession.mutate. + */ + +/** Time we wait for the initial handshake/sync before giving up. */ +const CONNECT_TIMEOUT_MS = 25000; +/** Time we wait for the server to acknowledge our write before giving up. */ +const PERSIST_TIMEOUT_MS = 20000; + +/** + * Tunables, read fresh from the environment on every acquire so tests (and a + * live rollback) can change them without reloading the module. Mirrors how + * http.ts parses MCP_SESSION_IDLE_MS. + * - MCP_COLLAB_SESSION_IDLE_MS: idle TTL, reset after every op. Default 60s. + * 0 (or negative) DISABLES the cache — every op opens its own provider and + * destroys it after the op, i.e. the exact legacy per-op-provider behavior + * (the rollback path). + * - MCP_COLLAB_SESSION_MAX_AGE_MS: hard lifetime checked at acquire; bounds + * the permission-staleness window. Default 10 min. + * - MCP_COLLAB_SESSION_MAX_ENTRIES: registry cap; the least-recently-used + * session is destroy-evicted when the cap is reached. Default 32. + */ +interface SessionConfig { + idleMs: number; + maxAgeMs: number; + maxEntries: number; +} + +function parseEnvInt(value: string | undefined, fallback: number): number { + const parsed = parseInt(value ?? "", 10); + return Number.isFinite(parsed) ? parsed : fallback; +} + +function readConfig(): SessionConfig { + // idleMs: allow 0 (disable). A malformed value falls back to the default. + const idleRaw = parseInt(process.env.MCP_COLLAB_SESSION_IDLE_MS ?? "", 10); + const idleMs = Number.isFinite(idleRaw) ? Math.max(0, idleRaw) : 60 * 1000; + const maxAgeMs = Math.max( + 0, + parseEnvInt(process.env.MCP_COLLAB_SESSION_MAX_AGE_MS, 10 * 60 * 1000), + ); + const maxEntriesRaw = parseEnvInt( + process.env.MCP_COLLAB_SESSION_MAX_ENTRIES, + 32, + ); + const maxEntries = maxEntriesRaw > 0 ? maxEntriesRaw : 32; + return { idleMs, maxAgeMs, maxEntries }; +} + +/** + * The subset of HocuspocusProvider this module depends on, so the provider can + * be replaced with a fake in unit tests (there is no server in the test env). + */ +export interface CollabProviderLike { + synced: boolean; + unsyncedChanges: number; + destroy(): void; + on(event: "unsyncedChanges", handler: (data: { number: number }) => void): void; + off(event: "unsyncedChanges", handler: (data: { number: number }) => void): void; +} + +/** The configuration object passed to the provider factory. */ +export interface CollabProviderConfig { + url: string; + name: string; + document: Y.Doc; + token: string; + WebSocketPolyfill: unknown; + onConnect: () => void; + onSynced: () => void; + onDisconnect: () => void; + onClose: () => void; + onAuthenticationFailed: () => void; +} + +export type CollabProviderFactory = ( + config: CollabProviderConfig, +) => CollabProviderLike; + +const defaultProviderFactory: CollabProviderFactory = (config) => + // @ts-ignore - WebSocketPolyfill is required for the Node.js environment. + new HocuspocusProvider(config) as unknown as CollabProviderLike; + +let providerFactory: CollabProviderFactory = defaultProviderFactory; + +/** + * TEST SEAM: swap the provider factory (pass null to restore the real one). + * Not part of the public API — used only by the unit tests, which cannot reach + * a real collaboration server. + */ +export function __setCollabProviderFactory( + factory: CollabProviderFactory | null, +): void { + providerFactory = factory ?? defaultProviderFactory; +} + +/** Optional per-acquire hooks (metrics), passed through from the call site. */ +export interface AcquireOptions { + /** Invoked when the initial connect handshake times out (CONNECT_TIMEOUT_MS). */ + onConnectTimeout?: () => void; +} + +type SessionState = "connecting" | "ready" | "dead"; + +/** + * One live provider + ydoc for a single (wsUrl, pageId, token) triple. + * + * Lifecycle: connecting -> ready -> dead. A session becomes `dead` on the first + * disconnect/close/auth-failure at ANY time, on an idle/eviction/max-age + * teardown, or on an explicit destroy(); death is terminal and removes the + * session from the registry so the next acquire opens a fresh one. We never use + * the provider's auto-reconnect — destroying on the first disconnect closes the + * "reconnect drove unsyncedChanges to 0 without retransmitting our write" class + * of false success. + */ +export class CollabSession { + readonly key: string; + readonly pageId: string; + readonly wsUrl: string; + readonly token: string; + readonly createdAt: number; + state: SessionState = "connecting"; + /** + * Set true on disconnect/close/auth-failure so a reconnect-driven + * unsyncedChanges->0 cannot be mistaken for a successful persist of our + * write (preserved verbatim from the one-shot machine). + */ + connectionLost = false; + + provider: CollabProviderLike | undefined; + private readonly ydoc: Y.Doc; + private readonly cfg: SessionConfig; + /** + * Ephemeral sessions (cache disabled, MCP_COLLAB_SESSION_IDLE_MS<=0) are never + * registered and self-destroy after their single op — the legacy + * provider-per-op behavior. + */ + private readonly ephemeral: boolean; + private readonly opts: AcquireOptions | undefined; + + private dead = false; + private connectTimer: ReturnType | undefined; + private idleTimer: ReturnType | undefined; + private openPromise: Promise | undefined; + private openResolve: (() => void) | undefined; + private openReject: ((err: Error) => void) | undefined; + private openSettled = false; + /** + * The rejector of the CURRENT in-flight mutate, if any. A disconnect/close/ + * auth-failure or timeout at ANY time rejects the in-flight op through this + * with the SAME error text the one-shot machine emitted. + */ + private inflightReject: ((err: Error) => void) | undefined; + + constructor( + key: string, + pageId: string, + wsUrl: string, + token: string, + cfg: SessionConfig, + ephemeral: boolean, + opts: AcquireOptions | undefined, + ) { + this.key = key; + this.pageId = pageId; + this.wsUrl = wsUrl; + this.token = token; + this.cfg = cfg; + this.ephemeral = ephemeral; + this.opts = opts; + this.createdAt = Date.now(); + this.ydoc = new Y.Doc(); + } + + /** + * A cached session may be reused only when it is fully ready, still synced, + * has not lost its connection, and has not exceeded its max age (invariant 5 + * "validate on reuse" + the max-age acquire check). + */ + isReusable(): boolean { + return ( + !this.dead && + this.state === "ready" && + !this.connectionLost && + !!this.provider && + this.provider.synced === true && + Date.now() - this.createdAt < this.cfg.maxAgeMs + ); + } + + /** + * Connect and wait for the initial sync (onSynced) within CONNECT_TIMEOUT_MS. + * Idempotent: repeated calls return the same in-flight/settled promise. + */ + open(): Promise { + if (this.openPromise) return this.openPromise; + this.openPromise = new Promise((resolve, reject) => { + this.openResolve = resolve; + this.openReject = reject; + + this.connectTimer = setTimeout(() => { + // The 25s connect timeout: the collab connection never became ready. + this.opts?.onConnectTimeout?.(); + this.teardown( + new Error("Connection timeout to collaboration server"), + false, + ); + }, CONNECT_TIMEOUT_MS); + + if (process.env.DEBUG) + console.error(`Connecting to WebSocket: ${this.wsUrl}`); + + this.provider = providerFactory({ + url: this.wsUrl, + name: `page.${this.pageId}`, + document: this.ydoc, + token: this.token, + WebSocketPolyfill: WebSocket, + onConnect: () => { + if (process.env.DEBUG) console.error("WS Connect"); + }, + // An unexpected disconnect/close at ANY time (during the connect-wait, + // between edits, or during a persistence wait) makes the session dead: + // surface it now instead of hanging, reject any in-flight op with the + // same error text as the one-shot machine, and remove ourselves from + // the registry so the next acquire opens fresh. `teardown` is idempotent + // so the onClose our own destroy() triggers is a harmless no-op. + onDisconnect: () => { + if (process.env.DEBUG) console.error("WS Disconnect"); + this.teardown( + new Error( + "Collaboration connection closed before the update was persisted/synced", + ), + true, + ); + }, + onClose: () => { + if (process.env.DEBUG) console.error("WS Close"); + this.teardown( + new Error( + "Collaboration connection closed before the update was persisted/synced", + ), + true, + ); + }, + onSynced: () => { + if (this.dead || this.openSettled) return; + if (process.env.DEBUG) console.error("Connected and synced!"); + if (this.connectTimer) { + clearTimeout(this.connectTimer); + this.connectTimer = undefined; + } + this.state = "ready"; + this.openSettled = true; + this.openResolve?.(); + }, + onAuthenticationFailed: () => { + this.teardown( + new Error("Authentication failed for collaboration connection"), + true, + ); + }, + }); + }); + return this.openPromise; + } + + /** + * Run one atomic read -> transform -> write against the LIVE doc and wait for + * the server to acknowledge the write. + * + * INVARIANT 1 (read->write atomicity): between `TiptapTransformer.fromYdoc` + * and `applyDocToFragment` there is NO `await`. Yjs applies remote updates + * only when the event loop yields, so this synchronous block sees a consistent + * live doc and no concurrent human edit can interleave and be clobbered — + * exactly as in the one-shot onSynced code, just on a persistent provider. + * + * INVARIANT 2 (per-edit ack): after the write, resolve immediately if + * unsyncedChanges is already 0, else wait for the unsyncedChanges->0 event + * (PERSIST_TIMEOUT_MS), guarded by connectionLost so a reconnect handshake + * cannot report a false success. + * + * CONCURRENCY: not safe to invoke concurrently on ONE session — the caller + * MUST serialize (hold the per-page lock), mirroring acquireCollabSession. + * The in-flight op is tracked in a single `inflightReject` field, so an + * overlapping second call would clobber the first's rejector and leave it + * hanging on disconnect. A fail-fast guard below rejects the overlap instead. + * Sequential (awaited) mutates are fine: localFinish clears inflightReject + * before the promise settles, so the guard is clear by the time the next runs. + */ + mutate( + transform: (liveDoc: any) => any | null, + ): Promise { + // Belt-and-suspenders (acquire already validated): refuse to write on a + // session that is not in a live, synced, ready state. + if ( + this.dead || + this.state !== "ready" || + this.connectionLost || + !this.provider || + this.provider.synced !== true + ) { + return Promise.reject( + new Error("Collaboration session is not in a ready state"), + ); + } + + // Fail-fast on concurrent use: a second overlapping mutate would overwrite + // the first's inflightReject, so a disconnect would only reject the second + // and hang the first until PERSIST_TIMEOUT_MS. Reject the overlap WITHOUT + // touching the in-flight op's state (no localFinish/teardown here). + if (this.inflightReject) { + return Promise.reject( + new Error( + "mutate already in-flight; caller must serialize (hold the page lock)", + ), + ); + } + + return new Promise((resolve, reject) => { + let settled = false; + let persistTimer: ReturnType | undefined; + let unsyncedHandler: + | ((data: { number: number }) => void) + | undefined; + // The verifiable result resolved on every success/abort path. Set on + // abort (no-op report) and after a real write (computed change report). + let mutationResult: MutationResult; + + const localFinish = (err: Error | null, value?: MutationResult) => { + if (settled) return; + settled = true; + if (persistTimer) clearTimeout(persistTimer); + if (unsyncedHandler && this.provider) { + try { + this.provider.off("unsyncedChanges", unsyncedHandler); + } catch (e) {} + } + this.inflightReject = undefined; + if (err) reject(err); + else resolve(value as MutationResult); + // Post-settle lifecycle: an ephemeral (cache-disabled) session dies with + // its single op; a cached session that is still alive re-arms its idle + // TTL so the clock starts from the LAST op. + if (this.ephemeral) { + this.destroy("ephemeral op complete"); + } else if (!this.dead) { + this.armIdle(); + } + }; + + // Register so a disconnect/close/auth-failure/teardown rejects THIS op + // with the connection-loss error text. localFinish's `settled` guard makes + // a racing teardown + normal resolve safe (first one wins). + this.inflightReject = (e: Error) => localFinish(e); + + // Resolve once the server acknowledges our update: the provider increments + // unsyncedChanges when the local update is sent and decrements it on the + // server's SyncStatus(applied=true); reaching 0 means the authoritative + // in-memory ydoc on the server now contains our write. + const waitForPersistence = () => { + if (settled) return; + // A missing provider is a failure, not a success: without it the write + // can never have been acknowledged. + if (!this.provider) { + localFinish(new Error("collab provider gone before persistence")); + return; + } + if (this.provider.unsyncedChanges === 0) { + localFinish(null, mutationResult); + return; + } + persistTimer = setTimeout(() => { + localFinish( + new Error( + "Timeout waiting for collaboration server to persist the update", + ), + ); + }, PERSIST_TIMEOUT_MS); + unsyncedHandler = (data: { number: number }) => { + // Only treat unsyncedChanges->0 as success when the connection is + // still up. A transient disconnect + reconnect handshake can drive the + // counter back to 0 without our write being re-transmitted; in that + // case let the disconnect/close error win instead. + if (data.number === 0 && !this.connectionLost) { + localFinish(null, mutationResult); + } + }; + this.provider.on("unsyncedChanges", unsyncedHandler); + }; + + // CRITICAL: everything between reading the live doc and writing it back + // must stay synchronous (no await). While the JS event loop is not + // yielded, no incoming remote update can interleave, so any already-synced + // concurrent edits are preserved in liveDoc. + let newDoc: any; + let beforeDoc: any; + try { + let liveDoc = TiptapTransformer.fromYdoc(this.ydoc, "default"); + if ( + !liveDoc || + typeof liveDoc !== "object" || + !Array.isArray(liveDoc.content) + ) { + liveDoc = { type: "doc", content: [] }; + } + + // Snapshot the before-doc for the change report. Docs are + // JSON-serializable, so this is a safe deep clone. + beforeDoc = JSON.parse(JSON.stringify(liveDoc)); + + newDoc = transform(liveDoc); + + if (newDoc == null) { + // Transform aborted — write nothing, return the live doc with a no-op + // change report. + mutationResult = { + doc: liveDoc, + verify: { + changed: false, + textInserted: 0, + textDeleted: 0, + blocksChanged: 0, + marks: {}, + summary: "no changes (transform aborted)", + }, + }; + localFinish(null, mutationResult); + return; + } + + // Structural diff into the live fragment (issue #152): preserves the Yjs + // ids of unchanged nodes, so an open editor's cursor is not yanked to the + // end of the document on every agent write. + applyDocToFragment(this.ydoc, newDoc); + } catch (e) { + // Includes errors thrown by transform (e.g. "afterText not found", + // "text not found"): propagate them verbatim to the caller. + localFinish(e instanceof Error ? e : new Error(String(e))); + return; + } + + // Compute the verifiable change report AFTER the transact write: it only + // needs the JSON before/after, so it cannot affect the atomic read->write + // window, and summarizeChange never throws. + mutationResult = { + doc: newDoc, + verify: summarizeChange(beforeDoc, newDoc), + }; + if (process.env.DEBUG) + console.error("Content written, waiting for server to persist..."); + waitForPersistence(); + }); + } + + /** (Re)arm the idle TTL so the clock starts from the most recent activity. */ + armIdle(): void { + if (this.dead || this.ephemeral) return; + if (this.idleTimer) clearTimeout(this.idleTimer); + if (this.cfg.idleMs > 0) { + this.idleTimer = setTimeout(() => { + this.destroy("idle timeout"); + }, this.cfg.idleMs); + // Never let the idle timer keep the process alive. + (this.idleTimer as any).unref?.(); + } + } + + /** + * Idempotent teardown: mark dead, clear timers, remove from the registry, fail + * any pending open/in-flight op, and destroy the provider. `inflightError` is + * the error a pending open or in-flight op is rejected with; `connectionLoss` + * marks the session as connection-lost so the ack guard cannot report a false + * success on a racing unsyncedChanges->0. + */ + private teardown(inflightError: Error | null, connectionLoss: boolean): void { + if (this.dead) return; + this.dead = true; + this.state = "dead"; + if (connectionLoss) this.connectionLost = true; + + if (this.connectTimer) { + clearTimeout(this.connectTimer); + this.connectTimer = undefined; + } + if (this.idleTimer) { + clearTimeout(this.idleTimer); + this.idleTimer = undefined; + } + + // Remove ourselves from the registry (only if we are still the live entry — + // a re-open under the same key must not be evicted by our teardown). + if (sessions.get(this.key) === this) { + sessions.delete(this.key); + } + + // Fail a pending open() and any in-flight mutate with the terminal error. + if (!this.openSettled) { + this.openSettled = true; + this.openReject?.( + inflightError ?? new Error("Collaboration session destroyed"), + ); + } + if (this.inflightReject) { + const rej = this.inflightReject; + this.inflightReject = undefined; + rej(inflightError ?? new Error("Collaboration session destroyed")); + } + + if (this.provider) { + try { + this.provider.destroy(); + } catch (e) {} + this.provider = undefined; + } + } + + /** + * Public idempotent teardown used by the acquire/eviction paths and by a + * caller that wants the session dropped after a failed op ("next call + * reconnects fresh"). + */ + destroy(reason: string): void { + if (this.dead) return; + if (process.env.DEBUG) + console.error(`Destroying collab session ${this.pageId}: ${reason}`); + this.teardown(new Error(`Collaboration session destroyed: ${reason}`), false); + } +} + +/** key = wsUrl + pageId + collabToken (identity isolation: invariant 4). */ +const sessions = new Map(); + +function sessionKey(wsUrl: string, pageId: string, token: string): string { + // The token is part of the key so sessions are NEVER shared between different + // users' MCP sessions (HTTP mode), and a token rotation makes a new entry + // while the old one idles out. + return `${wsUrl}${pageId}${token}`; +} + +/** + * Get a live, synced CollabSession for a page, reusing a cached one when it is + * still valid or opening a fresh one otherwise. Does NOT take the per-page lock + * — the caller MUST already hold it (both call sites run inside withPageLock, + * which is not reentrant, so acquiring the lock here would deadlock + * mutateLiveContentUnlocked). + */ +export async function acquireCollabSession( + pageId: string, + collabToken: string, + baseUrl: string, + opts?: AcquireOptions, +): Promise { + const cfg = readConfig(); + const wsUrl = buildCollabWsUrl(baseUrl); + + // Cache disabled (rollback path): open an unregistered ephemeral session that + // self-destroys after its single op — the exact legacy per-op-provider flow. + if (cfg.idleMs <= 0) { + const session = new CollabSession( + sessionKey(wsUrl, pageId, collabToken), + pageId, + wsUrl, + collabToken, + cfg, + true, + opts, + ); + await session.open(); + return session; + } + + const key = sessionKey(wsUrl, pageId, collabToken); + const existing = sessions.get(key); + if (existing) { + if (existing.isReusable()) { + // Reuse. Refresh LRU order (re-insert = most recently used) and re-arm the + // idle TTL so the reuse counts as activity. + sessions.delete(key); + sessions.set(key, existing); + existing.armIdle(); + if (process.env.DEBUG) + console.error(`Reusing collab session for page ${pageId}`); + return existing; + } + // Stale (not synced / past max age / lost): drop it and open fresh. + existing.destroy("stale on reuse"); + } + + // Enforce the registry cap before inserting: destroy-evict the least recently + // used (the first entry in insertion order) until there is room. + while (sessions.size >= cfg.maxEntries) { + const oldestKey: string | undefined = sessions.keys().next().value; + if (oldestKey === undefined) break; + const victim = sessions.get(oldestKey); + if (victim) victim.destroy("evicted (LRU cap)"); + // destroy() removes it from the map; guard against a no-op destroy. + if (sessions.has(oldestKey)) sessions.delete(oldestKey); + } + + const session = new CollabSession( + key, + pageId, + wsUrl, + collabToken, + cfg, + false, + opts, + ); + sessions.set(key, session); + try { + await session.open(); + } catch (e) { + // Failed connect/sync: make sure it is not left cached. + session.destroy("open failed"); + throw e; + } + session.armIdle(); + if (process.env.DEBUG) + console.error(`Opened new collab session for page ${pageId}`); + return session; +} + +/** + * Destroy every cached session. Wired into the process shutdown so a hanging + * session does not keep a doc loaded on the server past exit. + */ +export function destroyAllSessions(): void { + for (const session of [...sessions.values()]) { + session.destroy("process shutdown"); + } + sessions.clear(); +} + +/** TEST-ONLY: number of currently cached sessions. */ +export function __sessionCountForTests(): number { + return sessions.size; +} diff --git a/packages/mcp/src/lib/collaboration.ts b/packages/mcp/src/lib/collaboration.ts index c14afbf8..db58aceb 100644 --- a/packages/mcp/src/lib/collaboration.ts +++ b/packages/mcp/src/lib/collaboration.ts @@ -1,4 +1,3 @@ -import { HocuspocusProvider } from "@hocuspocus/provider"; import { TiptapTransformer } from "@hocuspocus/transformer"; import * as Y from "yjs"; import WebSocket from "ws"; @@ -16,7 +15,8 @@ import { docmostExtensions, docmostSchema } from "./docmost-schema.js"; import { withPageLock } from "./page-lock.js"; import { sanitizeForYjs, findUnstorableAttr } from "./node-ops.js"; import { canonicalizeFootnotes } from "./footnote-canonicalize.js"; -import { summarizeChange, VerifyReport } from "./diff.js"; +import { VerifyReport } from "./diff.js"; +import { acquireCollabSession } from "./collab-session.js"; export { markdownToProseMirror }; @@ -194,26 +194,27 @@ export function assertYjsEncodable(doc: any): void { } } -/** Time we wait for the initial handshake/sync before giving up. */ -const CONNECT_TIMEOUT_MS = 25000; -/** Time we wait for the server to acknowledge our write before giving up. */ -const PERSIST_TIMEOUT_MS = 20000; - /** * Safely mutate the live content of a page over the collaboration websocket. * * This is the single safe write path for every MCP content mutation. It: * 1. serializes per-page writes through withPageLock (no two MCP writes on * the same page overlap); - * 2. connects to Hocuspocus and waits for the initial sync so the local ydoc - * mirrors the authoritative server doc — INCLUDING edits/comments/images - * that are not yet in the debounced REST snapshot; - * 3. inside onSynced, SYNCHRONOUSLY reads the live doc, runs `transform`, and - * writes the result back — with no `await` between read and write so no - * remote update can interleave and clobber concurrent human edits; + * 2. acquires a LIVE, synced CollabSession for the page (issue #400) — a + * cached provider whose local ydoc mirrors the authoritative server doc + * (INCLUDING edits/comments/images not yet in the debounced REST snapshot), + * reused across a series of edits instead of a fresh connect/auth/sync per + * call; + * 3. SYNCHRONOUSLY reads the live doc, runs `transform`, and writes the result + * back — with no `await` between read and write so no remote update can + * interleave and clobber concurrent human edits (CollabSession.mutate); * 4. waits for the server to acknowledge the write (unsyncedChanges -> 0) * before resolving, so the next operation observes our change. * + * On any mutate failure the session is destroyed so the next call reconnects + * fresh; the page lock is held for the whole acquire+mutate so the session's + * synchronous read->write window never overlaps another MCP write on the page. + * * `transform` receives the live ProseMirror doc and returns the NEW full * ProseMirror doc to write, or `null` to abort with no write (a no-op). If * `transform` throws, the error is propagated to the caller (not swallowed). @@ -230,7 +231,7 @@ export async function mutatePageContent( baseUrl: string, transform: (liveDoc: any) => any | null, ): Promise { - return withPageLock(pageId, () => { + return withPageLock(pageId, async () => { if (process.env.DEBUG) { console.error(`Starting realtime content mutate for page ${pageId}`); // Token prefix is sensitive; only log it under DEBUG. @@ -239,202 +240,15 @@ export async function mutatePageContent( ); } - const ydoc = new Y.Doc(); - const wsUrl = buildCollabWsUrl(baseUrl); - if (process.env.DEBUG) console.error(`Connecting to WebSocket: ${wsUrl}`); - - return new Promise((resolve, reject) => { - let provider: HocuspocusProvider | undefined; - let applied = false; // onSynced may fire again on reconnect — apply once. - let settled = false; - // Set true on disconnect/close so a reconnect-driven unsyncedChanges->0 - // cannot be mistaken for a successful persist of our write. - let connectionLost = false; - let connectTimer: ReturnType | undefined; - let persistTimer: ReturnType | undefined; - let unsyncedHandler: ((data: { number: number }) => void) | undefined; - - const cleanup = () => { - if (connectTimer) clearTimeout(connectTimer); - if (persistTimer) clearTimeout(persistTimer); - if (provider) { - if (unsyncedHandler) { - try { - provider.off("unsyncedChanges", unsyncedHandler); - } catch (err) {} - } - try { - provider.destroy(); - } catch (err) {} - } - }; - - const finish = (err: Error | null, value?: MutationResult) => { - if (settled) return; - settled = true; - cleanup(); - if (err) reject(err); - else resolve(value as MutationResult); - }; - - connectTimer = setTimeout(() => { - finish(new Error("Connection timeout to collaboration server")); - }, CONNECT_TIMEOUT_MS); - - // Resolve once the server has acknowledged our update. The provider - // increments unsyncedChanges when our local update is sent and - // decrements it when the server replies with a SyncStatus(applied=true); - // reaching 0 means the authoritative in-memory ydoc on the server now - // contains our write. - const waitForPersistence = () => { - if (settled) return; - // A missing provider is a failure, not a success: without it the write - // can never have been acknowledged. Only an actual unsyncedChanges===0 - // on a live provider counts as persisted. - if (!provider) { - finish(new Error("collab provider gone before persistence")); - return; - } - if (provider.unsyncedChanges === 0) { - finish(null, mutationResult); - return; - } - persistTimer = setTimeout(() => { - finish( - new Error( - "Timeout waiting for collaboration server to persist the update", - ), - ); - }, PERSIST_TIMEOUT_MS); - unsyncedHandler = (data: { number: number }) => { - // Only treat unsyncedChanges->0 as success when the connection is - // still up. A transient disconnect + reconnect handshake can drive - // the counter back to 0 without our write being re-transmitted; in - // that case let the disconnect/close error win instead. - if (data.number === 0 && !connectionLost) { - finish(null, mutationResult); - } - }; - provider.on("unsyncedChanges", unsyncedHandler); - }; - - // The verifiable result resolved on every success/abort path. Set on - // abort (no-op report) and after a real write (computed change report). - let mutationResult: MutationResult; - - provider = new HocuspocusProvider({ - url: wsUrl, - name: `page.${pageId}`, - document: ydoc, - token: collabToken, - // @ts-ignore - Required for Node.js environment - WebSocketPolyfill: WebSocket, - onConnect: () => { - if (process.env.DEBUG) console.error("WS Connect"); - }, - // An unexpected disconnect/close while we are still waiting (during the - // connect-wait before onSynced, or during the persistence wait after the - // write) means the update will never be acknowledged — surface it now - // instead of hanging until the connect/persist timeout fires. `finish` - // is idempotent via the `settled` flag, so the onClose that our own - // cleanup()->provider.destroy() triggers (after settled=true is set) is - // a harmless no-op and cannot cause a double-resolve. - onDisconnect: () => { - if (process.env.DEBUG) console.error("WS Disconnect"); - // Mark BEFORE finish so the unsyncedChanges handler (if it races) - // sees the connection as lost and won't report a false success. - connectionLost = true; - finish( - new Error( - "Collaboration connection closed before the update was persisted/synced", - ), - ); - }, - onClose: () => { - if (process.env.DEBUG) console.error("WS Close"); - // Mark BEFORE finish so the unsyncedChanges handler (if it races) - // sees the connection as lost and won't report a false success. - connectionLost = true; - finish( - new Error( - "Collaboration connection closed before the update was persisted/synced", - ), - ); - }, - onSynced: () => { - if (applied || settled) return; - applied = true; - if (process.env.DEBUG) console.error("Connected and synced!"); - - // CRITICAL: everything between reading the live doc and writing it - // back must stay synchronous (no await). While the JS event loop is - // not yielded, no incoming remote update can interleave, so any - // already-synced concurrent edits are preserved in liveDoc. - let newDoc: any; - let beforeDoc: any; - try { - let liveDoc = TiptapTransformer.fromYdoc(ydoc, "default"); - if ( - !liveDoc || - typeof liveDoc !== "object" || - !Array.isArray(liveDoc.content) - ) { - liveDoc = { type: "doc", content: [] }; - } - - // Snapshot the before-doc for the change report. Docs are - // JSON-serializable, so this is a safe deep clone. - beforeDoc = JSON.parse(JSON.stringify(liveDoc)); - - newDoc = transform(liveDoc); - - if (newDoc == null) { - // Transform aborted — write nothing, return the live doc with a - // no-op change report. - mutationResult = { - doc: liveDoc, - verify: { - changed: false, - textInserted: 0, - textDeleted: 0, - blocksChanged: 0, - marks: {}, - summary: "no changes (transform aborted)", - }, - }; - finish(null, mutationResult); - return; - } - - // Structural diff into the live fragment (issue #152): preserves - // the Yjs ids of unchanged nodes, so an open editor's cursor is not - // yanked to the end of the document on every agent write. - applyDocToFragment(ydoc, newDoc); - } catch (e) { - // Includes errors thrown by transform (e.g. "afterText not found", - // "text not found"): propagate them verbatim to the caller. - finish(e instanceof Error ? e : new Error(String(e))); - return; - } - - // Compute the verifiable change report AFTER the transact write: it - // only needs the JSON before/after, so it cannot affect the atomic - // read->write window, and summarizeChange never throws. - mutationResult = { - doc: newDoc, - verify: summarizeChange(beforeDoc, newDoc), - }; - if (process.env.DEBUG) - console.error("Content written, waiting for server to persist..."); - waitForPersistence(); - }, - onAuthenticationFailed: () => { - finish( - new Error("Authentication failed for collaboration connection"), - ); - }, - }); - }); + const session = await acquireCollabSession(pageId, collabToken, baseUrl); + try { + return await session.mutate(transform); + } catch (e) { + // Drop the session on any failure so the next call reconnects fresh (this + // also closes the "reconnect drove the counter to 0" false-success class). + session.destroy("mutate failed"); + throw e; + } }); } diff --git a/packages/mcp/src/stdio.ts b/packages/mcp/src/stdio.ts index fcb18645..dc06bbaa 100644 --- a/packages/mcp/src/stdio.ts +++ b/packages/mcp/src/stdio.ts @@ -1,6 +1,7 @@ #!/usr/bin/env node import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js"; import { createDocmostMcpServer } from "./index.js"; +import { destroyAllSessions } from "./lib/collab-session.js"; // Standalone stdio entrypoint. This restores the original behavior of the // package when run as a CLI (`docmost-mcp`): it reads credentials from the @@ -33,6 +34,20 @@ async function run() { console.error("Uncaught exception:", error); }); + // Teardown hook (issue #400): destroy every cached live CollabSession on exit + // so a hanging session does not keep a doc loaded on the server (which would + // also defer the server's afterUnloadDocument cleanup). `exit` runs the + // synchronous idempotent teardown; SIGINT/SIGTERM also run it, then exit. + process.on("exit", () => { + destroyAllSessions(); + }); + for (const signal of ["SIGINT", "SIGTERM"] as const) { + process.on(signal, () => { + destroyAllSessions(); + process.exit(0); + }); + } + const server = createDocmostMcpServer({ apiUrl: API_URL!, email: EMAIL!, diff --git a/packages/mcp/test/mock/resolve-page-id-collab-doc-name.test.mjs b/packages/mcp/test/mock/resolve-page-id-collab-doc-name.test.mjs index 3952de5a..feaf94c6 100644 --- a/packages/mcp/test/mock/resolve-page-id-collab-doc-name.test.mjs +++ b/packages/mcp/test/mock/resolve-page-id-collab-doc-name.test.mjs @@ -19,6 +19,7 @@ import { WebSocketServer } from "ws"; import { Hocuspocus } from "@hocuspocus/server"; import { DocmostClient } from "../../build/client.js"; import { buildYDoc } from "../../build/lib/collaboration.js"; +import { destroyAllSessions } from "../../build/lib/collab-session.js"; // Import the SAME page-lock module instance that build/client.js imports. ESM // caches modules by resolved URL, so this `withPageLock` shares the very // per-page mutex map (`chains`) the client uses — letting the replaceImage test @@ -188,6 +189,10 @@ async function spawnCollabStack(opts = {}) { const openStacks = []; after(async () => { + // #400: tests now leave a cached live CollabSession per page. Destroy them + // first (closes the client ws) so the server.close() below is not racing an + // open collab connection. + destroyAllSessions(); await Promise.all( openStacks.map( ({ server, hocuspocus }) => @@ -270,17 +275,23 @@ test("a UUID input is passed through unchanged and triggers NO /pages/info fetch ); }); -test("a repeated slugId edit resolves the UUID only once (cache)", async () => { +test("repeated slugId edits reuse ONE live collab session and resolve the UUID only once (#400 cache)", async () => { const { state, baseURL } = await spawnCollabStack(); const client = new DocmostClient(baseURL, "user@example.com", "pw"); - // Each mock connection re-seeds a fresh "hello world" doc (the mock does not - // persist across connects), so both edits target "hello". The cache assertion - // only concerns the slugId->uuid resolution, not the document content. + // #400: a series of edits on the same page reuses ONE live CollabSession, so + // the connect/handshake happens once and the collab doc is OPENED a single + // time (not per edit). The live ydoc persists between edits (the whole point), + // so the second edit sees the first edit's result: after "hello" -> "hi world" + // it targets the still-present "world". await client.editPageText(SLUG, [{ find: "hello", replace: "hi" }]); - await client.editPageText(SLUG, [{ find: "hello", replace: "hey" }]); + await client.editPageText(SLUG, [{ find: "world", replace: "planet" }]); - assert.deepEqual(state.docNames, [`page.${UUID}`, `page.${UUID}`]); + assert.deepEqual( + state.docNames, + [`page.${UUID}`], + "the two edits must reuse one live collab session -> a single collab-doc open (#400)", + ); assert.equal( state.pagesInfoCalls.length, 1, @@ -325,8 +336,9 @@ test("replaceImage opens by the resolved UUID AND keys its page lock by that UUI await uploadStarted; // deterministic: replaceImage now holds its page lock. // (a) OPEN BY UUID: the only collab doc opened so far (the scan pass) used the - // canonical UUID, never the slugId. (The write pass opens a second time after - // we release the gate; asserted at the end.) + // canonical UUID, never the slugId. (#400: the write pass will REUSE this same + // live session rather than reopen, so docNames stays a single entry — asserted + // at the end.) assert.deepEqual( state.docNames, [`page.${UUID}`], @@ -378,8 +390,9 @@ test("replaceImage opens by the resolved UUID AND keys its page lock by that UUI assert.equal(res.success, true); assert.equal(res.replaced, 1, "the one seeded image must be repointed"); - // Both opens (scan pass + write pass) used the UUID; the slugId never appears. - assert.deepEqual(state.docNames, [`page.${UUID}`, `page.${UUID}`]); + // #400: the write pass REUSES the scan pass's live session, so the collab doc + // is opened ONCE across both passes (never reopened, never by the slugId). + assert.deepEqual(state.docNames, [`page.${UUID}`]); assert.ok( !state.docNames.includes(`page.${SLUG}`), "replaceImage must NEVER open the collab doc by the slugId (the #260 bug)", diff --git a/packages/mcp/test/unit/collab-session.test.mjs b/packages/mcp/test/unit/collab-session.test.mjs new file mode 100644 index 00000000..f5874721 --- /dev/null +++ b/packages/mcp/test/unit/collab-session.test.mjs @@ -0,0 +1,348 @@ +import { test, beforeEach, afterEach, mock } from "node:test"; +import assert from "node:assert/strict"; +import { EventEmitter } from "node:events"; + +import { + acquireCollabSession, + destroyAllSessions, + __setCollabProviderFactory, + __sessionCountForTests, +} from "../../build/lib/collab-session.js"; +import { withPageLock } from "../../build/lib/page-lock.js"; + +// A stand-in for HocuspocusProvider: it shares the ydoc (so the real yjs +// read/transform/write in CollabSession.mutate runs unchanged), auto-completes +// the connect+sync handshake on a microtask (unaffected by mock timers), and +// exposes hooks to drive disconnect/close/auth-failure and the unsyncedChanges +// ack. There is no collaboration server in the test env, so every test drives +// the provider through this fake. +class FakeProvider extends EventEmitter { + static instances = []; + static connectCount = 0; + + static reset() { + FakeProvider.instances = []; + FakeProvider.connectCount = 0; + } + + static last() { + return FakeProvider.instances[FakeProvider.instances.length - 1]; + } + + constructor(config, opts = {}) { + super(); + this.config = config; + this.ydoc = config.document; + this.synced = false; + this.unsyncedChanges = opts.unsynced ?? 0; + this.destroyed = false; + FakeProvider.instances.push(this); + FakeProvider.connectCount += 1; + if (opts.autoSync !== false) { + // Real HocuspocusProvider fires onSynced asynchronously after the + // handshake; a microtask reproduces that without depending on timers. + queueMicrotask(() => { + if (this.destroyed) return; + this.config.onConnect?.(); + this.synced = true; + this.config.onSynced?.(); + }); + } + } + + destroy() { + this.destroyed = true; + } + + // --- test drivers --- + _disconnect() { + this.config.onDisconnect?.(); + } + _close() { + this.config.onClose?.(); + } + _authFail() { + this.config.onAuthenticationFailed?.(); + } + _ack() { + this.unsyncedChanges = 0; + this.emit("unsyncedChanges", { number: 0 }); + } +} + +/** Build a provider factory that stamps every provider with the given opts. */ +function factory(opts = {}) { + return (config) => new FakeProvider(config, opts); +} + +/** A minimal, schema-valid ProseMirror doc for a write. */ +function docWith(text) { + return { + type: "doc", + content: [ + { type: "paragraph", content: [{ type: "text", text: String(text) }] }, + ], + }; +} + +const ENV_KEYS = [ + "MCP_COLLAB_SESSION_IDLE_MS", + "MCP_COLLAB_SESSION_MAX_AGE_MS", + "MCP_COLLAB_SESSION_MAX_ENTRIES", +]; +let savedEnv; + +beforeEach(() => { + savedEnv = {}; + for (const k of ENV_KEYS) savedEnv[k] = process.env[k]; + FakeProvider.reset(); + __setCollabProviderFactory(factory()); +}); + +afterEach(() => { + destroyAllSessions(); + __setCollabProviderFactory(null); + mock.timers.reset(); + for (const k of ENV_KEYS) { + if (savedEnv[k] === undefined) delete process.env[k]; + else process.env[k] = savedEnv[k]; + } +}); + +test("acquire opens one provider; reuse returns the SAME live session", async () => { + const a = await acquireCollabSession("page-1", "tok", "http://h/api"); + const b = await acquireCollabSession("page-1", "tok", "http://h/api"); + assert.equal(a, b, "same (wsUrl,page,token) must reuse the session"); + assert.equal(FakeProvider.connectCount, 1, "exactly one connect/sync"); + assert.equal(__sessionCountForTests(), 1); +}); + +test("N mutates on one page open the provider ONCE (coalesced series)", async () => { + const session = await acquireCollabSession("page-1", "tok", "http://h/api"); + for (let i = 0; i < 20; i++) { + const r = await session.mutate(() => docWith(`edit ${i}`)); + assert.ok(r && r.doc, "each mutate resolves a MutationResult"); + } + assert.equal( + FakeProvider.connectCount, + 1, + "20 mutates must not reconnect — one live provider for the whole series", + ); +}); + +test("mutate preserves the transform-abort no-op report (null transform)", async () => { + const session = await acquireCollabSession("page-1", "tok", "http://h/api"); + const r = await session.mutate(() => null); + assert.equal(r.verify.changed, false); + assert.equal(r.verify.summary, "no changes (transform aborted)"); +}); + +test("registry key includes the token: different tokens => different sessions", async () => { + const a = await acquireCollabSession("page-1", "tok-A", "http://h/api"); + const b = await acquireCollabSession("page-1", "tok-B", "http://h/api"); + assert.notEqual(a, b, "different tokens must never share a session"); + assert.equal(FakeProvider.connectCount, 2); + assert.equal(__sessionCountForTests(), 2); + // Same token reuses. + const a2 = await acquireCollabSession("page-1", "tok-A", "http://h/api"); + assert.equal(a, a2); + assert.equal(FakeProvider.connectCount, 2); +}); + +test("disconnect at any time kills the session and removes it from the registry", async () => { + const a = await acquireCollabSession("page-1", "tok", "http://h/api"); + assert.equal(__sessionCountForTests(), 1); + FakeProvider.last()._disconnect(); + assert.equal(__sessionCountForTests(), 0, "dead session is deregistered"); + // Re-acquire opens a brand new provider. + const b = await acquireCollabSession("page-1", "tok", "http://h/api"); + assert.notEqual(a, b); + assert.equal(FakeProvider.connectCount, 2); +}); + +test("an in-flight mutate rejects with the connection-closed text on disconnect", async () => { + __setCollabProviderFactory(factory({ unsynced: 1 })); // stay pending after write + const session = await acquireCollabSession("page-1", "tok", "http://h/api"); + const p = session.mutate(() => docWith("x")); + // The write happened synchronously; persistence is pending. Now the socket drops. + FakeProvider.last()._disconnect(); + await assert.rejects( + p, + /Collaboration connection closed before the update was persisted\/synced/, + ); +}); + +test("auth failure rejects the pending open with the auth error text", async () => { + __setCollabProviderFactory(factory({ autoSync: false })); + const p = acquireCollabSession("page-1", "tok", "http://h/api"); + // Let the provider be constructed, then fail auth. + await Promise.resolve(); + FakeProvider.last()._authFail(); + await assert.rejects( + p, + /Authentication failed for collaboration connection/, + ); + assert.equal(__sessionCountForTests(), 0); +}); + +test("mutate-error invalidates the session (caller destroys; re-acquire is fresh)", async () => { + const session = await acquireCollabSession("page-1", "tok", "http://h/api"); + await assert.rejects( + session.mutate(() => { + throw new Error("afterText not found"); + }), + /afterText not found/, + ); + // The production caller destroys on failure; mirror that here. + session.destroy("mutate failed"); + assert.equal(__sessionCountForTests(), 0); + const fresh = await acquireCollabSession("page-1", "tok", "http://h/api"); + assert.notEqual(fresh, session); + assert.equal(FakeProvider.connectCount, 2); +}); + +test("a pending write resolves when the server acks (unsyncedChanges -> 0)", async () => { + __setCollabProviderFactory(factory({ unsynced: 1 })); + const session = await acquireCollabSession("page-1", "tok", "http://h/api"); + const p = session.mutate(() => docWith("y")); + FakeProvider.last()._ack(); + const r = await p; + assert.ok(r.doc); +}); + +test("concurrent mutate on one session: the second rejects, the first is unaffected", async () => { + __setCollabProviderFactory(factory({ unsynced: 1 })); // first stays pending after write + const session = await acquireCollabSession("page-1", "tok", "http://h/api"); + // First mutate: write happens synchronously, persistence ack still pending. + const first = session.mutate(() => docWith("first")); + // Second overlapping mutate on the SAME session must fail fast. + await assert.rejects( + session.mutate(() => docWith("second")), + /mutate already in-flight; caller must serialize \(hold the page lock\)/, + ); + // The first op is untouched: its rejector was not clobbered. Ack it now. + FakeProvider.last()._ack(); + const r = await first; + assert.ok(r.doc, "the first in-flight mutate still resolves on its own ack"); + assert.equal(FakeProvider.connectCount, 1, "no reconnect from the rejected overlap"); +}); + +test("sequential mutates on one session both succeed (the guard doesn't break serialized use)", async () => { + const session = await acquireCollabSession("page-1", "tok", "http://h/api"); + // Await the first fully, then run the second: inflightReject was cleared by + // localFinish before the first settled, so the guard is clear for the second. + const r1 = await session.mutate(() => docWith("one")); + assert.ok(r1.doc, "first sequential mutate resolves"); + const r2 = await session.mutate(() => docWith("two")); + assert.ok(r2.doc, "second sequential mutate resolves — guard not tripped"); + assert.equal(FakeProvider.connectCount, 1, "sequential mutates reuse one provider"); +}); + +test("connect timeout rejects with the connect-timeout text and fires the metric hook", async () => { + mock.timers.enable({ apis: ["setTimeout"] }); + __setCollabProviderFactory(factory({ autoSync: false })); + let metricFired = 0; + const p = acquireCollabSession("page-1", "tok", "http://h/api", { + onConnectTimeout: () => { + metricFired += 1; + }, + }); + mock.timers.tick(25000); + await assert.rejects(p, /Connection timeout to collaboration server/); + assert.equal(metricFired, 1); + assert.equal(__sessionCountForTests(), 0); +}); + +test("idle TTL destroys the session; re-acquire reconnects", async () => { + mock.timers.enable({ apis: ["setTimeout"] }); + process.env.MCP_COLLAB_SESSION_IDLE_MS = "1000"; + const a = await acquireCollabSession("page-1", "tok", "http://h/api"); + assert.equal(__sessionCountForTests(), 1); + mock.timers.tick(1000); // idle fires + assert.equal(__sessionCountForTests(), 0, "idle timeout destroyed it"); + const b = await acquireCollabSession("page-1", "tok", "http://h/api"); + assert.notEqual(a, b); + assert.equal(FakeProvider.connectCount, 2); +}); + +test("max age is enforced at acquire (destroy + open fresh), idle held off", async () => { + mock.timers.enable({ apis: ["setTimeout", "Date"] }); + process.env.MCP_COLLAB_SESSION_MAX_AGE_MS = "1000"; + process.env.MCP_COLLAB_SESSION_IDLE_MS = "10000000"; // never fires in this test + const a = await acquireCollabSession("page-1", "tok", "http://h/api"); + mock.timers.tick(2000); // past max age, below idle + const b = await acquireCollabSession("page-1", "tok", "http://h/api"); + assert.notEqual(a, b, "a session past its max age must be replaced"); + assert.equal(FakeProvider.connectCount, 2); +}); + +test("registry cap: least-recently-used session is destroy-evicted", async () => { + process.env.MCP_COLLAB_SESSION_MAX_ENTRIES = "2"; + const s1 = await acquireCollabSession("page-1", "tok", "http://h/api"); + const p1prov = FakeProvider.last(); + const s2 = await acquireCollabSession("page-2", "tok", "http://h/api"); + const p2prov = FakeProvider.last(); + assert.equal(__sessionCountForTests(), 2); + + // Touch page-1 so page-2 becomes the least-recently-used entry. + assert.equal(await acquireCollabSession("page-1", "tok", "http://h/api"), s1); + assert.equal(FakeProvider.connectCount, 2, "reuse must not reconnect"); + + // A third distinct page (cap 2) must evict the LRU (page-2), not page-1. + const s3 = await acquireCollabSession("page-3", "tok", "http://h/api"); + assert.equal(__sessionCountForTests(), 2); + assert.equal(FakeProvider.connectCount, 3); + assert.equal(p2prov.destroyed, true, "the LRU provider was destroy-evicted"); + assert.equal(p1prov.destroyed, false, "the MRU page-1 survived"); + + // page-1 still reused (survived), page-3 still reused. + assert.equal(await acquireCollabSession("page-1", "tok", "http://h/api"), s1); + assert.equal(await acquireCollabSession("page-3", "tok", "http://h/api"), s3); + assert.equal(FakeProvider.connectCount, 3, "no extra reconnects"); +}); + +test("MCP_COLLAB_SESSION_IDLE_MS=0 disables the cache (legacy provider-per-op)", async () => { + process.env.MCP_COLLAB_SESSION_IDLE_MS = "0"; + const s1 = await acquireCollabSession("page-1", "tok", "http://h/api"); + // Never registered (ephemeral). + assert.equal(__sessionCountForTests(), 0); + const r1 = await s1.mutate(() => docWith("a")); + assert.ok(r1.doc); + // After its single op the ephemeral session self-destroyed. + assert.equal(FakeProvider.instances[0].destroyed, true); + // A second op opens a BRAND NEW provider (no reuse). + const s2 = await acquireCollabSession("page-1", "tok", "http://h/api"); + assert.notEqual(s1, s2); + assert.equal(FakeProvider.connectCount, 2); +}); + +test("replaceImage-shaped flow: acquire under an EXTERNAL page lock does not deadlock and reuses one session", async () => { + const pageId = "page-lock"; + // Mirror replaceImage: hold ONE withPageLock across scan (read-only) + write, + // each going through the non-locking acquireCollabSession. + const result = await withPageLock(pageId, async () => { + // pass 1: read-only scan (transform returns null -> no write) + const scan = await acquireCollabSession(pageId, "tok", "http://h/api"); + const scanRes = await scan.mutate(() => null); + assert.equal(scanRes.verify.changed, false); + // pass 2: the actual write, same held lock + const write = await acquireCollabSession(pageId, "tok", "http://h/api"); + return write.mutate(() => docWith("repointed")); + }); + assert.ok(result.doc, "the locked scan+write completed without deadlock"); + assert.equal( + FakeProvider.connectCount, + 1, + "both passes under the held lock reuse ONE live session", + ); +}); + +test("destroyAllSessions tears down every cached session", async () => { + await acquireCollabSession("page-1", "tok", "http://h/api"); + await acquireCollabSession("page-2", "tok", "http://h/api"); + assert.equal(__sessionCountForTests(), 2); + const provs = [...FakeProvider.instances]; + destroyAllSessions(); + assert.equal(__sessionCountForTests(), 0); + assert.ok(provs.every((p) => p.destroyed), "all providers destroyed"); +});