vvzvlad/gitmost
1
0
Fork 0
Code Issues 31 Pull Requests 15 Packages Projects Releases Activity

perf(mcp): кэш живого collab-соединения (CollabSession) — серия правок за один connect/sync #431

Merged
vvzvlad merged 2 commits from perf/400-collab-session into develop 2026-07-10 04:26:08 +03:00
Conversation 3 Commits 2 Files Changed 7
+1105 -392
7 changed files with 1105 additions and 392 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/mcp/src/client.ts
+22 -172
View File
@@ -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<MutationResult> {
const CONNECT_TIMEOUT_MS = 25000;
const PERSIST_TIMEOUT_MS = 20000;
const ydoc = new Y.Doc();
const wsUrl = buildCollabWsUrl(this.apiUrl);
return new Promise<MutationResult>((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<typeof setTimeout> | undefined;
let persistTimer: ReturnType<typeof setTimeout> | 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;
}
}
/**
packages/mcp/src/index.ts
+5
View File
@@ -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()).
packages/mcp/src/lib/collab-session.ts
+668
View File
@@ -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<typeof setTimeout> | undefined;
private idleTimer: ReturnType<typeof setTimeout> | undefined;
private openPromise: Promise<void> | 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<void> {
if (this.openPromise) return this.openPromise;
this.openPromise = new Promise<void>((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<MutationResult> {
// 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<MutationResult>((resolve, reject) => {
let settled = false;
let persistTimer: ReturnType<typeof setTimeout> | 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<string, CollabSession>();
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<CollabSession> {
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;
}
packages/mcp/src/lib/collaboration.ts
+24 -210
View File
@@ -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<MutationResult> {
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<MutationResult>((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<typeof setTimeout> | undefined;
let persistTimer: ReturnType<typeof setTimeout> | 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;
}
});
}
packages/mcp/src/stdio.ts
+15
View File
@@ -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!,
packages/mcp/test/mock/resolve-page-id-collab-doc-name.test.mjs
+23 -10
View File
@@ -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)",
packages/mcp/test/unit/collab-session.test.mjs
+348
View File
@@ -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");
});