Merge pull request 'perf(mcp): кэш collab-токена в клиенте — чтобы кэш CollabSession (#400) попадал (#435)' (#439) from perf/435-collab-token-cache into develop

Reviewed-on: #439

packages/mcp/src/client.ts

@@ -167,6 +167,35 @@ function isUuid(value: string): boolean {
   return typeof value === "string" && UUID_RE.test(value);
 }
 
+/**
+ * Collab-token cache TTL in milliseconds (issue #435). Read fresh from the
+ * environment on every mint — like collab-session.ts readConfig — so tests and a
+ * live rollback can change it without reloading the module.
+ *
+ * Why a cache at all: the live CollabSession registry (#400/#431) keys sessions
+ * on (wsUrl, pageId, collabToken) for identity isolation (invariant 4). But BOTH
+ * collab-token sources mint a FRESH token per mutation — the in-app provider
+ * re-signs a JWT whose iat/exp (seconds) changes every second, and the external
+ * MCP POSTs /auth/collab-token each call — so the token in the key changed on
+ * every op and the session was almost never reused (connect-storms, 25s
+ * timeouts, zombie sessions). Caching the token per-client keeps the key stable
+ * across a burst of mutations so ONE session is reused.
+ *
+ * Default 5 min: well under the 24h collab-token lifetime AND <= the collab
+ * session max-age (10 min, MCP_COLLAB_SESSION_MAX_AGE_MS), so the
+ * permission-staleness window is not widened beyond what #431 already accepted.
+ * The rollback knob is an EXPLICIT 0 (or a negative number): that DISABLES the
+ * cache — an exact fetch-per-call legacy path, mirroring how idleMs<=0 disables
+ * the session cache. Unset OR unparseable (e.g. a typo like "5min", "abc") falls
+ * back to the 5-min default with the cache ON — parseInt yields NaN, which is
+ * treated as "not configured", not as "disabled". So to turn the cache off you
+ * must set the value to exactly 0, not to garbage.
+ */
+function readCollabTokenTtlMs(): number {
+  const raw = parseInt(process.env.MCP_COLLAB_TOKEN_TTL_MS ?? "", 10);
+  return Number.isFinite(raw) ? Math.max(0, raw) : 5 * 60 * 1000;
+}
+
 export class DocmostClient {
   private client: AxiosInstance;
   private token: string | null = null;
@@ -205,6 +234,15 @@ export class DocmostClient {
   // resolvePageId), so only slugId->uuid entries are stored/read here.
   private pageIdCache = new Map<string, string>();
 
+  // Collab-token cache (issue #435): the last minted collab token plus the
+  // wall-clock time it was minted, so a burst of content mutations reuses ONE
+  // token and therefore ONE live CollabSession (whose registry key includes the
+  // token — #400 invariant 4). Per-instance: a DocmostClient is built per
+  // user/per chat request, so a cached token can never leak across identities.
+  // Reset whenever the client's identity changes (login() / this.token cleared);
+  // bypassed on a forced refresh (the 401/403 reauth path). null = no token yet.
+  private collabTokenCache: { token: string; mintedAt: number } | null = null;
+
   // Two construction forms:
   //  - new DocmostClient(config)                  // discriminated union (current)
   //  - new DocmostClient(baseURL, email, password) // legacy positional creds
@@ -273,8 +311,11 @@ export class DocmostClient {
 
         if (config && isAuthError && !config._retry && !isLoginRequest) {
           config._retry = true;
-          // Drop the stale token + Authorization header before re-login.
+          // Drop the stale token + Authorization header before re-login. Also
+          // clear the collab-token cache (#435): a new identity/login must not
+          // keep serving a collab token minted under the old one.
           this.token = null;
+          this.collabTokenCache = null;
           delete this.client.defaults.headers.common["Authorization"];
           try {
             await this.login();
@@ -323,6 +364,9 @@ export class DocmostClient {
             throw new Error("getToken returned an empty token");
           }
           this.token = token;
+          // Identity (re)established: drop any collab token minted under a
+          // previous identity so the #435 cache can never outlive it.
+          this.collabTokenCache = null;
           this.client.defaults.headers.common["Authorization"] =
             `Bearer ${token}`;
         })
@@ -345,8 +389,34 @@ export class DocmostClient {
    * by this.client's response interceptor; this helper replicates that
    * behaviour for collab-token requests: ensure a token, try once, and on an
    * expired-token auth error perform a fresh login and retry exactly once.
+   *
+   * Collab-token cache (issue #435): both sources — the getCollabToken provider
+   * (in-app agent) AND the REST /auth/collab-token endpoint (external MCP) — mint
+   * a FRESH token per call, whose string therefore changes every op. Since the
+   * live CollabSession registry keys on the token string (#400/#431 invariant 4),
+   * that churned the key and defeated session reuse. So we cache the last minted
+   * token per-client for readCollabTokenTtlMs() and hand it back for a burst of
+   * mutations, keeping the session key stable. `forceRefresh` bypasses the cache
+   * (the 401/403 reauth retry uses it, so the retry cannot be handed the same
+   * stale token that just failed — otherwise reauth would be a no-op). TTL 0
+   * disables the cache: exact fetch-per-call legacy behaviour.
    */
-  private async getCollabTokenWithReauth(): Promise<string> {
+  private async getCollabTokenWithReauth(
+    forceRefresh = false,
+  ): Promise<string> {
+    const ttl = readCollabTokenTtlMs();
+    // Serve the cached collab token while it is still fresh (identity isolation
+    // is preserved: the cache is a per-instance field on a client built per
+    // user/per chat request, and it is cleared on every identity change).
+    if (
+      !forceRefresh &&
+      ttl > 0 &&
+      this.collabTokenCache &&
+      Date.now() - this.collabTokenCache.mintedAt < ttl
+    ) {
+      return this.collabTokenCache.token;
+    }
+
     // Collab-token PROVIDER path: when a getCollabToken provider was supplied
     // (the internal agent's provenance collab token), use it instead of the
     // REST /auth/collab-token endpoint. Re-invoke it once on a 401/403 (e.g. the
@@ -357,23 +427,13 @@ export class DocmostClient {
         if (typeof token !== "string" || token.length === 0) {
           throw new Error("getCollabToken returned an empty token");
         }
-        return token;
+        return this.rememberCollabToken(token, ttl);
       } catch (e) {
-        const axiosStatus = axios.isAxiosError(e)
-          ? e.response?.status
-          : undefined;
-        const attachedStatus = (e as any)?.status;
-        const isAuthError =
-          axiosStatus === 401 ||
-          axiosStatus === 403 ||
-          attachedStatus === 401 ||
-          attachedStatus === 403;
-        if (isAuthError) {
-          const token = await this.getCollabTokenFn();
-          if (typeof token !== "string" || token.length === 0) {
-            throw new Error("getCollabToken returned an empty token");
-          }
-          return token;
+        // On an auth error retry EXACTLY once, forcing a refresh so the retry
+        // re-invokes the provider (bypassing the cache) for a genuinely fresh
+        // token. `!forceRefresh` bounds it to a single retry (no loop).
+        if (this.isCollabAuthError(e) && !forceRefresh) {
+          return this.getCollabTokenWithReauth(true);
         }
         throw e;
       }
@@ -381,28 +441,51 @@ export class DocmostClient {
 
     await this.ensureAuthenticated();
     try {
-      return await getCollabToken(this.apiUrl, this.token!);
+      const token = await getCollabToken(this.apiUrl, this.token!);
+      return this.rememberCollabToken(token, ttl);
     } catch (e) {
       // getCollabToken wraps the AxiosError in a plain Error but attaches the
-      // HTTP status as `.status`, so detect an auth failure via either the raw
-      // AxiosError shape OR the attached status.
-      const axiosStatus = axios.isAxiosError(e)
-        ? e.response?.status
-        : undefined;
-      const attachedStatus = (e as any)?.status;
-      const isAuthError =
-        axiosStatus === 401 ||
-        axiosStatus === 403 ||
-        attachedStatus === 401 ||
-        attachedStatus === 403;
-      if (isAuthError) {
+      // HTTP status as `.status`, so isCollabAuthError detects an auth failure
+      // via either the raw AxiosError shape OR the attached status.
+      if (this.isCollabAuthError(e) && !forceRefresh) {
+        // Fresh login (which clears this.token AND the collab-token cache), then
+        // retry exactly once with the cache bypassed via forceRefresh.
         await this.login();
-        return await getCollabToken(this.apiUrl, this.token!);
+        return this.getCollabTokenWithReauth(true);
       }
       throw e;
     }
   }
 
+  /**
+   * Store a freshly minted collab token in the per-client cache (issue #435) and
+   * return it unchanged. No-op write when the cache is disabled (ttl<=0) or the
+   * token is empty, so a disabled cache is exact fetch-per-call legacy behaviour
+   * and a bad token is never cached.
+   */
+  private rememberCollabToken(token: string, ttl: number): string {
+    if (ttl > 0 && typeof token === "string" && token.length > 0) {
+      this.collabTokenCache = { token, mintedAt: Date.now() };
+    }
+    return token;
+  }
+
+  /**
+   * True when an error carries a 401/403 — either as a raw AxiosError
+   * (`error.response.status`) or as the plain-Error `.status` that
+   * lib/auth-utils.getCollabToken attaches after wrapping the AxiosError.
+   */
+  private isCollabAuthError(e: unknown): boolean {
+    const axiosStatus = axios.isAxiosError(e) ? e.response?.status : undefined;
+    const attachedStatus = (e as any)?.status;
+    return (
+      axiosStatus === 401 ||
+      axiosStatus === 403 ||
+      attachedStatus === 401 ||
+      attachedStatus === 403
+    );
+  }
+
   /**
    * Connect to the collaboration websocket, read the live doc, apply
    * `transform`, write the result, and wait for the server to persist it —

packages/mcp/test/mock/collab-token-cache.test.mjs

@@ -0,0 +1,282 @@
+// Unit tests for the collab-token cache (issue #435). The live CollabSession
+// registry (#400/#431) keys sessions on (wsUrl, pageId, collabToken), so a token
+// string that changes every op defeats reuse. This cache holds the last minted
+// token per DocmostClient for MCP_COLLAB_TOKEN_TTL_MS so a burst of mutations
+// reuses ONE token -> ONE session. These tests exercise both mint sources:
+//   - the getCollabToken PROVIDER path (in-app agent), via a counting provider fn;
+//   - the REST /auth/collab-token path (external MCP), via a mock http server.
+// getCollabTokenWithReauth is private in TS but a plain method on the compiled
+// build, so the tests call it directly (same convention as reauth.test.mjs).
+import { test, afterEach, after } from "node:test";
+import assert from "node:assert/strict";
+import http from "node:http";
+import { DocmostClient } from "../../build/client.js";
+
+// Restore the env knob after each test so cases do not leak into one another.
+const ENV_KEY = "MCP_COLLAB_TOKEN_TTL_MS";
+afterEach(() => {
+  delete process.env[ENV_KEY];
+});
+
+// ---------------------------------------------------------------------------
+// Small mock server for the REST /auth/collab-token path. Counts collab-token
+// mints and can be told to 401 the first N of them (to drive the reauth retry).
+// ---------------------------------------------------------------------------
+function readBody(req) {
+  return new Promise((resolve) => {
+    let raw = "";
+    req.on("data", (c) => (raw += c));
+    req.on("end", () => resolve(raw));
+  });
+}
+function sendJson(res, status, obj, extra = {}) {
+  res.writeHead(status, { "Content-Type": "application/json", ...extra });
+  res.end(JSON.stringify(obj));
+}
+
+const openServers = [];
+after(async () => {
+  await Promise.all(
+    openServers.map((s) => new Promise((r) => s.close(r))),
+  );
+});
+
+// state: { collabCalls, loginCalls, unauthorizedCollabHits }
+function spawnCollabServer(state, { collabAuthFailsFor = 0 } = {}) {
+  return new Promise((resolve) => {
+    const server = http.createServer(async (req, res) => {
+      await readBody(req);
+      if (req.url === "/api/auth/login") {
+        state.loginCalls++;
+        // A fresh authToken per login so an identity change is observable.
+        sendJson(res, 200, { success: true }, {
+          "Set-Cookie": `authToken=login-${state.loginCalls}; Path=/; HttpOnly`,
+        });
+        return;
+      }
+      if (req.url === "/api/auth/collab-token") {
+        state.collabCalls++;
+        if (state.collabCalls <= collabAuthFailsFor) {
+          sendJson(res, 401, { message: "Unauthorized" });
+          return;
+        }
+        // Unique token per mint so a stale cached value is distinguishable.
+        sendJson(res, 200, { data: { token: `collab-${state.collabCalls}` } });
+        return;
+      }
+      sendJson(res, 404, { message: "not found" });
+    });
+    server.listen(0, "127.0.0.1", () => {
+      openServers.push(server);
+      resolve(`http://127.0.0.1:${server.address().port}/api`);
+    });
+  });
+}
+
+// ===========================================================================
+// PROVIDER path (in-app agent getCollabToken fn)
+// ===========================================================================
+
+// A counting provider that returns a distinct token each call so a cached
+// (reused) token is visibly the SAME string while a fresh mint is different.
+function countingProvider() {
+  let n = 0;
+  const fn = async () => {
+    n++;
+    return `provider-token-${n}`;
+  };
+  return {
+    fn,
+    get calls() {
+      return n;
+    },
+  };
+}
+
+test("within TTL, repeated calls return the SAME token and mint ONCE (provider path)", async () => {
+  process.env[ENV_KEY] = "300000"; // 5 min
+  const p = countingProvider();
+  const client = new DocmostClient({
+    apiUrl: "http://127.0.0.1:1/api",
+    getToken: async () => "access",
+    getCollabToken: p.fn,
+  });
+
+  const a = await client.getCollabTokenWithReauth();
+  const b = await client.getCollabTokenWithReauth();
+  const c = await client.getCollabTokenWithReauth();
+
+  assert.equal(a, "provider-token-1");
+  assert.equal(b, a, "second call reuses the cached token");
+  assert.equal(c, a, "third call reuses the cached token");
+  assert.equal(p.calls, 1, "the provider is invoked exactly once within the TTL");
+});
+
+test("after TTL expiry a new token is minted (provider path)", async () => {
+  process.env[ENV_KEY] = "20"; // 20ms TTL
+  const p = countingProvider();
+  const client = new DocmostClient({
+    apiUrl: "http://127.0.0.1:1/api",
+    getToken: async () => "access",
+    getCollabToken: p.fn,
+  });
+
+  const a = await client.getCollabTokenWithReauth();
+  await new Promise((r) => setTimeout(r, 40)); // let the TTL lapse
+  const b = await client.getCollabTokenWithReauth();
+
+  assert.equal(a, "provider-token-1");
+  assert.equal(b, "provider-token-2", "a fresh token is minted after expiry");
+  assert.equal(p.calls, 2);
+});
+
+test("MCP_COLLAB_TOKEN_TTL_MS=0 disables the cache: mint on EVERY call (provider path)", async () => {
+  process.env[ENV_KEY] = "0";
+  const p = countingProvider();
+  const client = new DocmostClient({
+    apiUrl: "http://127.0.0.1:1/api",
+    getToken: async () => "access",
+    getCollabToken: p.fn,
+  });
+
+  await client.getCollabTokenWithReauth();
+  await client.getCollabTokenWithReauth();
+  await client.getCollabTokenWithReauth();
+
+  assert.equal(p.calls, 3, "cache disabled -> exact fetch-per-call legacy path");
+});
+
+test("a 401 triggers the internal reauth retry, which bypasses the cache and mints fresh (provider path)", async () => {
+  process.env[ENV_KEY] = "300000";
+  let n = 0;
+  const provider = async () => {
+    n++;
+    if (n === 1) {
+      // The FIRST mint fails with an auth error; the internal reauth retry must
+      // re-invoke the provider (bypassing the empty cache) for a fresh token.
+      const err = new Error("collab token expired");
+      err.status = 401;
+      throw err;
+    }
+    return `provider-token-${n}`;
+  };
+  const client = new DocmostClient({
+    apiUrl: "http://127.0.0.1:1/api",
+    getToken: async () => "access",
+    getCollabToken: provider,
+  });
+
+  // Cache is empty: mint #1 401s -> the reauth retry mints #2 and caches it.
+  const tok = await client.getCollabTokenWithReauth();
+  assert.equal(tok, "provider-token-2", "the post-401 retry token wins");
+  assert.equal(n, 2, "exactly one failed mint + one retry, no loop");
+
+  // The retried token is what got cached (no extra mint on a cache hit).
+  const cached = await client.getCollabTokenWithReauth();
+  assert.equal(cached, "provider-token-2");
+  assert.equal(n, 2, "served from cache, provider not re-invoked");
+});
+
+test("forceRefresh=true bypasses a warm cache and mints a fresh token (provider path)", async () => {
+  process.env[ENV_KEY] = "300000";
+  const p = countingProvider();
+  const client = new DocmostClient({
+    apiUrl: "http://127.0.0.1:1/api",
+    getToken: async () => "access",
+    getCollabToken: p.fn,
+  });
+
+  const first = await client.getCollabTokenWithReauth(); // caches token-1
+  assert.equal(first, "provider-token-1");
+
+  // A forced refresh (what the reauth path passes) must NOT return the cached
+  // token-1; it mints a fresh token-2 and replaces the cache.
+  const forced = await client.getCollabTokenWithReauth(true);
+  assert.equal(forced, "provider-token-2", "cache bypassed on forceRefresh");
+  assert.equal(p.calls, 2);
+
+  const cached = await client.getCollabTokenWithReauth();
+  assert.equal(cached, "provider-token-2", "the fresh token replaced the cache");
+  assert.equal(p.calls, 2);
+});
+
+test("two consecutive mutations keep the SAME token, so the session key is stable (provider path)", async () => {
+  // The whole point of #435: acquireCollabSession keys on the token, so two
+  // acquire calls in a burst must be handed the identical token string.
+  process.env[ENV_KEY] = "300000";
+  const p = countingProvider();
+  const client = new DocmostClient({
+    apiUrl: "http://127.0.0.1:1/api",
+    getToken: async () => "access",
+    getCollabToken: p.fn,
+  });
+
+  const t1 = await client.getCollabTokenWithReauth();
+  const t2 = await client.getCollabTokenWithReauth();
+  assert.equal(t1, t2, "identical token across two mutations -> one session key");
+  assert.equal(p.calls, 1);
+});
+
+// ===========================================================================
+// REST /auth/collab-token path (external MCP)
+// ===========================================================================
+
+test("within TTL, the REST /auth/collab-token endpoint is hit ONCE", async () => {
+  process.env[ENV_KEY] = "300000";
+  const state = { collabCalls: 0, loginCalls: 0 };
+  const baseURL = await spawnCollabServer(state);
+  const client = new DocmostClient(baseURL, "user@example.com", "pw");
+
+  const a = await client.getCollabTokenWithReauth();
+  const b = await client.getCollabTokenWithReauth();
+
+  assert.equal(a, "collab-1");
+  assert.equal(b, a, "cached token reused");
+  assert.equal(state.collabCalls, 1, "POST /auth/collab-token called once");
+});
+
+test("TTL=0 hits the REST endpoint on every call", async () => {
+  process.env[ENV_KEY] = "0";
+  const state = { collabCalls: 0, loginCalls: 0 };
+  const baseURL = await spawnCollabServer(state);
+  const client = new DocmostClient(baseURL, "user@example.com", "pw");
+
+  await client.getCollabTokenWithReauth();
+  await client.getCollabTokenWithReauth();
+
+  assert.equal(state.collabCalls, 2, "cache disabled -> fetch each call");
+});
+
+test("401 on REST collab-token re-logs-in and refetches (cache bypassed)", async () => {
+  process.env[ENV_KEY] = "300000";
+  const state = { collabCalls: 0, loginCalls: 0 };
+  // The first collab-token mint 401s; the reauth path logs in and retries.
+  const baseURL = await spawnCollabServer(state, { collabAuthFailsFor: 1 });
+  const client = new DocmostClient(baseURL, "user@example.com", "pw");
+  // Pre-seed a token so the initial call does not perform an initial login.
+  client.token = "seed";
+  client.client.defaults.headers.common["Authorization"] = "Bearer seed";
+
+  const tok = await client.getCollabTokenWithReauth();
+  assert.equal(tok, "collab-2", "the post-reauth mint wins, not the failed one");
+  assert.equal(state.loginCalls, 1, "re-login happened exactly once");
+  assert.equal(state.collabCalls, 2, "one failed mint + one successful retry");
+});
+
+test("a fresh login clears the cache so a collab token cannot outlive the identity", async () => {
+  process.env[ENV_KEY] = "300000";
+  const state = { collabCalls: 0, loginCalls: 0 };
+  const baseURL = await spawnCollabServer(state);
+  const client = new DocmostClient(baseURL, "user@example.com", "pw");
+
+  const before = await client.getCollabTokenWithReauth();
+  assert.equal(before, "collab-1");
+
+  // Simulate an identity change (the 401 interceptor / re-login path calls
+  // login(), which must drop the cached collab token).
+  await client.login();
+
+  const after = await client.getCollabTokenWithReauth();
+  assert.equal(after, "collab-2", "cache was invalidated by login(); refetched");
+  assert.equal(state.collabCalls, 2);
+});