Merge remote-tracking branch 'gitea/develop' into fix/541-reconnect-timeout

При живом обрыве SSE вход в лестницу reconnect происходил только
после резолва getRun(cid) (ре-сид из персиста). Путь REJECT
обрабатывался через .catch, но ЗАВИСШИЙ getRun (соединение есть,
ответа нет) был не ограничен: axios-клиент (lib/api-client.ts) не
имеет timeout, а stalled-idle-cap взводится только ПОСЛЕ входа в
reconnecting/polling. Итог — FSM залипала в `streaming` без баннера
и без поллинга до сокет-таймаута браузера (минуты). Это тот самый
класс тихого зависания, который устраняет эпик #497.

Оборачиваю ожидание getRun гонкой с таймаутом
(RECONNECT_RESEED_TIMEOUT_MS = 4s): по таймауту берётся ТОТ ЖЕ
фолбэк, что и на reject — dropLivePartialAndReplayFromStart() +
enterReconnect(runId), так что лестница/поллинг стартуют и
stalled-idle-cap взводится. Локальный флаг `settled` делает ветви
resolve/reject/timeout взаимоисключающими: поздний резолв getRun
после сработавшего таймаута полностью игнорируется (не входит в
reconnect повторно, не перетирает replay-from-start устаревшим
ре-сидом, ничего не перевзводит). Таймер живёт в reseedTimerRef и
очищается при размонтировании (никаких висящих setTimeout).

Тесты: hang-кейс (getRun не резолвится -> по таймауту FSM в
reconnecting, live-partial сброшен, replay-from-start) и
late-resolve safety (поздний резолв — no-op). Мутация: замена гонки
на голый getRun краснит оба #541-теста.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

.env.example

@@ -251,16 +251,6 @@ MCP_DOCMOST_PASSWORD=
 # Default 120000 (2 min).
 # AI_MCP_CALL_TIMEOUT_MS=120000
 
-# Kill-switch for the agent API-key feature (#501). Default ON when unset — a
-# deploy that never sets it must NOT silently kill every agent. STRICT parse:
-# only the literals `true` / `false` are accepted; a typo like `=0`/`=off`/`=False`
-# FAILS AT BOOT by design (never silently read as "enabled"), so the switch is
-# guaranteed to actually flip when an operator flips it during an incident. When
-# set to `false`: all api-key auth is DENIED (every api-key token is rejected) and
-# the api-key management endpoints return 404. The resolved state is logged at boot
-# (`API keys: ENABLED/DISABLED (API_KEYS_ENABLED=...)`) so it is verifiable per deploy.
-# API_KEYS_ENABLED=true
-
 # Max JSON/urlencoded request body size (bytes). Fastify's 1 MiB default is too
 # small for a long AI-chat research turn: the client resends the FULL message
 # history (every tool call + search result) on each turn, so a deep conversation's

.github/workflows/develop.yml

@@ -226,13 +226,6 @@ jobs:
       - name: Build mcp
         run: pnpm --filter @docmost/mcp build
 
-      # apps/server imports @docmost/token-estimate at runtime (history-budget.ts,
-      # #490); its dist/ is gitignored and `test:e2e` type-checks + runs the code,
-      # so build it here or tsc fails with TS2307 Cannot find module
-      # '@docmost/token-estimate' (mirrors the editor-ext / mcp build steps above).
-      - name: Build token-estimate
-        run: pnpm --filter @docmost/token-estimate build
-
       - name: Run migrations
         run: pnpm --filter ./apps/server migration:latest
 

.github/workflows/test.yml

@@ -159,14 +159,6 @@ jobs:
       - name: Build prosemirror-markdown
         run: pnpm --filter @docmost/prosemirror-markdown build
 
-      # @docmost/token-estimate is a shared workspace package the client vitest
-      # suite resolves via its dist build (main: ./dist/index.js); dist/ is
-      # gitignored and `pnpm -r test` does NOT honour nx `dependsOn: ^build`, so
-      # build it before the recursive test run or the client suite fails with
-      # "Failed to resolve import '@docmost/token-estimate'" (#490).
-      - name: Build token-estimate
-        run: pnpm --filter @docmost/token-estimate build
-
       - name: Run unit tests
         run: pnpm -r test
 

CHANGELOG.md

@@ -392,25 +392,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   `@docmost/token-estimate` package) so they can never diverge. Deferred-tool
   activation is also cached in the chat metadata to avoid re-resolving it each
   turn. (#490)
-- **Cyrillic (and any non-ASCII) draw.io labels no longer turn into mojibake
-  when a diagram is opened in the draw.io editor.** Agent-created diagrams
-  (`drawioCreate`) and Confluence-imported diagrams stored their model in the
-  SVG's `content=` attribute as base64; the draw.io editor decodes that via
-  Latin-1 `atob` (no UTF-8 step), so every non-ASCII char (e.g. `Старт-бит`,
-  `ё`, `—`) split into garbage and the editor's autosave then persisted the
-  corrupted model, breaking the page preview too. Both write paths
-  (`buildDrawioSvg`, the import service's `createDrawioSvg`) now write `content=`
-  as XML-entity-escaped mxfile XML — draw.io's own native form, decoded by the
-  DOM as UTF-8 — so labels open intact. The decoder reads both the new
-  entity-encoded form and the old base64 form, so existing diagrams still open.
-  *Healing pre-fix diagrams:* only a diagram that still holds its original
-  (correct-UTF-8) base64 — i.e. one not yet opened/autosaved in the draw.io
-  editor — can be repaired in place by `drawioGet` → `drawioUpdate` with the
-  same XML (rewrites the attachment in the new form); no migration script is
-  needed. A diagram that was already opened in the editor persisted the
-  mojibake at rest, so `drawioGet` reads the already-corrupted text and
-  `drawioUpdate` faithfully rewrites it — that text is lost and is not
-  recoverable by a rewrite. (#507)
 - **A chat with one malformed message part no longer 500s on every turn, and a
   failed send no longer duplicates the user's message.** Incoming client parts
   are now whitelisted to `text` (a forged tool-result part can no longer reach

Dockerfile

@@ -59,14 +59,6 @@ COPY --from=builder /app/packages/mcp/data /app/packages/mcp/data
 COPY --from=builder /app/packages/prosemirror-markdown/build /app/packages/prosemirror-markdown/build
 COPY --from=builder /app/packages/prosemirror-markdown/package.json /app/packages/prosemirror-markdown/package.json
 
-# apps/server imports @docmost/token-estimate (workspace:*) at runtime
-# (history-budget.ts, #490). tsc emits only dist/ and dist/ is gitignored, so the
-# prod install would resolve a broken workspace symlink and the server would die
-# with ERR_MODULE_NOT_FOUND on the first history-budget call. Ship the built
-# package + its manifest, mirroring prosemirror-markdown above.
-COPY --from=builder /app/packages/token-estimate/dist /app/packages/token-estimate/dist
-COPY --from=builder /app/packages/token-estimate/package.json /app/packages/token-estimate/package.json
-
 # Copy root package files
 COPY --from=builder /app/package.json /app/package.json
 COPY --from=builder /app/pnpm*.yaml /app/

apps/client/src/features/ai-chat/components/chat-thread.test.tsx

@@ -1254,4 +1254,88 @@ describe("ChatThread — live reconnect + stalled", () => {
     });
     expect(onResumeFallback).toHaveBeenCalledWith(false); // POLL_IDLE_CAP -> idle -> disarm
   });
+
+  it("#541: getRun HANGS on a live disconnect — the timeout race still enters reconnect (no silent freeze in `streaming`)", async () => {
+    // MUTATION-VERIFY: revert the race to a bare `getRun(cid).then/.catch` and this
+    // reddens — a HUNG (never-settling, NOT rejected) getRun leaves the FSM stuck in
+    // `streaming` with no reconnect banner and no poll (the axios client sets no
+    // request timeout, and the stalled-idle cap only arms AFTER reconnecting/polling).
+    renderLive();
+    h.state.getRun.mockReset();
+    h.state.getRun.mockReturnValue(new Promise(() => {})); // getRun HANGS forever
+    await disconnect(); // live partial = liveMsg (id "a2")
+    expect(h.state.getRun).toHaveBeenCalledWith("c1");
+    // BEFORE the bound fires the FSM is still in the live turn — the very freeze bug.
+    expect(screen.queryByText(/reconnecting/i)).toBeNull();
+    // The recovery-start bound fires -> the SAME fallback as the reject path.
+    act(() => {
+      vi.advanceTimersByTime(4_000);
+    });
+    expect(screen.getByText(/reconnecting/i)).toBeTruthy();
+    // The live partial (a2) was DROPPED from the store (replay-from-start, no stale
+    // tail-apply base). Mirrors the getRun-REJECT test's inspection.
+    const removedLivePartial = (
+      h.state.setMessages as unknown as {
+        mock: { calls: [unknown][] };
+      }
+    ).mock.calls.some(([updater]) => {
+      if (typeof updater !== "function") return false;
+      const out = (updater as (p: { id: string }[]) => { id: string }[])([
+        { id: "a2" },
+        { id: "u1" },
+      ]);
+      return !out.some((m) => m.id === "a2");
+    });
+    expect(removedLivePartial).toBe(true);
+    // ...and the anchor was NULLED -> replay-from-start (no ?anchor=&n= over the partial).
+    advanceToAttempt(1);
+    expect(h.state.resumeStream).toHaveBeenCalledTimes(1);
+    expect(h.state.transport!.prepareReconnectToStreamRequest!().api).toBe(
+      "/api/ai-chat/runs/c1/stream",
+    );
+  });
+
+  it("#541: a getRun resolve AFTER the timeout already fired is IGNORED (no double reconnect, no stale re-seed)", async () => {
+    // The timeout wins first and enters the ladder via replay-from-start. When the
+    // hung getRun FINALLY answers, its late `.then` must be a full no-op: it must not
+    // re-seed the store from the (now stale) persisted row, must not re-set the
+    // anchor, and must not re-enter reconnect. The local `settled` flag makes the
+    // resolve/reject/timeout branches mutually exclusive.
+    // MUTATION-VERIFY: drop the `settled` guard (let the late `.then` run) and the
+    // late re-seed re-sets the anchor (URL regains ?anchor=a2&n=3) + calls setMessages.
+    renderLive();
+    let resolveGetRun!: (v: unknown) => void;
+    h.state.getRun.mockReset();
+    h.state.getRun.mockReturnValue(
+      new Promise((r) => {
+        resolveGetRun = r;
+      }),
+    );
+    await disconnect();
+    act(() => {
+      vi.advanceTimersByTime(4_000); // bound fires -> replay-from-start, reconnecting
+    });
+    expect(screen.getByText(/reconnecting/i)).toBeTruthy();
+    advanceToAttempt(1);
+    expect(h.state.resumeStream).toHaveBeenCalledTimes(1);
+    // Anchor is null -> replay-from-start URL (the pre-condition the late resolve must
+    // not undo).
+    expect(h.state.transport!.prepareReconnectToStreamRequest!().api).toBe(
+      "/api/ai-chat/runs/c1/stream",
+    );
+    const setMessagesCallsBefore = h.state.setMessages.mock.calls.length;
+    // NOW the hung getRun finally resolves with a persisted anchor (id a2, steps 3).
+    await act(async () => {
+      resolveGetRun(persistedAnchor());
+      await Promise.resolve();
+    });
+    // The late resolve did NOT re-seed the store...
+    expect(h.state.setMessages.mock.calls.length).toBe(setMessagesCallsBefore);
+    // ...did NOT re-set the anchor (URL stays replay-from-start, no ?anchor=&n=)...
+    expect(h.state.transport!.prepareReconnectToStreamRequest!().api).toBe(
+      "/api/ai-chat/runs/c1/stream",
+    );
+    // ...and did NOT trigger a fresh reconnect attach.
+    expect(h.state.resumeStream).toHaveBeenCalledTimes(1);
+  });
 });

apps/client/src/features/ai-chat/components/chat-thread.tsx

@@ -81,6 +81,17 @@ const STREAM_THROTTLE_MS = 50;
 // THREAD now (the FSM owns polling->stalled); the window just polls while armed.
 const DEGRADED_POLL_IDLE_MAX_MS = 10 * 60_000;
 
+// #541: the bound on the persist re-seed (getRun) wait when ENTERING the reconnect
+// ladder after a live SSE drop. The axios client (lib/api-client.ts) sets NO request
+// timeout, and the stalled-idle cap only arms AFTER the FSM enters reconnecting/
+// polling — which never happens if getRun HANGS (connection established, no response).
+// Without this bound a live drop + a hung getRun sticks the FSM in `streaming` with
+// no banner and no poll until the browser socket timeout (minutes) — the silent-freeze
+// class #497 eliminates. This is a deliberate RECOVERY-START bound (drop the live
+// partial + enter the ladder so the poll/idle-cap arms), intentionally much shorter
+// than any network socket timeout — not a network read timeout.
+const RECONNECT_RESEED_TIMEOUT_MS = 4_000;
+
 /** The #487 active (non-terminal) run statuses — mirrors the server's
  *  ACTIVE_RUN_STATUSES. A run-fact is "active" only for these. */
 function isActiveRunStatus(status: string | null | undefined): boolean {
@@ -286,6 +297,9 @@ export default function ChatThread({
   // stalled inactivity cap. Cleared by the cancelReconnect effect / the cap effect.
   const reconnectTimerRef = useRef<ReturnType<typeof setTimeout> | null>(null);
   const idleCapTimerRef = useRef<ReturnType<typeof setTimeout> | null>(null);
+  // #541: the persist re-seed (getRun) timeout race on the disconnect->reconnect
+  // entry. Held in a ref so unmount (DISPOSE cleanup) clears it — no dangling timer.
+  const reseedTimerRef = useRef<ReturnType<typeof setTimeout> | null>(null);
 
   // #491 tail-only: seed EVERY persisted row unchanged (no strip). The streaming
   // tail holds steps 0..N-1; the run-stream registry's tail (steps >= N) is APPENDED
@@ -748,37 +762,68 @@ export default function ChatThread({
           anchorRef.current = null;
         };
         if (cid) {
-          void getRun(cid)
-            .then((res) => {
-              if (!mountedRef.current) return;
-              const persisted = res.message;
-              if (persisted && persisted.role === "assistant") {
-                anchorRef.current = {
-                  id: persisted.id,
-                  stepsPersisted: stepsPersistedOf(persisted),
-                };
-                // Replace the live partial with the persisted row IN PLACE by id —
-                // the re-seed from persist. The attach's tail (steps >= N) then
-                // appends to a store holding EXACTLY steps 0..N-1: no duplication.
-                setMessages((prev) => mergeById(prev, rowToUiMessage(persisted)));
-              } else {
-                // No persisted assistant row (pre-first-frame break): drop the live
-                // partial + replay from start (no anchor/n) so nothing is duplicated.
-                dropLivePartialAndReplayFromStart();
-              }
-              enterReconnect(res.run?.id ?? runId);
-            })
-            .catch(() => {
-              if (!mountedRef.current) return;
-              // Persist read FAILED: we cannot re-seed from fresh persist, and a
-              // stale mount-time anchor over the live partial would tail-apply
-              // already-present steps -> duplication (a flaky-network blip:
-              // SSE + getRun both fail, network recovers in ~1s, the registry still
-              // covers from the mount frontier). Restore the removed-filter guarantee
-              // instead: drop the live partial + replay from start / 204 -> poll.
+          // #541: bound the persist re-seed wait with a timeout race. getRun goes
+          // through the axios client, which has NO request timeout; a HUNG getRun
+          // (connection open, no response) — distinct from a REJECT, which the
+          // `.catch` already handles — would otherwise never let us enter the ladder,
+          // leaving the FSM stuck in `streaming` with no banner and no poll until the
+          // browser socket timeout. `settled` makes the three branches (resolve /
+          // reject / timeout) mutually exclusive: whichever fires FIRST wins and the
+          // others become no-ops. So a LATE getRun resolve AFTER the timeout is fully
+          // ignored — it cannot (i) re-enter reconnect a second time, (ii) overwrite
+          // the timeout's replay-from-start with a stale re-seed, or (iii) re-arm any
+          // timer. On timeout we take the SAME fallback as the reject path (drop the
+          // live partial + enter the ladder, so the poll and the stalled-idle cap arm).
+          let settled = false;
+          const finishReseed = (apply: () => void): void => {
+            if (settled) return;
+            settled = true;
+            if (reseedTimerRef.current) {
+              clearTimeout(reseedTimerRef.current);
+              reseedTimerRef.current = null;
+            }
+            if (!mountedRef.current) return;
+            apply();
+          };
+          reseedTimerRef.current = setTimeout(() => {
+            finishReseed(() => {
               dropLivePartialAndReplayFromStart();
               enterReconnect(runId);
             });
+          }, RECONNECT_RESEED_TIMEOUT_MS);
+          void getRun(cid)
+            .then((res) => {
+              finishReseed(() => {
+                const persisted = res.message;
+                if (persisted && persisted.role === "assistant") {
+                  anchorRef.current = {
+                    id: persisted.id,
+                    stepsPersisted: stepsPersistedOf(persisted),
+                  };
+                  // Replace the live partial with the persisted row IN PLACE by id —
+                  // the re-seed from persist. The attach's tail (steps >= N) then
+                  // appends to a store holding EXACTLY steps 0..N-1: no duplication.
+                  setMessages((prev) => mergeById(prev, rowToUiMessage(persisted)));
+                } else {
+                  // No persisted assistant row (pre-first-frame break): drop the live
+                  // partial + replay from start (no anchor/n) so nothing is duplicated.
+                  dropLivePartialAndReplayFromStart();
+                }
+                enterReconnect(res.run?.id ?? runId);
+              });
+            })
+            .catch(() => {
+              finishReseed(() => {
+                // Persist read FAILED: we cannot re-seed from fresh persist, and a
+                // stale mount-time anchor over the live partial would tail-apply
+                // already-present steps -> duplication (a flaky-network blip:
+                // SSE + getRun both fail, network recovers in ~1s, the registry still
+                // covers from the mount frontier). Restore the removed-filter guarantee
+                // instead: drop the live partial + replay from start / 204 -> poll.
+                dropLivePartialAndReplayFromStart();
+                enterReconnect(runId);
+              });
+            });
         } else {
           dropLivePartialAndReplayFromStart();
           enterReconnect(runId);
@@ -903,6 +948,12 @@ export default function ChatThread({
     }
     return () => {
       mountedRef.current = false;
+      // #541: clear the in-flight persist re-seed timeout (not an FSM effect timer,
+      // so DISPOSE does not touch it) — no dangling setTimeout after unmount.
+      if (reseedTimerRef.current) {
+        clearTimeout(reseedTimerRef.current);
+        reseedTimerRef.current = null;
+      }
       dispatch({ type: "DISPOSE" }); // aborts attach + timers, bumps epoch (I5)
     };
     // Mount-only by design; the parent remounts per chat via `key`.

apps/client/src/features/ai-chat/components/message-item-memo.test.tsx

@@ -27,7 +27,6 @@ vi.mock("@/features/ai-chat/utils/markdown.ts", async () => {
 
 import MessageItem from "./message-item";
 import { messageSignature } from "@/features/ai-chat/utils/message-signature.ts";
-import { splitPlainChunks } from "./streaming-plain-text";
 
 // matchMedia (read by MantineProvider) is stubbed globally in vitest.setup.ts.
 
@@ -115,89 +114,3 @@ describe("MessageItem markdown memoization", () => {
     expect(queryByText("streamed answer")).not.toBeNull();
   });
 });
-
-// PERF SMOKE (#492): the whole point of the incremental streaming render is that
-// the ANSWER path costs O(number of markdown blocks), NOT O(number of throttled
-// ~20Hz ticks). Pre-#492 the finalized MarkdownPart re-parsed the WHOLE growing
-// answer on every delta — a synthetic ~100 KB stream measured 394 renderChatMarkdown
-// calls (one per tick). With the incremental render each STABILIZED block is parsed
-// exactly once (memoized in MarkdownChunk) and the live tail is cheap plain text, so
-// the call count collapses to ~= the block count regardless of tick granularity.
-describe("MessageItem streaming answer render is O(blocks), not O(ticks)", () => {
-  // ~100 KB answer. Each section is a heading + a paragraph — TWO blank-line
-  // delimited markdown blocks — so the safe-cut block count is ~2× the section
-  // count. The perf claim is about the BLOCK count (the memoization granularity),
-  // measured directly with splitPlainChunks below, not the section count.
-  const buildAnswer = () => {
-    const SECTIONS = 100;
-    const paragraphs: string[] = [];
-    for (let i = 0; i < SECTIONS; i++) {
-      paragraphs.push(`## Section ${i}\n\n` + "lorem ipsum dolor ".repeat(55));
-    }
-    const full = paragraphs.join("\n\n");
-    // The number of memoized markdown blocks the incremental render splits into
-    // (all but the live tail are parsed once each).
-    return { full, blocks: splitPlainChunks(full).length };
-  };
-
-  const streamMsg = (text: string, state: "streaming" | "done"): UIMessage =>
-    ({
-      id: "m1",
-      role: "assistant",
-      parts: [{ type: "text", text, state }],
-    }) as UIMessage;
-
-  it("parses each block ~once over a 100KB stream (≈blocks, ≪ ticks)", () => {
-    renderChatMarkdownSpy.mockClear();
-    const { full, blocks } = buildAnswer();
-    const CHUNK = 128; // a realistic ~20Hz throttled delta size
-    const ticks = Math.ceil(full.length / CHUNK);
-
-    let msg = streamMsg(full.slice(0, CHUNK), "streaming");
-    const { rerender } = render(
-      <MantineProvider>
-        <MessageItem
-          message={msg}
-          signature={messageSignature(msg)}
-          turnStreaming
-        />
-      </MantineProvider>,
-    );
-    for (let end = 2 * CHUNK; end < full.length; end += CHUNK) {
-      msg = streamMsg(full.slice(0, end), "streaming");
-      rerender(
-        <MantineProvider>
-          <MessageItem
-            message={msg}
-            signature={messageSignature(msg)}
-            turnStreaming
-          />
-        </MantineProvider>,
-      );
-    }
-    // Finalize: the streaming→done flip renders the whole answer through ONE
-    // canonical pass (visual parity), so the finished DOM matches the pre-#492
-    // output. This is the single extra parse on top of the per-block ones.
-    const done = streamMsg(full, "done");
-    rerender(
-      <MantineProvider>
-        <MessageItem message={done} signature={messageSignature(done)} />
-      </MantineProvider>,
-    );
-
-    const calls = renderChatMarkdownSpy.mock.calls.length;
-    // Sanity: the stream really had far more ticks than blocks (else the test is
-    // vacuous — the point is that calls scale with blocks, not ticks).
-    expect(ticks).toBeGreaterThan(blocks * 3);
-    // O(blocks): each stabilized block parsed once + the single final whole-text
-    // parse. A small constant absorbs the finalize render and the live-tail block;
-    // the load-bearing claim is the bound below.
-    expect(calls).toBeLessThanOrEqual(blocks + 2);
-    // ≪ ticks — and, non-vacuously, the blocks WERE parsed (not skipped entirely).
-    expect(calls).toBeLessThan(ticks / 3);
-    expect(calls).toBeGreaterThan(blocks / 2);
-    // MUTATION-VERIFY (documented, not run here): dropping the `memo()` wrapper on
-    // MarkdownChunk (so every stable block re-parses each tick) drives `calls`
-    // toward `ticks` (~394), reddening both upper-bound assertions above.
-  });
-});

apps/client/src/features/ai-chat/components/message-item.render.test.tsx

@@ -1,112 +0,0 @@
-import { describe, expect, it, vi } from "vitest";
-import { render } from "@testing-library/react";
-import { MantineProvider } from "@mantine/core";
-import type { UIMessage } from "@ai-sdk/react";
-
-// Stub react-i18next (the component reads `useTranslation`). Mirrors the other
-// message-item specs.
-vi.mock("react-i18next", () => ({
-  useTranslation: () => ({ t: (key: string) => key }),
-}));
-
-import MessageItem from "./message-item";
-import { messageSignature } from "@/features/ai-chat/utils/message-signature.ts";
-// The REAL canonical renderer (NOT the spy the memo test installs): this file
-// exercises the actual markdown output so the visual-regression assertions below
-// compare against genuine HTML (incl. the schema's `<li><p>` wrappers).
-import { renderChatMarkdown } from "@/features/ai-chat/utils/markdown.ts";
-import classes from "./ai-chat.module.css";
-
-const msg = (
-  parts: UIMessage["parts"],
-  extra?: Partial<UIMessage>,
-): UIMessage =>
-  ({ id: "m1", role: "assistant", parts, ...extra }) as UIMessage;
-
-const renderRow = (message: UIMessage, turnStreaming = false) =>
-  render(
-    <MantineProvider>
-      <MessageItem
-        message={message}
-        signature={messageSignature(message)}
-        turnStreaming={turnStreaming}
-      />
-    </MantineProvider>,
-  );
-
-// A rich multi-block answer that exercises headings, a list (the `<li><p>` case
-// the scoped CSS tightens), inline emphasis, and multiple paragraphs.
-const ANSWER = [
-  "# Заголовок",
-  "",
-  "Первый абзац с **жирным** и `кодом`.",
-  "",
-  "- пункт один",
-  "- пункт два",
-  "",
-  "Второй абзац.",
-].join("\n");
-
-describe("MessageItem final render — visual parity with the canonical pipeline", () => {
-  it("a finalized text part renders exactly renderChatMarkdown(text)", () => {
-    const { container } = renderRow(
-      msg([{ type: "text", text: ANSWER, state: "done" }]),
-    );
-    const block = container.querySelector(`.${classes.markdown}`);
-    expect(block).not.toBeNull();
-    // Byte-for-byte the canonical output (the SAME whole-text pass the pre-#492
-    // MarkdownPart produced), including `<li><p>…</p></li>` wrappers.
-    expect(block!.innerHTML).toBe(renderChatMarkdown(ANSWER, {}));
-    // The list wrapper is really present (guards against a vacuous empty render).
-    expect(container.querySelectorAll("li p").length).toBe(2);
-  });
-
-  it("the streaming incremental view CONVERGES to the canonical render on finish", () => {
-    // Mount mid-stream (live tail) — the DOM here is the incremental view.
-    const { container, rerender } = render(
-      <MantineProvider>
-        <MessageItem
-          message={msg([{ type: "text", text: ANSWER, state: "streaming" }])}
-          signature={messageSignature(
-            msg([{ type: "text", text: ANSWER, state: "streaming" }]),
-          )}
-          turnStreaming
-        />
-      </MantineProvider>,
-    );
-    // Finish the turn: state flips to done AND the turn is no longer streaming.
-    const done = msg([{ type: "text", text: ANSWER, state: "done" }]);
-    rerender(
-      <MantineProvider>
-        <MessageItem message={done} signature={messageSignature(done)} />
-      </MantineProvider>,
-    );
-    // After finish there is exactly ONE canonical markdown container whose HTML is
-    // the whole-text render — identical to the non-streaming path above.
-    const blocks = container.querySelectorAll(`.${classes.markdown}`);
-    expect(blocks.length).toBe(1);
-    expect(blocks[0].innerHTML).toBe(renderChatMarkdown(ANSWER, {}));
-  });
-
-  it("neutralizeInternalLinks is honored on the finalized render", () => {
-    const linkAnswer = "См. [страницу](/p/abc).";
-    const { container } = render(
-      <MantineProvider>
-        <MessageItem
-          message={msg([{ type: "text", text: linkAnswer, state: "done" }])}
-          signature={messageSignature(
-            msg([{ type: "text", text: linkAnswer, state: "done" }]),
-          )}
-          neutralizeInternalLinks
-        />
-      </MantineProvider>,
-    );
-    const block = container.querySelector(`.${classes.markdown}`);
-    expect(block!.innerHTML).toBe(
-      renderChatMarkdown(linkAnswer, { neutralizeInternalLinks: true }),
-    );
-    // The internal link was made inert (no href) by the neutralization flag.
-    const a = container.querySelector("a");
-    expect(a?.hasAttribute("href")).toBe(false);
-  });
-});

apps/client/src/features/ai-chat/components/message-item.tsx

@@ -4,7 +4,6 @@ import { useTranslation } from "react-i18next";
 import type { UIMessage } from "@ai-sdk/react";
 import ToolCallCard from "@/features/ai-chat/components/tool-call-card.tsx";
 import ReasoningBlock from "@/features/ai-chat/components/reasoning-block.tsx";
-import { StreamingMarkdownText } from "@/features/ai-chat/components/streaming-markdown-text.tsx";
 import ChatErrorAlert from "@/features/ai-chat/components/chat-error-alert.tsx";
 import ChatStoppedNotice from "@/features/ai-chat/components/chat-stopped-notice.tsx";
 import { ToolUiPart, isToolPart } from "@/features/ai-chat/utils/tool-parts.tsx";
@@ -87,39 +86,17 @@ interface MessageItemProps {
  * One assistant text part rendered as sanitized markdown. Memoized on its inputs
  * so a finalized text part is NOT re-parsed on every streamed delta: during a
  * turn only the actively-growing tail part changes its `text`, so every earlier
- * part hits the memo and skips the expensive canonical parse + DOMPurify pass.
- * Props are primitives, so React.memo's default shallow compare is exactly right
- * (the `text` string is compared by value).
- *
- * Streaming gate (#492) — mirrors ReasoningBlock:
- *  - `streaming` (this is the live, actively-growing tail part of an in-flight
- *    turn): render incrementally via StreamingMarkdownText — the stabilized blocks
- *    go through the canonical pipeline (each parsed ONCE, memoized) and only the
- *    live tail is cheap plain text. This makes the per-tick cost O(new blocks),
- *    not the pre-#492 O(ticks) whole-answer re-parse on every ~20Hz delta.
- *  - finalized (the common case, and the turn-end flip): render the WHOLE text
- *    through ONE canonical pass — byte-identical to the pre-#492 output (visual
- *    parity). The row re-renders on the streaming→done flip because
- *    `messageSignature` tracks each part's `state` (and `turnStreaming` flips at
- *    turn end), so the incremental view always converges to this single render.
+ * part hits the memo and skips the expensive marked + DOMPurify pass. Props are
+ * primitives, so React.memo's default shallow compare is exactly right (the
+ * `text` string is compared by value).
  */
 const MarkdownPart = memo(function MarkdownPart({
   text,
   neutralizeInternalLinks,
-  streaming,
 }: {
   text: string;
   neutralizeInternalLinks: boolean;
-  streaming: boolean;
 }) {
-  if (streaming) {
-    return (
-      <StreamingMarkdownText
-        text={text}
-        neutralizeInternalLinks={neutralizeInternalLinks}
-      />
-    );
-  }
   const html = renderChatMarkdown(text, { neutralizeInternalLinks });
   if (html) {
     return (
@@ -202,10 +179,47 @@ function MessageItem({
         {resolveAssistantName(assistantName) ?? t("AI agent")}
       </Text>
       {message.parts.map((part, index) => {
-        // Tool parts (`tool-*` / `dynamic-tool`) are template-literal kinds, so
-        // they cannot be a `switch` case; the runtime guard handles them, and the
-        // switch below covers every CLOSED (literal-typed) part kind with a
-        // compile-time exhaustiveness check in its default.
+        if (part.type === "reasoning") {
+          // Reasoning ("thinking") -> a collapsible block with its own token
+          // count. Empty/whitespace reasoning with no authoritative count carries
+          // nothing to show, so skip it (avoids an empty 0-token block).
+          const text = (part as { text?: string }).text ?? "";
+          if (!text.trim() && !(reasoningTokens && reasoningTokens > 0))
+            return null;
+          // Absent state (persisted rows) and "done" both mean finalized.
+          // `messageSignature` already includes each part's `state`, so the
+          // streaming→done flip changes the row signature and re-renders this
+          // row — which is what lets ReasoningBlock switch from chunked plain
+          // text to its one-time markdown parse (see reasoning-block.tsx).
+          // ALSO require the turn to be live: a part stranded at
+          // `state:"streaming"` after the turn ended (no `reasoning-end` — see
+          // the `turnStreaming` prop doc) must still finalize and parse.
+          const streaming =
+            turnStreaming && (part as { state?: string }).state === "streaming";
+          return (
+            <ReasoningBlock
+              key={index}
+              text={text}
+              tokens={reasoningTokens}
+              streaming={streaming}
+            />
+          );
+        }
+
+        if (part.type === "text") {
+          // Skip empty/whitespace-only text parts (a streaming message often
+          // starts with an empty text part before the first token arrives); the
+          // typing indicator covers that gap until real content streams in.
+          if (!part.text.trim()) return null;
+          return (
+            <MarkdownPart
+              key={index}
+              text={part.text}
+              neutralizeInternalLinks={neutralizeInternalLinks}
+            />
+          );
+        }
+
         if (isToolPart(part.type)) {
           return (
             <ToolCallCard
@@ -218,76 +232,7 @@ function MessageItem({
           );
         }
 
-        switch (part.type) {
-          case "reasoning": {
-            // Reasoning ("thinking") -> a collapsible block with its own token
-            // count. Empty/whitespace reasoning with no authoritative count
-            // carries nothing to show, so skip it (avoids an empty 0-token block).
-            const text = part.text ?? "";
-            if (!text.trim() && !(reasoningTokens && reasoningTokens > 0))
-              return null;
-            // Absent state (persisted rows) and "done" both mean finalized.
-            // `messageSignature` already includes each part's `state`, so the
-            // streaming→done flip changes the row signature and re-renders this
-            // row — which is what lets ReasoningBlock switch from chunked plain
-            // text to its one-time markdown parse (see reasoning-block.tsx).
-            // ALSO require the turn to be live: a part stranded at
-            // `state:"streaming"` after the turn ended (no `reasoning-end` — see
-            // the `turnStreaming` prop doc) must still finalize and parse.
-            const streaming = turnStreaming && part.state === "streaming";
-            return (
-              <ReasoningBlock
-                key={index}
-                text={text}
-                tokens={reasoningTokens}
-                streaming={streaming}
-              />
-            );
-          }
-
-          case "text": {
-            // Skip empty/whitespace-only text parts (a streaming message often
-            // starts with an empty text part before the first token arrives); the
-            // typing indicator covers that gap until real content streams in.
-            if (!part.text.trim()) return null;
-            // The live, actively-growing tail part of the in-flight turn renders
-            // incrementally (see MarkdownPart); a finalized part (persisted, or
-            // the turn-end flip) renders the whole text through one canonical
-            // pass. Same liveness rule as the reasoning branch above.
-            const streaming = turnStreaming && part.state === "streaming";
-            return (
-              <MarkdownPart
-                key={index}
-                text={part.text}
-                neutralizeInternalLinks={neutralizeInternalLinks}
-                streaming={streaming}
-              />
-            );
-          }
-
-          case "source-url":
-          case "source-document":
-          case "file":
-          case "step-start":
-            // Not surfaced in the chat bubble (v1) — same as the pre-#492 default.
-            return null;
-
-          default: {
-            // Compile-time exhaustiveness over the CLOSED union members: every
-            // literal-typed part kind is handled above, so the only kinds that
-            // can reach here are the OPEN template-literal ones (`tool-*` — caught
-            // by the guard at runtime — and `data-*`) plus `dynamic-tool`. Adding
-            // a NEW closed part kind to UIMessagePart makes this assignment fail
-            // to compile, forcing it to be handled instead of silently ignored
-            // (this replaces the pre-#492 fall-through `return null` + WARNING).
-            const _exhaustive:
-              | `tool-${string}`
-              | "dynamic-tool"
-              | `data-${string}` = part.type;
-            void _exhaustive;
-            return null;
-          }
-        }
+        return null;
       })}
       {/* A persisted turn error (server stored it in metadata.error). Rendered
           here so it survives a thread remount and shows in reopened history. */}

apps/client/src/features/ai-chat/components/streaming-markdown-text.tsx

@@ -1,96 +0,0 @@
-import { memo, useMemo } from "react";
-import { splitPlainChunks } from "@/features/ai-chat/components/streaming-plain-text.tsx";
-import { renderChatMarkdown } from "@/features/ai-chat/utils/markdown.ts";
-import classes from "@/features/ai-chat/components/ai-chat.module.css";
-
-/**
- * One STABILIZED markdown block, rendered through the canonical pipeline and
- * memoized on its string prop. During streaming only the TAIL chunk grows (the
- * `splitPlainChunks` append-only invariant guarantees every earlier chunk is
- * byte-identical across deltas), so React skips every stable block and each one
- * is parsed by `renderChatMarkdown` EXACTLY ONCE — turning the pre-#492
- * "re-parse the whole accumulated answer on every ~20Hz tick" (O(ticks)) into
- * O(number of blocks). The markup is DOMPurify-sanitized inside renderChatMarkdown
- * before it reaches `dangerouslySetInnerHTML`.
- *
- * NOTE (transient streaming-only artifact): a safe cut is a blank-line boundary,
- * so a construct that legitimately contains a blank line (e.g. a fenced code block
- * with an empty line) can be split across chunks and render oddly WHILE it is still
- * streaming. This is cosmetic and self-heals: the moment the part finalizes,
- * MarkdownPart renders the WHOLE text through one canonical pass (visual parity
- * with the pre-#492 output). The reasoning path makes the same trade (plain text
- * while streaming, one markdown parse at the end).
- */
-const MarkdownChunk = memo(function MarkdownChunk({
-  text,
-  neutralizeInternalLinks,
-}: {
-  text: string;
-  neutralizeInternalLinks: boolean;
-}) {
-  const html = renderChatMarkdown(text, { neutralizeInternalLinks });
-  if (html) {
-    return (
-      <div
-        className={classes.markdown}
-        // Sanitized by renderChatMarkdown (DOMPurify) before insertion.
-        dangerouslySetInnerHTML={{ __html: html }}
-      />
-    );
-  }
-  // Malformed/unsupported markdown could not render synchronously: raw text.
-  return (
-    <div className={classes.markdown} style={{ whiteSpace: "pre-wrap" }}>
-      {text}
-    </div>
-  );
-});
-
-/**
- * The cheap streaming-time stand-in for the finalized answer's one-time markdown
- * parse (see MarkdownPart in message-item.tsx). Mirrors StreamingPlainText's
- * chunked-memo pattern but renders the STABILIZED prefix as real markdown (each
- * block parsed once, memoized) and only the LIVE tail as flat plain text — so the
- * user sees formatted output for everything up to the last safe cut, and the not-
- * yet-stable tail (which markdown-parsing every tick would make O(ticks)) stays a
- * single cheap escaped text node until it stabilizes into a new block.
- *
- * `splitPlainChunks` yields chunks where, under append-only growth, every chunk
- * except the LAST is immutable; the last chunk is the live tail. Index keys are
- * therefore stable (a given index never changes to a different chunk's content).
- */
-export function StreamingMarkdownText({
-  text,
-  neutralizeInternalLinks,
-}: {
-  text: string;
-  neutralizeInternalLinks: boolean;
-}) {
-  const chunks = useMemo(() => splitPlainChunks(text), [text]);
-  return (
-    <>
-      {chunks.map((chunk, index) =>
-        index < chunks.length - 1 ? (
-          <MarkdownChunk
-            key={index}
-            text={chunk}
-            neutralizeInternalLinks={neutralizeInternalLinks}
-          />
-        ) : (
-          // The live tail: flat, React-escaped plain text (no markdown parse, no
-          // sanitizer, no innerHTML). `pre-wrap` preserves its newlines; trailing
-          // separator newlines are dropped at display time so the block gap comes
-          // from the markdown margins, not a doubled empty line (mirrors
-          // PlainChunk in streaming-plain-text.tsx).
-          <div
-            key={index}
-            className={classes.markdown}
-            style={{ whiteSpace: "pre-wrap" }}
-          >
-            {chunk.replace(/\n+$/, "")}
-          </div>
-        ),
-      )}
-    </>
-  );
-}

apps/client/src/features/page/tree/components/space-tree.tsx

@@ -281,12 +281,10 @@ const SpaceTree = forwardRef<SpaceTreeApi, SpaceTreeProps>(function SpaceTree(
       setOpenTreeNodes((prev) => ({ ...prev, [id]: isOpen }));
       if (isOpen) {
         const node = treeModel.find(data, id) as SpaceTreeNode | null;
-        // Same "unloaded branch" predicate the realtime insert paths use
-        // (`isUnloadedBranch`) so the lazy-load gate and the realtime inserts
-        // (`insertByPosition` / `placeByPosition`) can never disagree about what
-        // counts as unloaded (#525). Note: local raw `insert` (DnD/create-page)
-        // does not yet route through it — see #525 follow-up.
-        if (treeModel.isUnloadedBranch(node)) {
+        if (
+          node?.hasChildren &&
+          (!node.children || node.children.length === 0)
+        ) {
           const fetched = await fetchAllAncestorChildren({
             pageId: id,
             spaceId: node.spaceId,

apps/client/src/features/page/tree/model/tree-model.test.ts

@@ -74,48 +74,6 @@ describe("treeModel.isDescendant", () => {
   });
 });
 
-// #525: the single "is this branch unloaded?" predicate shared by the lazy-load
-// gate and the insert paths. Unloaded == server says hasChildren but none are
-// present locally (canonical form `children: []`, also `undefined`). A parent
-// without hasChildren is genuinely empty, not unloaded.
-describe("treeModel.isUnloadedBranch", () => {
-  type PH = TreeNode<{ name: string; hasChildren?: boolean }>;
-  it("true for hasChildren + empty array (canonical unloaded form)", () => {
-    const n: PH = { id: "p", name: "P", hasChildren: true, children: [] };
-    expect(treeModel.isUnloadedBranch(n)).toBe(true);
-  });
-  it("true for hasChildren + undefined children", () => {
-    const n: PH = { id: "p", name: "P", hasChildren: true };
-    expect(treeModel.isUnloadedBranch(n)).toBe(true);
-  });
-  it("false for hasChildren + already-loaded children", () => {
-    const n: PH = {
-      id: "p",
-      name: "P",
-      hasChildren: true,
-      children: [{ id: "c", name: "C" }],
-    };
-    expect(treeModel.isUnloadedBranch(n)).toBe(false);
-  });
-  it("false for a genuinely-empty parent (no hasChildren)", () => {
-    expect(
-      treeModel.isUnloadedBranch({
-        id: "p",
-        name: "P",
-        hasChildren: false,
-        children: [],
-      } as PH),
-    ).toBe(false);
-    expect(
-      treeModel.isUnloadedBranch({ id: "p", name: "P" } as PH),
-    ).toBe(false);
-  });
-  it("false for null/undefined", () => {
-    expect(treeModel.isUnloadedBranch(null)).toBe(false);
-    expect(treeModel.isUnloadedBranch(undefined)).toBe(false);
-  });
-});
-
 describe("treeModel.visible", () => {
   it("returns only root nodes when no openIds", () => {
     const v = treeModel.visible(fixture, new Set());
@@ -239,64 +197,43 @@ describe("treeModel.insertByPosition", () => {
     ]);
   });
 
-  type PH = TreeNode<{
-    name: string;
-    position?: string;
-    hasChildren?: boolean;
-  }>;
-
-  // #159 #1 / #525: inserting/moving a node under an UNLOADED parent must NOT
-  // materialize a partial `[node]` list — that would defeat the lazy-load gate and
-  // hide the parent's other real children. The canonical unloaded form here is
-  // `children: []` + `hasChildren: true` (from `pageToTreeNode` /
-  // `pruneCollapsedChildren`), which the pre-#525 `=== undefined` guard MISSED.
-  // The node is left to be lazy-loaded; the chevron stays enabled.
-  it("does NOT materialize a child under an UNLOADED parent (children: [], hasChildren: true)", () => {
-    const tree: PH[] = [
-      { id: "p", name: "P", position: "a0", hasChildren: true, children: [] },
-    ];
-    const node: PH = { id: "x", name: "X", position: "a1" };
-    const t = treeModel.insertByPosition(tree, "p", node);
-    const parent = treeModel.find(t, "p");
-    // The node was NOT inserted (children stay unloaded -> lazy-load fetches the
-    // full set, including this node, on expand). MUTATION: the pre-#525 predicate
-    // `children === undefined` does not fire for `[]`, so it would insert `[x]`
-    // here and reredden this expectation.
-    expect(parent?.children).toEqual([]);
-    expect(treeModel.find(t, "x")).toBeNull();
-    // ...and the chevron stays enabled so the user can expand to load it.
-    expect((parent as PH).hasChildren).toBe(true);
-  });
-
-  it("does NOT materialize a child under an UNLOADED parent (children undefined, hasChildren: true)", () => {
-    const tree: PH[] = [
-      { id: "p", name: "P", position: "a0", hasChildren: true }, // children: undefined
-    ];
-    const node: PH = { id: "x", name: "X", position: "a1" };
-    const t = treeModel.insertByPosition(tree, "p", node);
-    const parent = treeModel.find(t, "p");
-    expect(parent?.children).toBeUndefined();
-    expect(treeModel.find(t, "x")).toBeNull();
-    expect((parent as PH).hasChildren).toBe(true);
-  });
-
-  it("DOES insert under a genuinely-empty parent (children: [], hasChildren: false)", () => {
-    const tree: PH[] = [
-      { id: "p", name: "P", position: "a0", hasChildren: false, children: [] },
-    ];
-    const node: PH = { id: "x", name: "X", position: "a1" };
-    const t = treeModel.insertByPosition(tree, "p", node);
-    // No server children (`hasChildren: false`), so materializing the first child
-    // is correct — nothing is hidden.
-    expect(treeModel.find(t, "p")?.children?.map((n) => n.id)).toEqual(["x"]);
-  });
-
-  it("DOES insert under a genuinely-empty parent (children undefined, hasChildren: false)", () => {
+  // #159 #1: inserting/moving a node under a parent whose children are NOT
+  // loaded (`children === undefined`, e.g. a collapsed page) must NOT materialize
+  // a partial `[node]` list — that would defeat the lazy-load gate and hide the
+  // parent's other real children. The node is left to be lazy-loaded; only
+  // `hasChildren` is flagged so the chevron appears.
+  it("does NOT materialize a child under an UNLOADED parent (children undefined)", () => {
+    type PH = TreeNode<{
+      name: string;
+      position?: string;
+      hasChildren?: boolean;
+    }>;
     const tree: PH[] = [
       { id: "p", name: "P", position: "a0", hasChildren: false }, // children: undefined
     ];
     const node: PH = { id: "x", name: "X", position: "a1" };
     const t = treeModel.insertByPosition(tree, "p", node);
+    const parent = treeModel.find(t, "p");
+    // The node was NOT inserted (children stay unloaded -> lazy-load fetches the
+    // full set, including this node, on expand).
+    expect(parent?.children).toBeUndefined();
+    expect(treeModel.find(t, "x")).toBeNull();
+    // ...but the chevron is enabled so the user can expand to load it.
+    expect((parent as PH).hasChildren).toBe(true);
+  });
+
+  it("DOES insert under a LOADED-but-empty parent (children: [])", () => {
+    type PH = TreeNode<{
+      name: string;
+      position?: string;
+      hasChildren?: boolean;
+    }>;
+    const tree: PH[] = [
+      { id: "p", name: "P", position: "a0", hasChildren: false, children: [] },
+    ];
+    const node: PH = { id: "x", name: "X", position: "a1" };
+    const t = treeModel.insertByPosition(tree, "p", node);
+    // A loaded (empty) child list is complete, so the node IS inserted.
     expect(treeModel.find(t, "p")?.children?.map((n) => n.id)).toEqual(["x"]);
   });
 

apps/client/src/features/page/tree/model/tree-model.ts

@@ -43,26 +43,6 @@ export const treeModel = {
     };
   },
 
-  // A branch is "unloaded" when the server says it HAS children (`hasChildren`)
-  // but none are present locally. The canonical unloaded form in this codebase
-  // is `children: []` (produced by `pageToTreeNode` and by `pruneCollapsedChildren`
-  // resetting collapsed branches), NOT `children: undefined` — so a predicate that
-  // only checks `=== undefined` misses the real case and materializes a misleading
-  // partial list (#525). This is the SINGLE source of truth for "should a
-  // fetch/materialize be deferred?", shared by the lazy-load gate (`handleToggle`)
-  // and the realtime insert paths (`insertByPosition` / `placeByPosition`), so they
-  // can never drift apart again. (The local raw `insert` primitive and its DnD/
-  // create-page callers do NOT yet route through this predicate — see #525
-  // follow-up.) A parent WITHOUT `hasChildren` is genuinely empty
-  // (no server children) — inserting its first child is correct, not deferred.
-  isUnloadedBranch<T extends object>(
-    node: TreeNode<T> | null | undefined,
-  ): boolean {
-    if (!node) return false;
-    const hasChildren = (node as { hasChildren?: boolean }).hasChildren === true;
-    return hasChildren && (node.children == null || node.children.length === 0);
-  },
-
   isDescendant<T extends object>(
     tree: TreeNode<T>[],
     ancestorId: string,
@@ -147,15 +127,14 @@ export const treeModel = {
     }
     const parent = treeModel.find(tree, parentId);
     // The parent is in the tree but its children have NOT been lazy-loaded yet
-    // (`hasChildren` set + children absent/empty — see `isUnloadedBranch`; the
-    // canonical unloaded form is `children: []`, NOT just `undefined`). Inserting
+    // (`children === undefined`, distinct from a loaded-but-empty `[]`). Inserting
     // here would MATERIALIZE a misleading partial child list (`[node]`) that
     // defeats the lazy-load gate — which fetches only when children are
     // absent/empty — so the parent's OTHER real children would never load and the
     // moved/added node would be the only one shown (a silent data loss, #159 #1).
     // Instead, leave the children unloaded and just flag `hasChildren` so the
     // chevron appears; expanding fetches the FULL set (including this node).
-    if (parent && treeModel.isUnloadedBranch(parent)) {
+    if (parent && parent.children === undefined) {
       return treeModel.update(
         tree,
         parentId,

apps/client/src/features/websocket/tree-socket-reducers.test.ts

@@ -82,19 +82,17 @@ describe("applyMoveTreeNode", () => {
     ]);
   });
 
-  it("does NOT create a partial child list when the destination is loaded-but-collapsed (children unloaded) — keeps it lazy-loadable (#159 #525)", () => {
-    // `dstCollapsed` is in the tree but its children were never lazy-loaded. The
-    // CANONICAL unloaded form here is `hasChildren: true` + `children: []` (from
-    // `pageToTreeNode` / `pruneCollapsedChildren`), NOT `children: undefined`.
-    // The pre-#525 predicate (`children === undefined`) missed this form and
-    // inserted `src` as the ONLY child ([src]), defeating the lazy-load gate and
-    // HIDING the parent's other real children. Now the move leaves children
-    // unloaded (so expanding fetches the FULL set, including src).
+  it("does NOT create a partial child list when the destination is loaded-but-collapsed (children unloaded) — keeps it lazy-loadable (#159)", () => {
+    // `dstCollapsed` is in the tree but its children were never lazy-loaded
+    // (children === undefined). The OLD behavior inserted `src` as the ONLY
+    // child ([src]), which defeated the lazy-load gate and HID the parent's
+    // other real children. Now the move leaves children unloaded (so expanding
+    // fetches the FULL set, including src) and just flags hasChildren.
     const tree: SpaceTreeNode[] = [
       node("dstCollapsed", {
         position: "a0",
-        hasChildren: true,
-        children: [],
+        hasChildren: false,
+        children: undefined as unknown as SpaceTreeNode[],
       }),
       node("src", { position: "a9" }),
     ];
@@ -107,10 +105,9 @@ describe("applyMoveTreeNode", () => {
       pageData: {},
     });
     const dst = treeModel.find(next, "dstCollapsed");
-    // Children stay unloaded ([] not materialized to [src]) -> the lazy-load gate
-    // fetches the FULL set (incl. src) on expand. MUTATION: the pre-#525
-    // `=== undefined` predicate would insert [src] here and redden this.
-    expect(dst?.children).toEqual([]);
+    // Children stay unloaded -> the lazy-load gate fetches the FULL set (incl.
+    // src) on expand, rather than showing a misleading partial [src] list.
+    expect(dst?.children).toBeUndefined();
     expect(dst?.hasChildren).toBe(true);
     // src moved away from its old root slot (it lives under dstCollapsed
     // server-side and reappears when the parent is expanded/loaded).

apps/server/src/collaboration/collaboration.module.ts

@@ -16,7 +16,6 @@ import { TransclusionService } from '../core/page/transclusion/transclusion.serv
 import { TransclusionModule } from '../core/page/transclusion/transclusion.module';
 import { StorageModule } from '../integrations/storage/storage.module';
 import { EnvironmentModule } from '../integrations/environment/environment.module';
-import { ApiKeyModule } from '../core/api-key/api-key.module';
 
 @Module({
   providers: [
@@ -32,7 +31,6 @@ import { ApiKeyModule } from '../core/api-key/api-key.module';
   exports: [CollaborationGateway],
   imports: [
     TokenModule,
-    ApiKeyModule,
     WatcherModule,
     StorageModule.forRootAsync({
       imports: [EnvironmentModule],

apps/server/src/collaboration/extensions/authentication.extension.spec.ts

@@ -52,7 +52,6 @@ describe('AuthenticationExtension.onAuthenticate', () => {
   let pageRepo: { findById: jest.Mock };
   let spaceMemberRepo: { getUserSpaceRoles: jest.Mock };
   let pagePermissionRepo: { canUserEditPage: jest.Mock };
-  let apiKeyService: { validate: jest.Mock };
 
   // Build the hocuspocus onAuthenticate payload. connectionConfig.readOnly
   // starts false; the extension flips it to true on a read-only downgrade.
@@ -80,15 +79,12 @@ describe('AuthenticationExtension.onAuthenticate', () => {
       }),
     };
 
-    apiKeyService = { validate: jest.fn().mockResolvedValue({ user: {}, workspace: {} }) };
-
     ext = new AuthenticationExtension(
       tokenService as any,
       userRepo as any,
       pageRepo as any,
       spaceMemberRepo as any,
       pagePermissionRepo as any,
-      apiKeyService as any,
     );
     // Silence the extension's logger (it warns/debugs on denial branches).
     jest.spyOn(ext['logger'], 'warn').mockImplementation(() => undefined);
@@ -235,73 +231,4 @@ describe('AuthenticationExtension.onAuthenticate', () => {
     // No internal ai_chats row for an MCP/service-account collab edit → null.
     expect(ctx.aiChatId).toBeNull();
   });
-
-  // --- #501: api-key laundering guard (fail-closed discriminator) ----------
-  describe('api-key laundering guard', () => {
-    it('api_key principal → row-checks the key on connect (valid key proceeds)', async () => {
-      tokenService.verifyJwt.mockResolvedValue(
-        buildJwt({ principal: 'api_key', apiKeyId: 'key-1' }),
-      );
-      const data = buildData();
-      await ext.onAuthenticate(data as any);
-
-      expect(apiKeyService.validate).toHaveBeenCalledTimes(1);
-      expect(apiKeyService.validate).toHaveBeenCalledWith(
-        expect.objectContaining({ apiKeyId: 'key-1', type: JwtType.API_KEY }),
-      );
-    });
-
-    it('REVOKED api_key → Unauthorized on connect, BEFORE any page/user lookup', async () => {
-      tokenService.verifyJwt.mockResolvedValue(
-        buildJwt({ principal: 'api_key', apiKeyId: 'key-1' }),
-      );
-      // The shared validator denies a revoked key.
-      apiKeyService.validate.mockRejectedValue(new UnauthorizedException());
-
-      await expect(ext.onAuthenticate(buildData() as any)).rejects.toThrow(
-        UnauthorizedException,
-      );
-      // No new collab connection: the key check gates before page access.
-      expect(pageRepo.findById).not.toHaveBeenCalled();
-    });
-
-    it('api_key principal missing apiKeyId → Unauthorized (malformed)', async () => {
-      tokenService.verifyJwt.mockResolvedValue(buildJwt({ principal: 'api_key' }));
-      await expect(ext.onAuthenticate(buildData() as any)).rejects.toThrow(
-        UnauthorizedException,
-      );
-      expect(apiKeyService.validate).not.toHaveBeenCalled();
-    });
-
-    it('session principal → NO api-key check (session-backed, incl. internal agent)', async () => {
-      tokenService.verifyJwt.mockResolvedValue(buildJwt({ principal: 'session' }));
-      await ext.onAuthenticate(buildData() as any);
-      expect(apiKeyService.validate).not.toHaveBeenCalled();
-    });
-
-    it('claimless token WITHIN the grace window → trusted (legacy pre-rollout)', async () => {
-      // Default rolloutAt = now, so we are inside the grace window.
-      tokenService.verifyJwt.mockResolvedValue(buildJwt()); // no principal
-      await expect(ext.onAuthenticate(buildData() as any)).resolves.toBeDefined();
-      expect(apiKeyService.validate).not.toHaveBeenCalled();
-    });
-
-    it('claimless token AFTER the grace window → Unauthorized (fail-closed)', async () => {
-      // Move the rollout reference far into the past so the grace has elapsed.
-      (ext as any).rolloutAt = Date.now() - 25 * 60 * 60 * 1000;
-      tokenService.verifyJwt.mockResolvedValue(buildJwt()); // no principal
-      await expect(ext.onAuthenticate(buildData() as any)).rejects.toThrow(
-        UnauthorizedException,
-      );
-    });
-
-    it('infra error from the api-key row-check propagates (not masked)', async () => {
-      tokenService.verifyJwt.mockResolvedValue(
-        buildJwt({ principal: 'api_key', apiKeyId: 'key-1' }),
-      );
-      const boom = new Error('db down');
-      apiKeyService.validate.mockRejectedValue(boom);
-      await expect(ext.onAuthenticate(buildData() as any)).rejects.toBe(boom);
-    });
-  });
 });

apps/server/src/collaboration/extensions/authentication.extension.ts

@@ -14,37 +14,20 @@ import { findHighestUserSpaceRole } from '@docmost/db/repos/space/utils';
 import { SpaceRole } from '../../common/helpers/types/permission';
 import { isUserDisabled } from '../../common/helpers';
 import { getPageId } from '../collaboration.util';
-import {
-  JwtApiKeyPayload,
-  JwtCollabPayload,
-  JwtType,
-} from '../../core/auth/dto/jwt-payload';
+import { JwtCollabPayload, JwtType } from '../../core/auth/dto/jwt-payload';
 import { resolveProvenance } from '../../common/decorators/auth-provenance.decorator';
 import { observeCollabAuth } from '../../integrations/metrics/metrics.registry';
-import { ApiKeyService } from '../../core/api-key/api-key.service';
-
-// Max lifetime of a collab token (generateCollabToken uses expiresIn '24h'). Used
-// as the rollout grace window below: once this long has elapsed since this
-// process started serving the #501 code, every STILL-VALID collab token was
-// necessarily minted post-rollout and MUST carry the `principal` discriminator,
-// so a claimless one is a bug and is rejected (fail-closed) rather than trusted.
-const COLLAB_TOKEN_GRACE_MS = 24 * 60 * 60 * 1000;
 
 @Injectable()
 export class AuthenticationExtension implements Extension {
   private readonly logger = new Logger(AuthenticationExtension.name);
 
-  // Reference instant for the claimless-rejection grace window. Overridable so a
-  // unit test can drive the pre-/post-grace boundary without wall-clock waits.
-  protected rolloutAt = Date.now();
-
   constructor(
     private tokenService: TokenService,
     private userRepo: UserRepo,
     private pageRepo: PageRepo,
     private readonly spaceMemberRepo: SpaceMemberRepo,
     private readonly pagePermissionRepo: PagePermissionRepo,
-    private readonly apiKeyService: ApiKeyService,
   ) {}
 
   async onAuthenticate(data: onAuthenticatePayload) {
@@ -71,36 +54,6 @@ export class AuthenticationExtension implements Extension {
       throw new UnauthorizedException('Invalid collab token');
     }
 
-    // #501 — fail-closed api-key laundering guard. A collab token minted by an
-    // api-key principal carries principal='api_key' + apiKeyId; re-check the key
-    // on connect so a REVOKED key gets NO new collab connections (a collab token
-    // outlives its 24h, but a revoked key can no longer open fresh ones). An
-    // api-key token missing its apiKeyId is malformed → reject. A claimless token
-    // (no recognized principal) is trusted only DURING the rollout grace window
-    // (a legacy pre-rollout session token, which api keys could never mint);
-    // once the grace has elapsed every valid token must carry the discriminator,
-    // so a claimless one is a bug and is rejected (not silently trusted for 24h).
-    const principal = jwtPayload.principal;
-    if (principal === 'api_key') {
-      if (!jwtPayload.apiKeyId) {
-        throw new UnauthorizedException();
-      }
-      // Row-check via the SHARED validator: throws Unauthorized on a revoked/
-      // expired/disabled key; an infra error propagates (not masked). No new
-      // connection for a dead key.
-      await this.apiKeyService.validate({
-        sub: jwtPayload.sub,
-        workspaceId: jwtPayload.workspaceId,
-        apiKeyId: jwtPayload.apiKeyId,
-        type: JwtType.API_KEY,
-      } as JwtApiKeyPayload);
-    } else if (principal !== 'session') {
-      // Unrecognized/absent discriminator: reject once past the grace window.
-      if (Date.now() - this.rolloutAt >= COLLAB_TOKEN_GRACE_MS) {
-        throw new UnauthorizedException();
-      }
-    }
-
     const userId = jwtPayload.sub;
     const workspaceId = jwtPayload.workspaceId;
 

apps/server/src/core/ai-chat/ai-chat.controller.ts

@@ -35,7 +35,6 @@ import {
 import { PaginationOptions } from '@docmost/db/pagination/pagination-options';
 import { AiChatRepo } from '@docmost/db/repos/ai-chat/ai-chat.repo';
 import { AiChatMessageRepo } from '@docmost/db/repos/ai-chat/ai-chat-message.repo';
-import { AiChatRunStepRepo } from '@docmost/db/repos/ai-chat/ai-chat-run-step.repo';
 import { PageRepo } from '@docmost/db/repos/page/page.repo';
 import { UserThrottlerGuard } from '../../integrations/throttle/user-throttler.guard';
 import { AI_CHAT_THROTTLER } from '../../integrations/throttle/throttler-names';
@@ -44,8 +43,6 @@ import {
   AiChatRunHooks,
   AiChatService,
   AiChatStreamBody,
-  rowHasInlineParts,
-  hydrateAssistantParts,
 } from './ai-chat.service';
 import { AiChatRunService } from './ai-chat-run.service';
 import { AiTranscriptionService } from './ai-transcription.service';
@@ -132,39 +129,8 @@ export class AiChatController {
     // production. Only touched on the resumable-stream (flag-on) path.
     private readonly streamRegistry?: AiChatStreamRegistryService,
     private readonly environment?: EnvironmentService,
-    // #492: reconstruct a #492 mid-run record's parts from the steps table before
-    // returning rows to the client / export. OPTIONAL so positional controller
-    // specs compile unchanged; when absent, hydration is skipped (old-era rows
-    // already carry inline parts, so nothing to reconstruct).
-    private readonly aiChatRunStepRepo?: AiChatRunStepRepo,
   ) {}
 
-  /**
-   * Reconstruct parts for any assistant rows that don't carry them INLINE — a
-   * #492 mid-run record whose per-step parts live in `ai_chat_run_steps` (the
-   * append-persist backend). Every FINISHED row (old-era + #492) and every old-era
-   * streaming snapshot already has inline `metadata.parts`, so the common path
-   * fetches NOTHING and returns the rows untouched; only an actively-streaming
-   * new-style row triggers the batch step fetch. Consumers (seed/poll/export) read
-   * `metadata.parts` off the returned rows exactly as before — the era switch is
-   * invisible to them (reconstructRunParts contract).
-   */
-  private async withReconstructedParts(
-    rows: AiChatMessage[],
-    workspaceId: string,
-  ): Promise<AiChatMessage[]> {
-    if (!this.aiChatRunStepRepo) return rows;
-    const needy = rows.filter(
-      (r) => r.role === 'assistant' && !rowHasInlineParts(r),
-    );
-    if (needy.length === 0) return rows;
-    const stepsByMessage = await this.aiChatRunStepRepo.findByMessageIds(
-      needy.map((r) => r.id),
-      workspaceId,
-    );
-    return hydrateAssistantParts(rows, stepsByMessage);
-  }
-
   /** List the requesting user's chats in this workspace (paginated). */
   @HttpCode(HttpStatus.OK)
   @Post('chats')
@@ -218,17 +184,11 @@ export class AiChatController {
     @AuthWorkspace() workspace: Workspace,
   ) {
     await this.assertOwnedChat(dto.chatId, user, workspace);
-    const page = await this.aiChatMessageRepo.findByChat(
+    return this.aiChatMessageRepo.findByChat(
       dto.chatId,
       workspace.id,
       pagination,
     );
-    // #492: reconstruct parts for any active new-style row so the client seed sees
-    // `metadata.parts` unchanged (a no-op for the finished rows that fill a page).
-    return {
-      ...page,
-      items: await this.withReconstructedParts(page.items, workspace.id),
-    };
   }
 
   /**
@@ -265,10 +225,7 @@ export class AiChatController {
       workspace.id,
     );
     return {
-      // #492: the delta of an actively-streaming new-style row carries its parts
-      // reconstructed from the steps table, so the degraded poll shows persisted
-      // progress exactly as the pre-#492 full-row snapshot did.
-      rows: await this.withReconstructedParts(rows, workspace.id),
+      rows,
       cursor,
       run: run ? { id: run.id, status: run.status } : null,
     };
@@ -290,10 +247,8 @@ export class AiChatController {
     @AuthWorkspace() workspace: Workspace,
   ): Promise<{ markdown: string }> {
     const chat = await this.assertOwnedChat(dto.chatId, user, workspace);
-    const rows = await this.withReconstructedParts(
-      await this.aiChatMessageRepo.findAllByChat(dto.chatId, workspace.id),
-      // #492: an interrupted-but-still-active turn exports its persisted steps
-      // (reconstructed from the steps table) just like the pre-#492 full row did.
+    const rows = await this.aiChatMessageRepo.findAllByChat(
+      dto.chatId,
       workspace.id,
     );
     const markdown = buildChatMarkdown({
@@ -333,13 +288,7 @@ export class AiChatController {
           workspace.id,
         )
       : undefined;
-    // #492: reconnect to an IN-FLIGHT run reconstructs the projection row's parts
-    // from the steps table (the row itself carries only the step marker mid-run);
-    // a finished run's row already has inline parts, so this is a no-op.
-    const [hydrated] = message
-      ? await this.withReconstructedParts([message], workspace.id)
-      : [undefined];
-    return { run, message: hydrated ?? null };
+    return { run, message: message ?? null };
   }
 
   /**

apps/server/src/core/ai-chat/ai-chat.service.ts

@@ -22,7 +22,6 @@ import { AiSettingsService } from '../../integrations/ai/ai-settings.service';
 import { describeProviderError } from '../../integrations/ai/ai-error.util';
 import { AiChatRepo } from '@docmost/db/repos/ai-chat/ai-chat.repo';
 import { AiChatMessageRepo } from '@docmost/db/repos/ai-chat/ai-chat-message.repo';
-import { AiChatRunStepRepo } from '@docmost/db/repos/ai-chat/ai-chat-run-step.repo';
 import { AiChatPageSnapshotRepo } from '@docmost/db/repos/ai-chat/ai-chat-page-snapshot.repo';
 import { AiAgentRoleRepo } from '@docmost/db/repos/ai-agent-roles/ai-agent-roles.repo';
 import { PageRepo } from '@docmost/db/repos/page/page.repo';
@@ -190,11 +189,10 @@ export function stepBudgetWarning(stepNumber: number): string {
 //
 // `system` is the in-scope system prompt; we CONCATENATE so the original
 // persona/context is preserved — a bare `system` override would REPLACE the
-// whole system prompt for the step. `activatedTools` is a closure Set grown by
-// loadTools and owned by the streaming loop; the caller seeds it from and
-// persists it to the chat's metadata across turns (#490), but this function only
-// READS the Set it is handed, so it stays a pure function of its arguments (not
-// module-global).
+// whole system prompt for the step. `activatedTools` is PER-TURN mutable state
+// owned by the streaming loop (a closure Set grown by loadTools); it is passed
+// in (not module-global, not persisted) so this stays a pure function of its
+// arguments.
 //
 // NOTE: at AI SDK v7 the per-step `system` field is renamed to `instructions`.
 // On v6 (`^6.0.134`) `system` is the correct field — adjust when bumping.
@@ -519,12 +517,6 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
     // constructions compile unchanged; Nest always injects the real singleton, so
     // reconcile sees the SAME in-memory active/zombie maps the runner mutates.
     private readonly aiChatRunService?: AiChatRunService,
-    // #492 append-persist: per-step INSERT into the lightweight steps table (the
-    // O(Σ steps) replacement for the O(n²) full-row `metadata.parts` rewrite).
-    // OPTIONAL so existing positional constructions (int-specs) compile unchanged;
-    // Nest injects the real singleton. When ABSENT the per-step path falls back to
-    // the pre-#492 full-row flush (no regression, only no WAL win).
-    private readonly aiChatRunStepRepo?: AiChatRunStepRepo,
   ) {}
 
   // #487: periodic reconcile timer (single-process phase 1). Started in
@@ -1121,34 +1113,8 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
         chatId,
         workspace.id,
       );
-      // #492: HYDRATE needy assistant rows from the steps table BEFORE the replay
-      // map. A #492 mid-run assistant row carries only a step marker
-      // (metadata.parts:[]); its real per-step parts live in `ai_chat_run_steps`.
-      // The graceful terminal callbacks (onFinish/onError/onAbort -> flushAssistant)
-      // assemble the full inline parts, so a normally-ended turn already has them.
-      // But a HARD crash mid-run (SIGKILL/OOM) fires NO terminal callback, so the
-      // row stays parts:[]; without this, rowToUiMessage falls back to an empty
-      // text part and the partial tool-calls/results/text — durable in the steps
-      // table — would DROP OUT of the model's replay context (regressing #183
-      // step-granular durability for the model consumer). Mirrors the controller's
-      // withReconstructedParts EXACTLY (same needy predicate + hydration helper).
-      // Guarded on the optional repo: absent (positional test builds) degrades to
-      // the current behavior rather than crashing.
-      let replayHistory = oldHistory;
-      if (this.aiChatRunStepRepo) {
-        const needy = oldHistory.filter(
-          (r) => r.role === 'assistant' && !rowHasInlineParts(r),
-        );
-        if (needy.length > 0) {
-          const stepsByMessage = await this.aiChatRunStepRepo.findByMessageIds(
-            needy.map((r) => r.id),
-            workspace.id,
-          );
-          replayHistory = hydrateAssistantParts(oldHistory, stepsByMessage);
-        }
-      }
       const uiMessages: Array<Omit<UIMessage, 'id'> & { id: string }> = [
-        ...replayHistory.map(rowToUiMessage),
+        ...oldHistory.map(rowToUiMessage),
         {
           id: 'pending-user',
           role: 'user',
@@ -1187,9 +1153,7 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
       // hint — confirm it against the persisted history (the preceding assistant
       // turn must really be aborted/streaming) so a spoofed flag cannot inject the
       // interrupt note onto an ordinary turn. The partial output the model needs is
-      // already in `messages`: a #492 mid-run row's per-step parts live only in the
-      // `ai_chat_run_steps` table and were hydrated into the replay history above,
-      // so the aborted assistant turn replays WITH its partial parts intact.
+      // already in `messages` (the aborted assistant row replays via findRecent).
       // Append the new user turn (shape-only) so index -2 is the prior assistant.
       const interrupted = isInterruptResume(
         [...oldHistory, { role: 'user', status: null, metadata: null }],
@@ -1446,11 +1410,10 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
       const baseTools = { ...external.tools, ...docmostTools };
 
       // Deferred tool loading state (#332), scoped to THIS streaming loop:
-      //  - `activatedTools` is a fresh closure Set per streamText call (not
-      //    module-global), SEEDED from the chat's persisted metadata.activatedTools
-      //    (#490, just below) so activation carries across turns. loadTools.execute
-      //    adds to it; prepareAgentStep reads it to widen `activeTools` on the NEXT
-      //    step; turn end persists it back.
+      //  - `activatedTools` is per-TURN mutable state — a fresh closure Set created
+      //    per streamText call, NOT module-global and NOT persisted, so a new turn
+      //    starts cold. loadTools.execute adds to it; prepareAgentStep reads it to
+      //    widen `activeTools` on the NEXT step.
       //  - `validDeferredNames` = every tool that is NOT core (the in-app deferred
       //    tools + ALL external MCP tools), computed from the ACTUAL toolset so an
       //    external tool is loadable by its namespaced name. loadTools rejects any
@@ -1594,57 +1557,17 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
         // connection when finalize runs, so the SQL `WHERE status='streaming'`
         // (not this flag) is what prevents it clobbering the terminal row.
         if (finalized) return null;
-        // The count derives from capturedSteps.length at THIS instant, so the
-        // returned value is EXACTLY the persisted `stepsPersisted` the ring rotates
-        // on (whether we take the append-persist path or the legacy fallback).
-        const stepsPersisted = capturedSteps.length;
+        // Build the flush ONCE so the returned count is EXACTLY the persisted
+        // `stepsPersisted` (both derive from capturedSteps.length at this instant).
+        const flushed = flushAssistant(capturedSteps, '', 'streaming', {
+          pageChanged,
+          partsCache,
+        });
+        const stepsPersisted = flushed.metadata.stepsPersisted as number;
         try {
-          if (this.aiChatRunStepRepo) {
-            // #492 APPEND-PERSIST: write only THIS finished step's parts to the
-            // steps table (O(step) WAL), then bump the row's CHEAP step marker —
-            // NO growing `metadata.parts` blob (that O(n²) full-row rewrite is
-            // exactly what this removes). The full `metadata.parts` is assembled
-            // once at finalize; a mid-run resume seed is reconstructed from the
-            // step rows (reconstructRunParts). The INSERT is idempotent
-            // (ON CONFLICT DO NOTHING), so a re-fired step never doubles the parts.
-            const index = stepsPersisted - 1;
-            if (index >= 0) {
-              const stepParts = assistantParts(
-                [capturedSteps[index]],
-                '',
-                partsCache,
-              );
-              await this.aiChatRunStepRepo.insertStep(
-                assistantId,
-                workspace.id,
-                index,
-                stepParts,
-              );
-            }
-            // Marker UPDATE: advance stepsPersisted + keep the toolTrace era marker
-            // (bumps updatedAt so the delta poll observes the step, and carries the
-            // frontier a resuming client attaches from). Scoped onlyIfStreaming so a
-            // late marker never clobbers the terminal finalize.
-            await this.aiChatMessageRepo.update(
-              assistantId,
-              workspace.id,
-              { metadata: stepMarkerMetadata(stepsPersisted) },
-              { onlyIfStreaming: true },
-            );
-          } else {
-            // Legacy fallback (no steps table wired — positional test builds): the
-            // pre-#492 full-row flush, so parts still land inline on the row.
-            const flushed = flushAssistant(capturedSteps, '', 'streaming', {
-              pageChanged,
-              partsCache,
-            });
-            await this.aiChatMessageRepo.update(
-              assistantId,
-              workspace.id,
-              flushed,
-              { onlyIfStreaming: true },
-            );
-          }
+          await this.aiChatMessageRepo.update(assistantId, workspace.id, flushed, {
+            onlyIfStreaming: true,
+          });
           return stepsPersisted;
         } catch (err) {
           this.logger.warn(
@@ -2824,122 +2747,6 @@ export function rowToUiMessage(row: AiChatMessage): Omit<UIMessage, 'id'> & {
   return { id: row.id, role, parts: parts as UIMessage['parts'] };
 }
 
-/**
- * Cheap step-marker metadata for the #492 per-step UPDATE. Advances
- * `stepsPersisted` (the resume attach frontier) and keeps the `toolTraceVersion`
- * era marker, WITHOUT the growing `parts` blob (those live in the steps table
- * now; the full `metadata.parts` is assembled once at finalize by flushAssistant).
- * `parts: []` is kept for shape stability — it reads as an empty inline-parts row,
- * which is exactly the discriminator that routes reconstruction to the steps table.
- */
-export function stepMarkerMetadata(
-  stepsPersisted: number,
-): Record<string, unknown> {
-  return { parts: [], toolTraceVersion: 2, stepsPersisted };
-}
-
-/**
- * Whether an assistant row already carries its full UI parts INLINE on the row
- * (`metadata.parts`). TRUE for every FINISHED row — old-era rows AND #492 rows,
- * whose full parts are assembled once at finalize — and for old-era streaming
- * snapshots (the pre-#492 per-step full-row flush). FALSE for a #492 MID-RUN
- * record, whose per-step parts live in the `ai_chat_run_steps` table. This is the
- * era discriminator the reconstruct seam branches on — no schema flag needed.
- */
-export function rowHasInlineParts(row: { metadata?: unknown }): boolean {
-  const meta = (row.metadata ?? {}) as { parts?: unknown };
-  return Array.isArray(meta.parts) && meta.parts.length > 0;
-}
-
-/**
- * Concatenate persisted per-step parts (in `stepIndex` order) into the turn's UI
- * parts (#492). Reproduces EXACTLY what flushAssistant → assistantParts would have
- * written to `metadata.parts` for those finished steps, since each step row stored
- * `assistantParts([step])` at persist time.
- */
-export function assembleStepParts(
-  stepRows: ReadonlyArray<{ stepIndex: number; parts: unknown }>,
-): UIMessage['parts'] {
-  const parts: Array<Record<string, unknown>> = [];
-  for (const step of [...stepRows].sort((a, b) => a.stepIndex - b.stepIndex)) {
-    if (Array.isArray(step.parts)) {
-      parts.push(...(step.parts as Array<Record<string, unknown>>));
-    }
-  }
-  return parts as UIMessage['parts'];
-}
-
-/**
- * reconstructRunParts (#492) — the single backend-switch seam. Given an assistant
- * ROW and its persisted step rows, return the turn's UI `parts` + the persisted
- * step count, reading from the ROW when it already carries inline parts (old-era
- * records AND every finished record) and from the STEPS TABLE otherwise (a #492
- * mid-run record). The higher-level consumers (attach seed, delta poll, export)
- * route their row→parts through this / {@link hydrateAssistantParts}, so old and
- * new records reconstruct identically WITHOUT the consumers branching on the era.
- */
-export function reconstructRunParts(
-  row: { metadata?: unknown; content?: string | null },
-  stepRows: ReadonlyArray<{ stepIndex: number; parts: unknown }>,
-): { parts: UIMessage['parts']; stepsPersisted: number } {
-  if (rowHasInlineParts(row)) {
-    const meta = row.metadata as {
-      parts: UIMessage['parts'];
-      stepsPersisted?: number;
-    };
-    return {
-      parts: meta.parts,
-      stepsPersisted:
-        typeof meta.stepsPersisted === 'number'
-          ? meta.stepsPersisted
-          : stepRows.length,
-    };
-  }
-  if (stepRows.length > 0) {
-    return {
-      parts: assembleStepParts(stepRows),
-      stepsPersisted: stepRows.length,
-    };
-  }
-  // No inline parts and no step rows: an old-era seed / empty streaming row. Fall
-  // back to a single text part from `content` (mirrors rowToUiMessage).
-  return {
-    parts: textPart(row.content ?? '') as UIMessage['parts'],
-    stepsPersisted: 0,
-  };
-}
-
-/**
- * Fill each assistant row's `metadata.parts` from its step rows when the row does
- * not already carry them inline (a #492 mid-run record), so a consumer that reads
- * `metadata.parts` off the RAW row (the client seed/poll, the Markdown export)
- * sees the reconstructed parts with NO change to itself. Rows that already have
- * inline parts (old-era + finished) and non-assistant rows pass through untouched.
- * Pure: returns new row objects, never mutates the inputs.
- */
-export function hydrateAssistantParts<
-  T extends { id: string; role?: string; metadata?: unknown },
->(
-  rows: ReadonlyArray<T>,
-  stepsByMessage: Map<
-    string,
-    ReadonlyArray<{ stepIndex: number; parts: unknown }>
-  >,
-): T[] {
-  return rows.map((row) => {
-    if (row.role !== 'assistant' || rowHasInlineParts(row)) return row;
-    const steps = stepsByMessage.get(row.id);
-    if (!steps || steps.length === 0) return row;
-    return {
-      ...row,
-      metadata: {
-        ...((row.metadata ?? {}) as Record<string, unknown>),
-        parts: assembleStepParts(steps),
-      },
-    };
-  });
-}
-
 /**
  * The persisted-row patch shape produced by {@link flushAssistant}. It is the
  * SAME shape the assistant repo insert/update consume (content + toolCalls +

apps/server/src/core/api-key/api-key.controller.spec.ts

@@ -1,193 +0,0 @@
-import { ForbiddenException, NotFoundException } from '@nestjs/common';
-import { ApiKeyController } from './api-key.controller';
-import {
-  WorkspaceCaslAction,
-  WorkspaceCaslSubject,
-} from '../casl/interfaces/workspace-ability.type';
-
-/**
- * Authorization contract for the /api-keys management surface.
- *
- *  - A token cannot manage tokens (an api_key PRINCIPAL is 403 on every method)
- *    — GitHub-PAT semantics closing post-revocation laundering.
- *  - admin (CASL Manage on API) sees/revokes all workspace keys; a member only
- *    their own.
- *  - kill-switch OFF -> the surface 404s (looks like the feature does not exist).
- */
-
-function makeController(over: any = {}) {
-  const apiKeyService = {
-    create: jest.fn().mockResolvedValue({
-      token: 'tok',
-      key: { id: 'k1', name: 'n', expiresAt: null, createdAt: new Date() },
-    }),
-    revoke: jest.fn().mockResolvedValue(undefined),
-    ...(over.apiKeyService ?? {}),
-  };
-  const apiKeyRepo = {
-    findAllInWorkspace: jest.fn().mockResolvedValue([]),
-    findByCreator: jest.fn().mockResolvedValue([]),
-    findById: jest.fn(),
-    ...(over.apiKeyRepo ?? {}),
-  };
-  const workspaceAbility = {
-    createForUser: jest.fn().mockReturnValue({
-      can: (_a: any, _s: any) => over.canManage ?? false,
-    }),
-    ...(over.workspaceAbility ?? {}),
-  };
-  const environmentService = {
-    isApiKeysEnabled: jest.fn().mockReturnValue(over.enabled ?? true),
-    ...(over.environmentService ?? {}),
-  };
-  const auditService = { log: jest.fn() };
-  const controller = new ApiKeyController(
-    apiKeyService as any,
-    apiKeyRepo as any,
-    workspaceAbility as any,
-    environmentService as any,
-    auditService as any,
-  );
-  return {
-    controller,
-    apiKeyService,
-    apiKeyRepo,
-    workspaceAbility,
-    environmentService,
-    auditService,
-  };
-}
-
-const user = { id: 'u-1', workspaceId: 'ws-1' } as any;
-const workspace = { id: 'ws-1' } as any;
-const reqAccess = () => ({ raw: {}, ip: '1.2.3.4', socket: {} }) as any;
-const reqApiKey = () =>
-  ({ raw: { authType: 'api_key', apiKeyId: 'k-self' }, ip: '1.2.3.4', socket: {} }) as any;
-
-describe('ApiKeyController — a token cannot manage tokens', () => {
-  it('403 on create for an api_key principal', async () => {
-    const { controller, apiKeyService } = makeController();
-    await expect(
-      controller.create({ name: 'x' } as any, user, reqApiKey()),
-    ).rejects.toBeInstanceOf(ForbiddenException);
-    expect(apiKeyService.create).not.toHaveBeenCalled();
-  });
-
-  it('403 on list for an api_key principal', async () => {
-    const { controller } = makeController();
-    await expect(
-      controller.list(user, workspace, reqApiKey()),
-    ).rejects.toBeInstanceOf(ForbiddenException);
-  });
-
-  it('403 on revoke for an api_key principal', async () => {
-    const { controller } = makeController();
-    await expect(
-      controller.revoke({ id: 'k1' } as any, user, workspace, reqApiKey()),
-    ).rejects.toBeInstanceOf(ForbiddenException);
-  });
-});
-
-describe('ApiKeyController — kill-switch OFF → 404', () => {
-  it('create/list/revoke all 404 when disabled', async () => {
-    const { controller } = makeController({ enabled: false });
-    await expect(
-      controller.create({ name: 'x' } as any, user, reqAccess()),
-    ).rejects.toBeInstanceOf(NotFoundException);
-    await expect(
-      controller.list(user, workspace, reqAccess()),
-    ).rejects.toBeInstanceOf(NotFoundException);
-    await expect(
-      controller.revoke({ id: 'k1' } as any, user, workspace, reqAccess()),
-    ).rejects.toBeInstanceOf(NotFoundException);
-  });
-});
-
-describe('ApiKeyController — list scoping', () => {
-  it('admin (Manage on API) lists ALL workspace keys', async () => {
-    const { controller, apiKeyRepo } = makeController({ canManage: true });
-    await controller.list(user, workspace, reqAccess());
-    expect(apiKeyRepo.findAllInWorkspace).toHaveBeenCalledWith('ws-1');
-    expect(apiKeyRepo.findByCreator).not.toHaveBeenCalled();
-  });
-
-  it('member lists ONLY their own keys', async () => {
-    const { controller, apiKeyRepo } = makeController({ canManage: false });
-    await controller.list(user, workspace, reqAccess());
-    expect(apiKeyRepo.findByCreator).toHaveBeenCalledWith('u-1', 'ws-1');
-    expect(apiKeyRepo.findAllInWorkspace).not.toHaveBeenCalled();
-  });
-});
-
-describe('ApiKeyController — revoke scoping', () => {
-  it("member cannot revoke another user's key (403)", async () => {
-    const { controller, apiKeyRepo, apiKeyService } = makeController({
-      canManage: false,
-    });
-    apiKeyRepo.findById.mockResolvedValue({
-      id: 'k1',
-      creatorId: 'someone-else',
-      workspaceId: 'ws-1',
-    });
-    await expect(
-      controller.revoke({ id: 'k1' } as any, user, workspace, reqAccess()),
-    ).rejects.toBeInstanceOf(ForbiddenException);
-    expect(apiKeyService.revoke).not.toHaveBeenCalled();
-  });
-
-  it('member CAN revoke their own key', async () => {
-    const { controller, apiKeyRepo, apiKeyService } = makeController({
-      canManage: false,
-    });
-    apiKeyRepo.findById.mockResolvedValue({
-      id: 'k1',
-      creatorId: 'u-1',
-      workspaceId: 'ws-1',
-    });
-    await controller.revoke({ id: 'k1' } as any, user, workspace, reqAccess());
-    expect(apiKeyService.revoke).toHaveBeenCalledWith('k1', 'ws-1');
-  });
-
-  it("admin CAN revoke another user's key", async () => {
-    const { controller, apiKeyRepo, apiKeyService } = makeController({
-      canManage: true,
-    });
-    apiKeyRepo.findById.mockResolvedValue({
-      id: 'k1',
-      creatorId: 'someone-else',
-      workspaceId: 'ws-1',
-    });
-    await controller.revoke({ id: 'k1' } as any, user, workspace, reqAccess());
-    expect(apiKeyService.revoke).toHaveBeenCalledWith('k1', 'ws-1');
-  });
-
-  it('404 when the key does not exist', async () => {
-    const { controller, apiKeyRepo } = makeController({ canManage: true });
-    apiKeyRepo.findById.mockResolvedValue(undefined);
-    await expect(
-      controller.revoke({ id: 'k1' } as any, user, workspace, reqAccess()),
-    ).rejects.toBeInstanceOf(NotFoundException);
-  });
-});
-
-describe('ApiKeyController — create emits audit + returns token once', () => {
-  it('logs API_KEY_CREATED and returns the token + computed expiry', async () => {
-    const { controller, auditService } = makeController();
-    const res = await controller.create(
-      { name: 'ci' } as any,
-      user,
-      reqAccess(),
-    );
-    expect(res.token).toBe('tok');
-    expect(res.apiKey).toMatchObject({ id: 'k1', name: 'n' });
-    expect(auditService.log).toHaveBeenCalledWith(
-      expect.objectContaining({ event: 'api_key.created' }),
-    );
-  });
-});
-
-// Sanity: the CASL subject used is the workspace API subject.
-it('uses WorkspaceCaslSubject.API with the Manage action', () => {
-  expect(WorkspaceCaslSubject.API).toBe('api_key');
-  expect(WorkspaceCaslAction.Manage).toBeDefined();
-});

apps/server/src/core/api-key/api-key.controller.ts

@@ -1,201 +0,0 @@
-import {
-  Body,
-  Controller,
-  ForbiddenException,
-  HttpCode,
-  HttpStatus,
-  Inject,
-  Logger,
-  NotFoundException,
-  Post,
-  Req,
-  UseGuards,
-} from '@nestjs/common';
-import { Throttle } from '@nestjs/throttler';
-import { FastifyRequest } from 'fastify';
-import { ApiKeyService } from './api-key.service';
-import { ApiKeyRepo } from '@docmost/db/repos/api-key/api-key.repo';
-import { CreateApiKeyDto } from './dto/create-api-key.dto';
-import { RevokeApiKeyDto } from './dto/revoke-api-key.dto';
-import { AuthUser } from '../../common/decorators/auth-user.decorator';
-import { AuthWorkspace } from '../../common/decorators/auth-workspace.decorator';
-import { JwtAuthGuard } from '../../common/guards/jwt-auth.guard';
-import { User, Workspace } from '@docmost/db/types/entity.types';
-import WorkspaceAbilityFactory from '../casl/abilities/workspace-ability.factory';
-import {
-  WorkspaceCaslAction,
-  WorkspaceCaslSubject,
-} from '../casl/interfaces/workspace-ability.type';
-import { AuditEvent, AuditResource } from '../../common/events/audit-events';
-import {
-  AUDIT_SERVICE,
-  IAuditService,
-} from '../../integrations/audit/audit.service';
-import { EnvironmentService } from '../../integrations/environment/environment.service';
-import { UserThrottlerGuard } from '../../integrations/throttle/user-throttler.guard';
-import { AUTH_THROTTLER } from '../../integrations/throttle/throttler-names';
-
-@UseGuards(JwtAuthGuard)
-@Controller('api-keys')
-export class ApiKeyController {
-  private readonly logger = new Logger(ApiKeyController.name);
-
-  constructor(
-    private readonly apiKeyService: ApiKeyService,
-    private readonly apiKeyRepo: ApiKeyRepo,
-    private readonly workspaceAbility: WorkspaceAbilityFactory,
-    private readonly environmentService: EnvironmentService,
-    @Inject(AUDIT_SERVICE) private readonly auditService: IAuditService,
-  ) {}
-
-  // Kill-switch OFF: the issuance surface disappears entirely (404), not a 403 —
-  // it must look like the feature does not exist.
-  private assertEnabled(): void {
-    if (!this.environmentService.isApiKeysEnabled()) {
-      throw new NotFoundException();
-    }
-  }
-
-  // A token cannot manage tokens (GitHub-PAT semantics): an api-key principal is
-  // refused on the management surface, so a leaked key cannot mint a replacement
-  // or revoke the keys that would lock it out (closes post-revocation laundering).
-  private rejectApiKeyPrincipal(req: FastifyRequest): void {
-    if ((req.raw as { authType?: string }).authType === 'api_key') {
-      throw new ForbiddenException('API keys cannot manage API keys');
-    }
-  }
-
-  private clientIp(req: FastifyRequest): string {
-    return req.ip ?? req.socket?.remoteAddress ?? 'unknown';
-  }
-
-  @HttpCode(HttpStatus.OK)
-  @UseGuards(JwtAuthGuard, UserThrottlerGuard)
-  @Throttle({ [AUTH_THROTTLER]: { limit: 10, ttl: 60_000 } })
-  @Post('create')
-  async create(
-    @Body() dto: CreateApiKeyDto,
-    @AuthUser() user: User,
-    @Req() req: FastifyRequest,
-  ) {
-    this.assertEnabled();
-    this.rejectApiKeyPrincipal(req);
-
-    // undefined -> service default (1 year); null -> unlimited; string -> Date.
-    const expiresAt =
-      dto.expiresAt === undefined
-        ? undefined
-        : dto.expiresAt === null
-          ? null
-          : new Date(dto.expiresAt);
-
-    const { token, key } = await this.apiKeyService.create(
-      user,
-      dto.name,
-      expiresAt,
-    );
-
-    this.auditService.log({
-      event: AuditEvent.API_KEY_CREATED,
-      resourceType: AuditResource.API_KEY,
-      resourceId: key.id,
-    });
-    // The durable trail lives in container logs (AUDIT_SERVICE is a Noop in this
-    // build). No token material — the JWT is only ever returned in the response.
-    this.logger.log(
-      `API key created: id=${key.id} name=${JSON.stringify(
-        key.name,
-      )} actor=${user.id} expiresAt=${
-        key.expiresAt ? new Date(key.expiresAt).toISOString() : 'never'
-      } ip=${this.clientIp(req)}`,
-    );
-
-    // Return the token ONCE (never retrievable again) and the computed expiry so
-    // the caller/UI can surface "expires <date>" (the year-default time-bomb
-    // early-warning).
-    return {
-      token,
-      apiKey: {
-        id: key.id,
-        name: key.name,
-        expiresAt: key.expiresAt,
-        createdAt: key.createdAt,
-      },
-    };
-  }
-
-  @HttpCode(HttpStatus.OK)
-  @Post('list')
-  async list(
-    @AuthUser() user: User,
-    @AuthWorkspace() workspace: Workspace,
-    @Req() req: FastifyRequest,
-  ) {
-    this.assertEnabled();
-    this.rejectApiKeyPrincipal(req);
-
-    const ability = this.workspaceAbility.createForUser(user, workspace);
-    // Admin (CASL Manage on API): every key in the workspace, attributed to its
-    // creator (closes "a leaked key of a departed employee is invisible"). A
-    // member sees only their own.
-    const keys = ability.can(
-      WorkspaceCaslAction.Manage,
-      WorkspaceCaslSubject.API,
-    )
-      ? await this.apiKeyRepo.findAllInWorkspace(workspace.id)
-      : await this.apiKeyRepo.findByCreator(user.id, workspace.id);
-
-    return keys.map((k) => ({
-      id: k.id,
-      name: k.name,
-      expiresAt: k.expiresAt,
-      lastUsedAt: k.lastUsedAt,
-      createdAt: k.createdAt,
-      creator: k.creator,
-    }));
-  }
-
-  @HttpCode(HttpStatus.OK)
-  @Post('revoke')
-  async revoke(
-    @Body() dto: RevokeApiKeyDto,
-    @AuthUser() user: User,
-    @AuthWorkspace() workspace: Workspace,
-    @Req() req: FastifyRequest,
-  ) {
-    this.assertEnabled();
-    this.rejectApiKeyPrincipal(req);
-
-    const key = await this.apiKeyRepo.findById(dto.id, workspace.id);
-    // Uniform 404 for a missing/already-revoked key regardless of who asks — no
-    // existence oracle for keys the caller may not own.
-    if (!key) {
-      throw new NotFoundException();
-    }
-
-    const ability = this.workspaceAbility.createForUser(user, workspace);
-    const canManageAll = ability.can(
-      WorkspaceCaslAction.Manage,
-      WorkspaceCaslSubject.API,
-    );
-    // Admin may revoke any key in the workspace; a member only their own.
-    if (!canManageAll && key.creatorId !== user.id) {
-      throw new ForbiddenException();
-    }
-
-    await this.apiKeyService.revoke(key.id, workspace.id);
-
-    this.auditService.log({
-      event: AuditEvent.API_KEY_DELETED,
-      resourceType: AuditResource.API_KEY,
-      resourceId: key.id,
-    });
-    this.logger.log(
-      `API key revoked: id=${key.id} name=${JSON.stringify(
-        key.name,
-      )} actor=${user.id} ip=${this.clientIp(req)}`,
-    );
-
-    return { success: true };
-  }
-}

apps/server/src/core/api-key/api-key.module.ts

@@ -1,19 +0,0 @@
-import { Module } from '@nestjs/common';
-import { ApiKeyService } from './api-key.service';
-import { ApiKeyController } from './api-key.controller';
-import { TokenModule } from '../auth/token.module';
-
-// Core (non-EE) API-key feature: issuance REST endpoints + the shared validator
-// consumed by jwt.strategy (REST) and McpService (the /mcp Bearer router).
-// DatabaseModule (global) provides ApiKeyRepo/UserRepo/WorkspaceRepo; CaslModule
-// (global) provides WorkspaceAbilityFactory; TokenModule provides TokenService
-// (the no-exp api-key signer). ApiKeyService is exported so AuthModule (for
-// jwt.strategy) and McpModule (for the /mcp router) can inject it directly,
-// replacing the absent EE `ee/api-key` dynamic require.
-@Module({
-  imports: [TokenModule],
-  controllers: [ApiKeyController],
-  providers: [ApiKeyService],
-  exports: [ApiKeyService],
-})
-export class ApiKeyModule {}

apps/server/src/core/api-key/api-key.service.spec.ts

@@ -1,299 +0,0 @@
-import { UnauthorizedException } from '@nestjs/common';
-import { ApiKeyService } from './api-key.service';
-import { JwtType } from '../auth/dto/jwt-payload';
-
-/**
- * Security contract for ApiKeyService.validate — the single validator shared by
- * jwt.strategy (REST) and the /mcp Bearer router.
- *
- * Invariants under test:
- *  - Anti-enumeration: every DEFINITE deny (missing/revoked/expired row, disabled
- *    user, workspace mismatch, kill-switch off) throws the SAME bare
- *    UnauthorizedException — an agent cannot distinguish expired from revoked.
- *  - deny-on-decision / 5xx-on-infra: an UNEXPECTED (infra) error PROPAGATES as
- *    itself (→ 5xx), never masked as a 401.
- *  - Expiry is read from the ROW, never a JWT exp claim.
- *  - No validate cache (each call re-reads the row → immediate revocation).
- */
-
-const APP_SECRET = 'secret';
-
-function makeDeps(over: Partial<Record<string, any>> = {}) {
-  const apiKeyRepo = {
-    findById: jest.fn(),
-    insert: jest.fn(),
-    softDelete: jest.fn(),
-    touchLastUsed: jest.fn().mockResolvedValue(undefined),
-    ...(over.apiKeyRepo ?? {}),
-  };
-  const userRepo = {
-    findById: jest.fn(),
-    ...(over.userRepo ?? {}),
-  };
-  const workspaceRepo = {
-    findById: jest.fn().mockResolvedValue({ id: 'ws-1' }),
-    ...(over.workspaceRepo ?? {}),
-  };
-  const tokenService = {
-    generateApiToken: jest.fn().mockResolvedValue('minted.jwt.token'),
-    ...(over.tokenService ?? {}),
-  };
-  const environmentService = {
-    isApiKeysEnabled: jest.fn().mockReturnValue(true),
-    getApiKeysEnabledRaw: jest.fn().mockReturnValue(undefined),
-    getAppSecret: jest.fn().mockReturnValue(APP_SECRET),
-    ...(over.environmentService ?? {}),
-  };
-  const service = new (ApiKeyService as unknown as new (...a: any[]) => ApiKeyService)(
-    apiKeyRepo,
-    userRepo,
-    workspaceRepo,
-    tokenService,
-    environmentService,
-  );
-  return { service, apiKeyRepo, userRepo, workspaceRepo, tokenService, environmentService };
-}
-
-const payload = (over: Record<string, any> = {}) => ({
-  sub: 'u-1',
-  workspaceId: 'ws-1',
-  apiKeyId: 'key-1',
-  type: JwtType.API_KEY,
-  ...over,
-});
-
-const activeRow = (over: Record<string, any> = {}) => ({
-  id: 'key-1',
-  name: 'k',
-  creatorId: 'u-1',
-  workspaceId: 'ws-1',
-  expiresAt: null,
-  lastUsedAt: null,
-  deletedAt: null,
-  ...over,
-});
-
-const activeUser = (over: Record<string, any> = {}) => ({
-  id: 'u-1',
-  workspaceId: 'ws-1',
-  deactivatedAt: null,
-  deletedAt: null,
-  isAgent: false,
-  ...over,
-});
-
-describe('ApiKeyService.validate', () => {
-  it('returns { user, workspace } for a valid active key', async () => {
-    const { service, apiKeyRepo, userRepo } = makeDeps();
-    apiKeyRepo.findById.mockResolvedValue(activeRow());
-    userRepo.findById.mockResolvedValue(activeUser());
-
-    const res = await service.validate(payload() as any);
-
-    expect(res.user).toMatchObject({ id: 'u-1' });
-    expect(res.workspace).toMatchObject({ id: 'ws-1' });
-    // Loads isAgent so downstream provenance does not silently degrade.
-    expect(userRepo.findById).toHaveBeenCalledWith('u-1', 'ws-1', {
-      includeIsAgent: true,
-    });
-  });
-
-  // --- Anti-enumeration: every deny is the SAME bare 401 -------------------
-  const denyCases: Array<[string, (d: ReturnType<typeof makeDeps>) => void]> = [
-    [
-      'missing/orphaned row',
-      (d) => d.apiKeyRepo.findById.mockResolvedValue(undefined),
-    ],
-    [
-      'revoked (soft-deleted → findById returns undefined)',
-      (d) => d.apiKeyRepo.findById.mockResolvedValue(undefined),
-    ],
-    [
-      'expired (expires_at in the past)',
-      (d) => {
-        d.apiKeyRepo.findById.mockResolvedValue(
-          activeRow({ expiresAt: new Date(Date.now() - 1000) }),
-        );
-        d.userRepo.findById.mockResolvedValue(activeUser());
-      },
-    ],
-    [
-      'disabled user',
-      (d) => {
-        d.apiKeyRepo.findById.mockResolvedValue(activeRow());
-        d.userRepo.findById.mockResolvedValue(
-          activeUser({ deactivatedAt: new Date() }),
-        );
-      },
-    ],
-    [
-      'user not found',
-      (d) => {
-        d.apiKeyRepo.findById.mockResolvedValue(activeRow());
-        d.userRepo.findById.mockResolvedValue(undefined);
-      },
-    ],
-    [
-      'creator/sub mismatch',
-      (d) => {
-        d.apiKeyRepo.findById.mockResolvedValue(activeRow({ creatorId: 'other' }));
-        d.userRepo.findById.mockResolvedValue(activeUser());
-      },
-    ],
-    [
-      'workspace row missing',
-      (d) => {
-        d.apiKeyRepo.findById.mockResolvedValue(activeRow());
-        d.userRepo.findById.mockResolvedValue(activeUser());
-        d.workspaceRepo.findById.mockResolvedValue(undefined);
-      },
-    ],
-    [
-      'kill-switch OFF',
-      (d) => d.environmentService.isApiKeysEnabled.mockReturnValue(false),
-    ],
-    [
-      'malformed payload (missing apiKeyId)',
-      () => undefined,
-    ],
-  ];
-
-  it.each(denyCases)(
-    'throws a bare UnauthorizedException with NO message for: %s',
-    async (label, arrange) => {
-      const deps = makeDeps();
-      arrange(deps);
-      const p =
-        label === 'malformed payload (missing apiKeyId)'
-          ? (payload({ apiKeyId: undefined }) as any)
-          : (payload() as any);
-
-      const err = await deps.service.validate(p).catch((e) => e);
-      expect(err).toBeInstanceOf(UnauthorizedException);
-      // Anti-enumeration: identical, class-less message for every failure mode.
-      expect(err.message).toBe('Unauthorized');
-    },
-  );
-
-  it('kill-switch OFF denies BEFORE any DB read', async () => {
-    const { service, apiKeyRepo, environmentService } = makeDeps();
-    environmentService.isApiKeysEnabled.mockReturnValue(false);
-
-    await expect(service.validate(payload() as any)).rejects.toBeInstanceOf(
-      UnauthorizedException,
-    );
-    expect(apiKeyRepo.findById).not.toHaveBeenCalled();
-  });
-
-  // --- Infra error propagates (5xx), NOT masked as 401 ---------------------
-  it('PROPAGATES an infra (DB) error instead of masking it as 401', async () => {
-    const { service, apiKeyRepo } = makeDeps();
-    const boom = new Error('connection terminated');
-    apiKeyRepo.findById.mockRejectedValue(boom);
-
-    await expect(service.validate(payload() as any)).rejects.toBe(boom);
-  });
-
-  // --- No cache: revocation is immediate on the next call ------------------
-  it('re-reads the row on every call (no validate cache → immediate revocation)', async () => {
-    const { service, apiKeyRepo, userRepo } = makeDeps();
-    apiKeyRepo.findById.mockResolvedValue(activeRow());
-    userRepo.findById.mockResolvedValue(activeUser());
-    await service.validate(payload() as any);
-
-    // Now the key is revoked (row invisible) — the very next call denies.
-    apiKeyRepo.findById.mockResolvedValue(undefined);
-    await expect(service.validate(payload() as any)).rejects.toBeInstanceOf(
-      UnauthorizedException,
-    );
-    expect(apiKeyRepo.findById).toHaveBeenCalledTimes(2);
-  });
-
-  // --- last_used_at is throttled + fire-and-forget -------------------------
-  it('touches last_used_at when stale and skips when fresh', async () => {
-    const { service, apiKeyRepo, userRepo } = makeDeps();
-    userRepo.findById.mockResolvedValue(activeUser());
-
-    // Stale (never used) -> touch.
-    apiKeyRepo.findById.mockResolvedValue(activeRow({ lastUsedAt: null }));
-    await service.validate(payload() as any);
-    expect(apiKeyRepo.touchLastUsed).toHaveBeenCalledTimes(1);
-
-    // Fresh (used 1 minute ago) -> skip.
-    apiKeyRepo.touchLastUsed.mockClear();
-    apiKeyRepo.findById.mockResolvedValue(
-      activeRow({ lastUsedAt: new Date(Date.now() - 60_000) }),
-    );
-    await service.validate(payload() as any);
-    expect(apiKeyRepo.touchLastUsed).not.toHaveBeenCalled();
-  });
-
-  it('a failing last_used_at touch never fails the request', async () => {
-    const { service, apiKeyRepo, userRepo } = makeDeps();
-    apiKeyRepo.findById.mockResolvedValue(activeRow({ lastUsedAt: null }));
-    userRepo.findById.mockResolvedValue(activeUser());
-    apiKeyRepo.touchLastUsed.mockRejectedValue(new Error('write failed'));
-
-    await expect(service.validate(payload() as any)).resolves.toMatchObject({
-      user: { id: 'u-1' },
-    });
-  });
-});
-
-describe('ApiKeyService.create (mint-then-insert, no exp)', () => {
-  it('mints the JWT BEFORE inserting the row and returns the token once', async () => {
-    const { service, apiKeyRepo, tokenService } = makeDeps();
-    const order: string[] = [];
-    tokenService.generateApiToken.mockImplementation(async () => {
-      order.push('mint');
-      return 'tok';
-    });
-    apiKeyRepo.insert.mockImplementation(async (row: any) => {
-      order.push('insert');
-      return { ...row, createdAt: new Date() };
-    });
-
-    const user = { id: 'u-1', workspaceId: 'ws-1' } as any;
-    const res = await service.create(user, 'my key');
-
-    expect(order).toEqual(['mint', 'insert']);
-    expect(res.token).toBe('tok');
-    // The id is generated before minting and reused for the row.
-    const mintArg = tokenService.generateApiToken.mock.calls[0][0];
-    const insertArg = apiKeyRepo.insert.mock.calls[0][0];
-    expect(mintArg.apiKeyId).toBe(insertArg.id);
-  });
-
-  it('applies the 1-year default when expiresAt is undefined', async () => {
-    const { service, apiKeyRepo } = makeDeps();
-    apiKeyRepo.insert.mockImplementation(async (row: any) => row);
-    const user = { id: 'u-1', workspaceId: 'ws-1' } as any;
-
-    await service.create(user, 'k');
-
-    const insertArg = apiKeyRepo.insert.mock.calls[0][0];
-    const days =
-      (insertArg.expiresAt.getTime() - Date.now()) / (24 * 60 * 60 * 1000);
-    expect(days).toBeGreaterThan(364);
-    expect(days).toBeLessThan(367);
-  });
-
-  it('honors an explicit null (unlimited)', async () => {
-    const { service, apiKeyRepo } = makeDeps();
-    apiKeyRepo.insert.mockImplementation(async (row: any) => row);
-    const user = { id: 'u-1', workspaceId: 'ws-1' } as any;
-
-    await service.create(user, 'k', null);
-
-    expect(apiKeyRepo.insert.mock.calls[0][0].expiresAt).toBeNull();
-  });
-
-  it('does NOT insert a row when minting fails (mint-then-insert → inert)', async () => {
-    const { service, apiKeyRepo, tokenService } = makeDeps();
-    tokenService.generateApiToken.mockRejectedValue(new Error('mint failed'));
-    const user = { id: 'u-1', workspaceId: 'ws-1' } as any;
-
-    await expect(service.create(user, 'k')).rejects.toThrow('mint failed');
-    expect(apiKeyRepo.insert).not.toHaveBeenCalled();
-  });
-});

apps/server/src/core/api-key/api-key.service.ts

@@ -1,191 +0,0 @@
-import {
-  Injectable,
-  Logger,
-  OnModuleInit,
-  UnauthorizedException,
-} from '@nestjs/common';
-import { v7 as uuid7 } from 'uuid';
-import { ApiKeyRepo } from '@docmost/db/repos/api-key/api-key.repo';
-import { UserRepo } from '@docmost/db/repos/user/user.repo';
-import { WorkspaceRepo } from '@docmost/db/repos/workspace/workspace.repo';
-import { TokenService } from '../auth/services/token.service';
-import { EnvironmentService } from '../../integrations/environment/environment.service';
-import { JwtApiKeyPayload } from '../auth/dto/jwt-payload';
-import { ApiKey, User, Workspace } from '@docmost/db/types/entity.types';
-import { isUserDisabled } from '../../common/helpers';
-
-// Default lifetime for a new key when the caller does not specify one: 1 year.
-// The owner runs a homelab where agents live for years; forcing rotation is
-// operational pain, so an explicit `null` (unlimited) is also allowed.
-const DEFAULT_LIFETIME_MS = 365 * 24 * 60 * 60 * 1000;
-
-// last_used_at is a best-effort forensics stamp, not an access record; we only
-// refresh it when it is older than this to avoid a write on every request. The
-// 1h resolution is a deliberate constant (forensics granularity, not accounting).
-const LAST_USED_THROTTLE_MS = 60 * 60 * 1000;
-
-/**
- * Core API-key lifecycle service. Owns minting (create), the single validator
- * shared by BOTH the REST jwt.strategy path and the /mcp Bearer path (validate),
- * and revocation (revoke). The `api_keys` ROW is the sole source of truth for a
- * key's lifetime and revocation — never the JWT, which carries no `exp` claim.
- */
-@Injectable()
-export class ApiKeyService implements OnModuleInit {
-  private readonly logger = new Logger(ApiKeyService.name);
-
-  constructor(
-    private readonly apiKeyRepo: ApiKeyRepo,
-    private readonly userRepo: UserRepo,
-    private readonly workspaceRepo: WorkspaceRepo,
-    private readonly tokenService: TokenService,
-    private readonly environmentService: EnvironmentService,
-  ) {}
-
-  onModuleInit() {
-    // Boot log so the kill-switch state after each deploy is verifiable in logs.
-    const enabled = this.environmentService.isApiKeysEnabled();
-    const raw = this.environmentService.getApiKeysEnabledRaw();
-    this.logger.log(
-      `API keys: ${enabled ? 'ENABLED' : 'DISABLED'} (API_KEYS_ENABLED=${
-        raw ?? 'unset'
-      })`,
-    );
-  }
-
-  /**
-   * Mint a new key for `user`. mint-then-insert ordering (R1):
-   *   1. generate the id first — it must be in the JWT payload before the row.
-   *   2. mint the JWT (no `exp` claim). A mint failure aborts before any row is
-   *      written (inert), so a half-created key cannot exist.
-   *   3. insert the row last. A lost response leaves an orphaned row that is
-   *      visible in `list` and self-heals (the user revokes it).
-   * The token is returned ONCE and never stored — the JWT is self-contained, so
-   * no token material lives in the table.
-   *
-   * `expiresAt`: `undefined` -> default 1 year; `null` -> unlimited (explicit);
-   * a Date -> that instant (a past date is rejected at the DTO layer).
-   */
-  async create(
-    user: User,
-    name: string,
-    expiresAt?: Date | null,
-  ): Promise<{ token: string; key: ApiKey }> {
-    const resolvedExpiresAt =
-      expiresAt === undefined
-        ? new Date(Date.now() + DEFAULT_LIFETIME_MS)
-        : expiresAt;
-
-    const apiKeyId = uuid7();
-
-    const token = await this.tokenService.generateApiToken({
-      apiKeyId,
-      user,
-      workspaceId: user.workspaceId,
-    });
-
-    const key = await this.apiKeyRepo.insert({
-      id: apiKeyId,
-      name,
-      creatorId: user.id,
-      workspaceId: user.workspaceId,
-      expiresAt: resolvedExpiresAt,
-    });
-
-    return { token, key };
-  }
-
-  /**
-   * The single validator for an api-key principal, shared by jwt.strategy and the
-   * /mcp Bearer router. Returns `{ user, workspace }` (the same shape the access
-   * path returns) so the AuthUser/AuthWorkspace decorators and MCP identity work
-   * unchanged.
-   *
-   * Failure semantics (R4, anti-enumeration): a DEFINITE negative fact — feature
-   * disabled, missing/revoked/expired row, workspace mismatch, disabled user —
-   * throws a bare `UnauthorizedException` (a single generic 401 for every case;
-   * an agent cannot distinguish expired from revoked, and its reaction is
-   * identical). An UNEXPECTED (infra) error is NOT caught here: it propagates so
-   * the surface returns 5xx, never a masked 401 (deny-on-decision / 5xx-on-infra).
-   * There is NO validate cache: it is 2–3 PK lookups (~1ms), two orders of
-   * magnitude cheaper than the bcrypt it replaces; a cache would only add a
-   * Date-serialization trap and a revocation lag. Revocation is immediate.
-   */
-  async validate(
-    payload: JwtApiKeyPayload,
-  ): Promise<{ user: User; workspace: Workspace }> {
-    // Kill-switch OFF: deny unconditionally (same generic 401). Note the
-    // endpoints additionally 404 at the controller; here we deny the token.
-    if (!this.environmentService.isApiKeysEnabled()) {
-      throw new UnauthorizedException();
-    }
-
-    if (!payload?.apiKeyId || !payload?.sub || !payload?.workspaceId) {
-      throw new UnauthorizedException();
-    }
-
-    const row = await this.apiKeyRepo.findById(
-      payload.apiKeyId,
-      payload.workspaceId,
-    );
-    // Absent row = revoked (soft-deleted, invisible to findById), orphaned
-    // (creator/workspace cascade-deleted), or never existed. All terminal deny.
-    if (!row) {
-      throw new UnauthorizedException();
-    }
-
-    // Expiry is read from the ROW, never an `exp` JWT claim.
-    if (row.expiresAt && row.expiresAt.getTime() <= Date.now()) {
-      throw new UnauthorizedException();
-    }
-
-    const user = await this.userRepo.findById(
-      payload.sub,
-      payload.workspaceId,
-      { includeIsAgent: true },
-    );
-    if (!user || isUserDisabled(user)) {
-      throw new UnauthorizedException();
-    }
-
-    // The key acts only as its creator (defence in depth against a token whose
-    // signed `sub` ever drifted from the row's owner).
-    if (row.creatorId !== user.id) {
-      throw new UnauthorizedException();
-    }
-
-    const workspace = await this.workspaceRepo.findById(payload.workspaceId);
-    if (!workspace) {
-      throw new UnauthorizedException();
-    }
-
-    // Best-effort, throttled, fire-and-forget forensics stamp AFTER all checks.
-    this.touchLastUsed(row);
-
-    return { user, workspace };
-  }
-
-  /**
-   * Revoke (soft-delete) a key. Authorization/ownership is decided by the caller
-   * (the controller, via CASL); this only performs the terminal write. Idempotent:
-   * a second revoke is a no-op (the row is already invisible).
-   */
-  async revoke(id: string, workspaceId: string): Promise<void> {
-    await this.apiKeyRepo.softDelete(id, workspaceId);
-  }
-
-  // Throttled best-effort last_used_at bump: skip if it was touched within the
-  // window; otherwise fire-and-forget so a stamp write never fails or slows the
-  // request (mirrors SessionActivityService.trackActivity).
-  private touchLastUsed(row: ApiKey): void {
-    const last = row.lastUsedAt ? new Date(row.lastUsedAt).getTime() : 0;
-    if (Date.now() - last < LAST_USED_THROTTLE_MS) return;
-    void this.apiKeyRepo.touchLastUsed(row.id).catch((err) => {
-      this.logger.warn(
-        `Failed to update api_key last_used_at for ${row.id}: ${
-          (err as Error)?.message ?? err
-        }`,
-      );
-    });
-  }
-}

apps/server/src/core/api-key/dto/create-api-key.dto.ts

@@ -1,51 +0,0 @@
-import {
-  IsDateString,
-  IsNotEmpty,
-  IsOptional,
-  IsString,
-  MaxLength,
-  registerDecorator,
-  ValidationOptions,
-} from 'class-validator';
-
-/**
- * `expiresAt` must be strictly in the future. Skips validation for `null`
- * (explicit "unlimited") and `undefined` (server applies the 1-year default),
- * so only an actually-supplied date is range-checked. Rejecting a PAST date at
- * the DTO layer means a caller cannot mint an already-dead key.
- */
-function IsFutureDateString(options?: ValidationOptions) {
-  return function (object: object, propertyName: string) {
-    registerDecorator({
-      name: 'isFutureDateString',
-      target: object.constructor,
-      propertyName,
-      options: {
-        message: 'expiresAt must be a date in the future',
-        ...options,
-      },
-      validator: {
-        validate(value: unknown) {
-          if (value === null || value === undefined) return true;
-          if (typeof value !== 'string') return false;
-          const t = Date.parse(value);
-          return !Number.isNaN(t) && t > Date.now();
-        },
-      },
-    });
-  };
-}
-
-export class CreateApiKeyDto {
-  @IsString()
-  @IsNotEmpty()
-  @MaxLength(255)
-  name: string;
-
-  // undefined -> default 1 year (applied server-side); null -> unlimited
-  // (explicit); an ISO date string -> that instant, which must be in the future.
-  @IsOptional()
-  @IsDateString()
-  @IsFutureDateString()
-  expiresAt?: string | null;
-}

apps/server/src/core/api-key/dto/revoke-api-key.dto.ts

@@ -1,6 +0,0 @@
-import { IsUUID } from 'class-validator';
-
-export class RevokeApiKeyDto {
-  @IsUUID()
-  id: string;
-}

apps/server/src/core/auth/auth.controller.spec.ts

@@ -20,71 +20,3 @@ describe('AuthController', () => {
     expect(controller).toBeDefined();
   });
 });
-
-// The collab-token handler is the ARM-SEAM of the #501 anti-laundering defense:
-// it derives the api-key origin args ({ apiKeyId }) from the SIGNED-derived
-// `req.raw` fields (stamped by jwt.strategy) and threads them into the collab
-// token mint, forcing principal='api_key' so a later key revoke rejects NEW
-// collab connections. A regression here (reading `body`, dropping `apiKeyId`,
-// or inverting the ternary) would silently mint api-key requests as
-// principal='session' and reopen the laundering hole with a green suite. These
-// tests pin the EXACT args the controller passes to authService.getCollabToken.
-describe('AuthController.collabToken arms the api-key origin (#501)', () => {
-  let controller: AuthController;
-  let getCollabToken: jest.Mock;
-
-  const user = { id: 'u-1', workspaceId: 'ws-1' } as any;
-  const workspace = { id: 'ws-1' } as any;
-
-  beforeEach(() => {
-    getCollabToken = jest.fn().mockResolvedValue({ token: 'ct' });
-    controller = new AuthController(
-      { getCollabToken } as any, // authService
-      {} as any, // sessionService
-      {} as any, // environmentService
-      {} as any, // moduleRef
-      {} as any, // auditService
-    );
-  });
-
-  it('threads { apiKeyId } when req.raw marks an api_key principal (ARMED)', async () => {
-    const req = { raw: { authType: 'api_key', apiKeyId: 'k-1' } } as any;
-
-    await controller.collabToken(user, workspace, req);
-
-    expect(getCollabToken).toHaveBeenCalledTimes(1);
-    expect(getCollabToken).toHaveBeenCalledWith(user, 'ws-1', {
-      apiKeyId: 'k-1',
-    });
-  });
-
-  it('passes undefined for a normal session request (NOT armed)', async () => {
-    const req = { raw: {} } as any;
-
-    await controller.collabToken(user, workspace, req);
-
-    expect(getCollabToken).toHaveBeenCalledTimes(1);
-    expect(getCollabToken).toHaveBeenCalledWith(user, 'ws-1', undefined);
-  });
-
-  it('does NOT arm when api_key authType lacks an apiKeyId', async () => {
-    const req = { raw: { authType: 'api_key' } } as any;
-
-    await controller.collabToken(user, workspace, req);
-
-    expect(getCollabToken).toHaveBeenCalledWith(user, 'ws-1', undefined);
-  });
-
-  it('reads the SIGNED req.raw, never a spoofable client body', async () => {
-    // A client-supplied body claiming api_key must be ignored: only the
-    // jwt.strategy-stamped req.raw can arm the api-key origin.
-    const req = {
-      raw: {},
-      body: { authType: 'api_key', apiKeyId: 'attacker' },
-    } as any;
-
-    await controller.collabToken(user, workspace, req);
-
-    expect(getCollabToken).toHaveBeenCalledWith(user, 'ws-1', undefined);
-  });
-});

apps/server/src/core/auth/auth.controller.ts

@@ -207,20 +207,8 @@ export class AuthController {
   async collabToken(
     @AuthUser() user: User,
     @AuthWorkspace() workspace: Workspace,
-    @Req() req: FastifyRequest,
   ) {
-    // Thread the api-key origin (#501): when the requester authenticated with an
-    // api key (jwt.strategy stamped req.raw.authType/apiKeyId), the minted collab
-    // token carries principal='api_key' + apiKeyId so a later revoke of the key
-    // rejects NEW collab connections. A normal session request mints a
-    // principal='session' token. Reading the SIGNED-derived req.raw fields (never
-    // a client body) keeps it unspoofable.
-    const raw = req.raw as { authType?: string; apiKeyId?: string };
-    const apiKey =
-      raw.authType === 'api_key' && raw.apiKeyId
-        ? { apiKeyId: raw.apiKeyId }
-        : undefined;
-    return this.authService.getCollabToken(user, workspace.id, apiKey);
+    return this.authService.getCollabToken(user, workspace.id);
   }
 
   @SkipThrottle({ [AUTH_THROTTLER]: true })

apps/server/src/core/auth/auth.module.ts

@@ -5,13 +5,9 @@ import { JwtStrategy } from './strategies/jwt.strategy';
 import { WorkspaceModule } from '../workspace/workspace.module';
 import { SignupService } from './services/signup.service';
 import { TokenModule } from './token.module';
-import { ApiKeyModule } from '../api-key/api-key.module';
 
 @Module({
-  // ApiKeyModule supplies ApiKeyService, injected into JwtStrategy so an
-  // api_key Bearer/cookie token is validated directly (replacing the absent EE
-  // `ee/api-key` dynamic require).
-  imports: [TokenModule, WorkspaceModule, ApiKeyModule],
+  imports: [TokenModule, WorkspaceModule],
   controllers: [AuthController],
   providers: [AuthService, SignupService, JwtStrategy],
   exports: [SignupService, AuthService],

apps/server/src/core/auth/dto/jwt-payload.ts

@@ -33,17 +33,6 @@ export type JwtPayload = {
   aiChatId?: string | null;
 };
 
-// The AUTH-PRINCIPAL kind behind a collab token — distinct from `actor`
-// (provenance). Stamped into EVERY newly-minted collab token (#501): 'session'
-// for a normal user/session (incl. the internal AI agent, which is session-
-// backed), 'api_key' when the token was minted by an api-key principal (an
-// external MCP agent). The discriminator keys on the token's ORIGIN, so the
-// collab seam can re-check a revoked api key on connect and reject a claimless
-// token after the rollout grace window. NOT keyed on `actor:'agent'` — the
-// internal agent is 'agent' but session-backed, so it must stay on the no-check
-// path.
-export type CollabPrincipal = 'session' | 'api_key';
-
 export type JwtCollabPayload = {
   sub: string;
   workspaceId: string;
@@ -55,13 +44,6 @@ export type JwtCollabPayload = {
   // Nullable: an external MCP agent has no internal ai_chats row, so it carries
   // an 'agent' actor with a null aiChatId.
   aiChatId?: string | null;
-  // Auth-principal discriminator (#501). Present on every post-rollout token;
-  // its absence on a still-valid token past the grace window is treated as an
-  // error, not trust (fail-closed).
-  principal?: CollabPrincipal;
-  // Only when principal === 'api_key': the minting key's id, so the collab seam
-  // can row-check (and reject) a revoked key on connect.
-  apiKeyId?: string;
 };
 
 export type JwtExchangePayload = {

apps/server/src/core/auth/services/auth.service.ts

@@ -375,20 +375,10 @@ export class AuthService {
     }
   }
 
-  async getCollabToken(
-    user: User,
-    workspaceId: string,
-    // Origin of the request minting this collab token (#501). When the caller is
-    // an api-key principal, its apiKeyId is threaded into the token so the collab
-    // seam can re-check the key on connect (closing api-key -> long-lived-collab
-    // laundering). Absent for a normal session/human request.
-    apiKey?: { apiKeyId: string },
-  ) {
+  async getCollabToken(user: User, workspaceId: string) {
     const token = await this.tokenService.generateCollabToken(
       user,
       workspaceId,
-      undefined,
-      apiKey,
     );
     return { token };
   }

apps/server/src/core/auth/services/token.service.behavior.spec.ts

@@ -1,5 +1,4 @@
 import { ForbiddenException, UnauthorizedException } from '@nestjs/common';
-import * as jwt from 'jsonwebtoken';
 import { TokenService } from './token.service';
 import { JwtType } from '../dto/jwt-payload';
 
@@ -214,175 +213,4 @@ describe('TokenService.generateCollabToken', () => {
       aiChatId: 'chat-456',
     });
   });
-
-  // #501 fail-closed discriminator: EVERY collab token carries a principal.
-  it("defaults principal to 'session' with NO apiKeyId (normal/internal-agent path)", async () => {
-    const { service, jwtService } = makeTokenService();
-    await service.generateCollabToken(makeUser() as never, 'ws-1');
-    const [payload] = jwtService.sign.mock.calls[0];
-    expect(payload.principal).toBe('session');
-    expect(payload).not.toHaveProperty('apiKeyId');
-  });
-
-  it("the internal agent (provenance, NO apiKey) still gets principal='session'", async () => {
-    const { service, jwtService } = makeTokenService();
-    await service.generateCollabToken(
-      makeUser() as never,
-      'ws-1',
-      { actor: 'agent', aiChatId: 'chat-1' },
-    );
-    const [payload] = jwtService.sign.mock.calls[0];
-    // Keyed on api-key ORIGIN, not actor: an is_agent session token is 'session'.
-    expect(payload.principal).toBe('session');
-    expect(payload).not.toHaveProperty('apiKeyId');
-  });
-
-  it("stamps principal='api_key' + apiKeyId when minted by an api-key principal", async () => {
-    const { service, jwtService } = makeTokenService();
-    await service.generateCollabToken(
-      makeUser() as never,
-      'ws-1',
-      undefined,
-      { apiKeyId: 'key-9' },
-    );
-    const [payload] = jwtService.sign.mock.calls[0];
-    expect(payload.principal).toBe('api_key');
-    expect(payload.apiKeyId).toBe('key-9');
-  });
 });
-
-/**
- * API-key token minting MUST carry NO `exp` claim — the ONLY source of truth for
- * a key's lifetime/revocation is its `api_keys` row, checked on every request.
- *
- * This is a LIVE-bug regression: the shared JwtService is registered with a
- * global `signOptions.expiresIn` (default '90d') that merges into every sign(),
- * so an api-key minted through it silently gets exp=now+90d and an "unlimited"
- * key dies in 90 days. TokenService.generateApiToken mints through a dedicated
- * no-expiry signer instead. These tests construct the REAL signer (a real secret
- * via the stubbed EnvironmentService) and decode the produced JWT to assert the
- * observable property: no `exp`.
- */
-describe('TokenService.generateApiToken (no exp claim ever)', () => {
-  const APP_SECRET_LOCAL = 'apikey-secret';
-
-  function makeRealSignerService() {
-    // Give the SHARED jwtService a global expiresIn so a regression (minting
-    // through it) would show up as an exp claim — the exact live bug.
-    const { JwtService } = require('@nestjs/jwt');
-    const sharedJwt = new JwtService({
-      secret: APP_SECRET_LOCAL,
-      signOptions: { expiresIn: '90d', issuer: 'Docmost' },
-    });
-    const environmentService = {
-      getAppSecret: () => APP_SECRET_LOCAL,
-    };
-    const service = new (TokenService as unknown as new (
-      ...args: unknown[]
-    ) => TokenService)(sharedJwt, environmentService);
-    return { service };
-  }
-
-  const user = makeUser({ id: 'svc-1', workspaceId: 'ws-1' });
-
-  it('mints an api-key JWT with NO exp claim and issuer Docmost', async () => {
-    const { service } = makeRealSignerService();
-
-    const token = await service.generateApiToken({
-      apiKeyId: 'key-1',
-      user: user as never,
-      workspaceId: 'ws-1',
-    });
-
-    const decoded = jwt.decode(token) as Record<string, unknown>;
-    // The observable security property: no expiry lives in the JWT.
-    expect(decoded.exp).toBeUndefined();
-    expect(decoded).toMatchObject({
-      sub: 'svc-1',
-      apiKeyId: 'key-1',
-      workspaceId: 'ws-1',
-      type: JwtType.API_KEY,
-      iss: 'Docmost',
-    });
-  });
-
-  it('demonstrates the live bug it guards: the SHARED signer WOULD add exp', () => {
-    const { JwtService } = require('@nestjs/jwt');
-    const sharedJwt = new JwtService({
-      secret: APP_SECRET_LOCAL,
-      signOptions: { expiresIn: '90d', issuer: 'Docmost' },
-    });
-    // Even with empty per-call options the global expiresIn merges in.
-    const leaky = sharedJwt.sign({ sub: 'x', type: JwtType.API_KEY }, {});
-    expect((jwt.decode(leaky) as Record<string, unknown>).exp).toBeDefined();
-  });
-
-  it('refuses to mint for a disabled user', async () => {
-    const { service } = makeRealSignerService();
-    await expect(
-      service.generateApiToken({
-        apiKeyId: 'key-1',
-        user: makeUser({ deactivatedAt: new Date() }) as never,
-        workspaceId: 'ws-1',
-      }),
-    ).rejects.toBeInstanceOf(ForbiddenException);
-  });
-});
-
-/**
- * verifyJwtOneOf is the type-routing primitive: verify the signature once and
- * assert the token type is on an explicit allowlist. It must NOT degrade into a
- * "return whatever type" helper — a token whose type is off the allowlist is
- * rejected with the same generic error as a single-type mismatch.
- */
-describe('TokenService.verifyJwtOneOf (allowlist type-routing)', () => {
-  it('returns the payload when the type is on the allowlist', async () => {
-    const verifyAsync = jest
-      .fn()
-      .mockResolvedValue({ type: JwtType.API_KEY, sub: 'u-1' });
-    const { service } = makeTokenService({ verifyAsync });
-
-    const payload = await service.verifyJwtOneOf(token123(), [
-      JwtType.ACCESS,
-      JwtType.API_KEY,
-    ]);
-
-    expect(payload).toMatchObject({ type: JwtType.API_KEY, sub: 'u-1' });
-    expect(verifyAsync).toHaveBeenCalledTimes(1);
-  });
-
-  it('accepts the OTHER allowed type too', async () => {
-    const verifyAsync = jest
-      .fn()
-      .mockResolvedValue({ type: JwtType.ACCESS, sub: 'u-1' });
-    const { service } = makeTokenService({ verifyAsync });
-
-    await expect(
-      service.verifyJwtOneOf(token123(), [JwtType.ACCESS, JwtType.API_KEY]),
-    ).resolves.toMatchObject({ type: JwtType.ACCESS });
-  });
-
-  it('rejects a token whose type is OFF the allowlist (confused-deputy guard)', async () => {
-    const verifyAsync = jest
-      .fn()
-      .mockResolvedValue({ type: JwtType.COLLAB, sub: 'u-1' });
-    const { service } = makeTokenService({ verifyAsync });
-
-    await expect(
-      service.verifyJwtOneOf(token123(), [JwtType.ACCESS, JwtType.API_KEY]),
-    ).rejects.toBeInstanceOf(UnauthorizedException);
-  });
-
-  it('verifies the signature exactly ONCE', async () => {
-    const verifyAsync = jest
-      .fn()
-      .mockResolvedValue({ type: JwtType.ACCESS });
-    const { service } = makeTokenService({ verifyAsync });
-    await service.verifyJwtOneOf(token123(), [JwtType.ACCESS]);
-    expect(verifyAsync).toHaveBeenCalledTimes(1);
-  });
-});
-
-function token123(): string {
-  return 'a.b.c';
-}

apps/server/src/core/auth/services/token.service.ts

@@ -4,6 +4,7 @@ import {
   UnauthorizedException,
 } from '@nestjs/common';
 import { JwtService } from '@nestjs/jwt';
+import type { StringValue } from 'ms';
 import { EnvironmentService } from '../../../integrations/environment/environment.service';
 import {
   JwtApiKeyPayload,
@@ -61,12 +62,6 @@ export class TokenService {
     // token carries no actor/aiChatId and is treated as 'user' downstream.
     // aiChatId is nullable for an external agent with no internal ai_chats row.
     provenance?: { actor: 'agent'; aiChatId: string | null },
-    // Optional api-key origin (#501). When the collab token is minted by an
-    // api-key principal (an external MCP agent), the caller passes the key id so
-    // the token carries principal='api_key' + apiKeyId and the collab seam can
-    // re-check the key on connect. Absent -> principal='session' (a normal
-    // user/session, including the internal session-backed AI agent).
-    apiKey?: { apiKeyId: string },
   ): Promise<string> {
     if (isUserDisabled(user)) {
       throw new ForbiddenException();
@@ -76,10 +71,6 @@ export class TokenService {
       sub: user.id,
       workspaceId,
       type: JwtType.COLLAB,
-      // Fail-closed discriminator on EVERY minted token: 'api_key' when minted by
-      // an api-key principal, else 'session'.
-      principal: apiKey ? 'api_key' : 'session',
-      ...(apiKey ? { apiKeyId: apiKey.apiKeyId } : {}),
       ...(provenance
         ? { actor: provenance.actor, aiChatId: provenance.aiChatId }
         : {}),
@@ -132,8 +123,9 @@ export class TokenService {
     apiKeyId: string;
     user: User;
     workspaceId: string;
+    expiresIn?: StringValue | number;
   }): Promise<string> {
-    const { apiKeyId, user, workspaceId } = opts;
+    const { apiKeyId, user, workspaceId, expiresIn } = opts;
     if (isUserDisabled(user)) {
       throw new ForbiddenException();
     }
@@ -145,32 +137,7 @@ export class TokenService {
       type: JwtType.API_KEY,
     };
 
-    // API-key tokens carry NO `exp` claim EVER — the ONLY source of truth for a
-    // key's lifetime and revocation is its `api_keys` row (checked on every
-    // request), not the JWT. This CANNOT use `this.jwtService`: TokenModule
-    // registers it with a global `signOptions.expiresIn` (JWT_TOKEN_EXPIRES_IN,
-    // default '90d'), which merges into EVERY sign() call — even `sign(payload,
-    // {})` — and `{ expiresIn: undefined }` THROWS rather than stripping it
-    // (verified empirically). So an "unlimited" key minted through the shared
-    // signer would silently get exp=now+90d and die in 90 days regardless of its
-    // row. We mint through a dedicated no-expiry signer, re-stamping only
-    // `issuer: 'Docmost'` for claim parity with the shared signer.
-    return this.apiKeyJwtService().sign(payload);
-  }
-
-  // Lazily-built JWT signer for API-key tokens: same APP_SECRET, same 'Docmost'
-  // issuer, but WITHOUT the global `expiresIn` — so minted API-key tokens have no
-  // `exp` claim. Built once and cached. Verification still goes through the
-  // shared verifier (same secret); `verifyAsync` does not require an `exp`.
-  private _apiKeyJwtService?: JwtService;
-  private apiKeyJwtService(): JwtService {
-    if (!this._apiKeyJwtService) {
-      this._apiKeyJwtService = new JwtService({
-        secret: this.environmentService.getAppSecret(),
-        signOptions: { issuer: 'Docmost' },
-      });
-    }
-    return this._apiKeyJwtService;
+    return this.jwtService.sign(payload, expiresIn ? { expiresIn } : {});
   }
 
   async generatePdfRenderToken(
@@ -210,31 +177,4 @@ export class TokenService {
 
     return payload;
   }
-
-  /**
-   * Verify a token's signature ONCE and assert its `type` is one of `allowed`.
-   *
-   * This is the type-routing primitive for surfaces that legitimately accept
-   * more than one token type on the same Bearer slot (the /mcp Bearer path
-   * accepts both an ACCESS and an API_KEY token). It is deliberately NOT a
-   * "verify-and-return-whatever-type" helper — that would be a reusable
-   * confused-deputy footgun (any caller could then feed an attachment/collab
-   * token where an access token is expected). An explicit allowlist preserves
-   * the type-pinning property of `verifyJwt`: a token whose `type` is not in the
-   * allowlist is rejected with the SAME generic error as a type mismatch, and
-   * the signature is verified exactly once (no double-verify).
-   */
-  async verifyJwtOneOf(token: string, allowed: JwtType[]) {
-    const payload = await this.jwtService.verifyAsync(token, {
-      secret: this.environmentService.getAppSecret(),
-    });
-
-    if (!allowed.includes(payload.type)) {
-      throw new UnauthorizedException(
-        'Invalid JWT token. Token type does not match.',
-      );
-    }
-
-    return payload;
-  }
 }

apps/server/src/core/auth/strategies/jwt.strategy.spec.ts

@@ -22,8 +22,7 @@ describe('JwtStrategy — provenance derivation', () => {
     const userSessionRepo: any = { findActiveById: jest.fn() };
     const sessionActivityService: any = { trackActivity: jest.fn() };
     const environmentService: any = { getAppSecret: () => 'test-secret' };
-    // ACCESS-path tests never touch the api-key seam; a bare stub suffices.
-    const apiKeyService: any = { validate: jest.fn() };
+    const moduleRef: any = {};
 
     const strategy = new JwtStrategy(
       userRepo,
@@ -31,7 +30,7 @@ describe('JwtStrategy — provenance derivation', () => {
       userSessionRepo,
       sessionActivityService,
       environmentService,
-      apiKeyService,
+      moduleRef,
     );
     return { strategy, userRepo };
   }
@@ -123,29 +122,25 @@ describe('JwtStrategy — provenance derivation', () => {
 });
 
 /**
- * Provenance derivation on the API-KEY path (jwt.strategy.validateApiKey, #486
- * + #501).
+ * Provenance derivation on the API-KEY path (jwt.strategy.validateApiKey, #486).
  *
  * The access-token path stamped provenance; the API-key path returned early
  * WITHOUT it, so an is_agent API key's REST writes recorded no 'agent' marker.
  * The API-key payload carries no signed claim, so provenance is resolved from the
- * SERVER-SIDE user returned by ApiKeyService.validate: isAgent -> 'agent',
+ * SERVER-SIDE user returned by ApiKeyService.validateApiKey: isAgent -> 'agent',
  * otherwise 'user'; aiChatId is always null (an API key has no ai_chats row).
  *
- * #501 wires the CORE ApiKeyService (the EE `ee/api-key` module is absent in the
- * fork) directly into the strategy — no dynamic require. The strategy also stamps
- * `req.raw.authType='api_key'` + `req.raw.apiKeyId` for the "a token cannot manage
- * tokens" guard on the /api-keys surface.
+ * The enterprise ApiKeyService is not bundled in the OSS build, so the strategy
+ * loads it through an overridable `resolveApiKeyService` seam that we stub here.
  */
-describe('JwtStrategy — API-key provenance derivation (#486/#501)', () => {
-  function makeApiKeyStrategy(validateImpl: (p: any) => Promise<any>) {
+describe('JwtStrategy — API-key provenance derivation (#486)', () => {
+  function makeApiKeyStrategy(validateApiKeyImpl: (p: any) => Promise<any>) {
     const userRepo: any = { findById: jest.fn() };
     const workspaceRepo: any = { findById: jest.fn() };
     const userSessionRepo: any = { findActiveById: jest.fn() };
     const sessionActivityService: any = { trackActivity: jest.fn() };
     const environmentService: any = { getAppSecret: () => 'test-secret' };
-    const validate = jest.fn(validateImpl);
-    const apiKeyService: any = { validate };
+    const moduleRef: any = {};
 
     const strategy = new JwtStrategy(
       userRepo,
@@ -153,9 +148,14 @@ describe('JwtStrategy — API-key provenance derivation (#486/#501)', () => {
       userSessionRepo,
       sessionActivityService,
       environmentService,
-      apiKeyService,
+      moduleRef,
     );
-    return { strategy, validate };
+    // Stub the EE ApiKeyService seam (the real module is not in the OSS build).
+    const validateApiKey = jest.fn(validateApiKeyImpl);
+    jest
+      .spyOn(strategy as any, 'resolveApiKeyService')
+      .mockReturnValue({ validateApiKey });
+    return { strategy, validateApiKey };
   }
 
   const makeReq = () => ({ raw: {} as Record<string, any> });
@@ -166,23 +166,22 @@ describe('JwtStrategy — API-key provenance derivation (#486/#501)', () => {
     type: JwtType.API_KEY,
   });
 
-  it("stamps actor='agent' + authType/apiKeyId for an is_agent API key", async () => {
+  it("stamps actor='agent' for an is_agent API key (from the validated user)", async () => {
     const validated = {
       user: { id: 'svc-1', isAgent: true },
       workspace: { id: 'ws-1' },
     };
-    const { strategy, validate } = makeApiKeyStrategy(async () => validated);
+    const { strategy, validateApiKey } = makeApiKeyStrategy(
+      async () => validated,
+    );
     const req = makeReq();
 
     const result = await strategy.validate(req, apiKeyPayload() as any);
 
-    expect(validate).toHaveBeenCalledTimes(1);
+    expect(validateApiKey).toHaveBeenCalledTimes(1);
     expect(req.raw.actor).toBe('agent');
     // API keys carry no internal ai_chats row -> null.
     expect(req.raw.aiChatId).toBeNull();
-    // Principal-kind markers for the management-surface guard.
-    expect(req.raw.authType).toBe('api_key');
-    expect(req.raw.apiKeyId).toBe('key-1');
     // The validated auth object is returned unchanged (req.user shape preserved).
     expect(result).toBe(validated);
   });
@@ -198,19 +197,25 @@ describe('JwtStrategy — API-key provenance derivation (#486/#501)', () => {
 
     expect(req.raw.actor).toBe('user');
     expect(req.raw.aiChatId).toBeNull();
-    expect(req.raw.authType).toBe('api_key');
   });
 
-  it('propagates a validate() rejection and stamps nothing', async () => {
-    const { strategy } = makeApiKeyStrategy(async () => {
-      throw new UnauthorizedException();
-    });
+  it('throws Unauthorized (and stamps nothing) when the EE module is missing', async () => {
+    const userRepo: any = { findById: jest.fn() };
+    const strategy = new JwtStrategy(
+      userRepo,
+      { findById: jest.fn() } as any,
+      { findActiveById: jest.fn() } as any,
+      { trackActivity: jest.fn() } as any,
+      { getAppSecret: () => 'test-secret' } as any,
+      {} as any,
+    );
+    // EE not bundled: the seam returns null.
+    jest.spyOn(strategy as any, 'resolveApiKeyService').mockReturnValue(null);
     const req = makeReq();
 
     await expect(
       strategy.validate(req, apiKeyPayload() as any),
     ).rejects.toThrow(UnauthorizedException);
     expect(req.raw.actor).toBeUndefined();
-    expect(req.raw.authType).toBeUndefined();
   });
 });

apps/server/src/core/auth/strategies/jwt.strategy.ts

@@ -1,4 +1,4 @@
-import { Injectable, UnauthorizedException } from '@nestjs/common';
+import { Injectable, Logger, UnauthorizedException } from '@nestjs/common';
 import { PassportStrategy } from '@nestjs/passport';
 import { Strategy } from 'passport-jwt';
 import { EnvironmentService } from '../../../integrations/environment/environment.service';
@@ -9,18 +9,20 @@ import { UserSessionRepo } from '@docmost/db/repos/session/user-session.repo';
 import { SessionActivityService } from '../../session/session-activity.service';
 import { FastifyRequest } from 'fastify';
 import { extractBearerTokenFromHeader, isUserDisabled } from '../../../common/helpers';
+import { ModuleRef } from '@nestjs/core';
 import { resolveProvenance } from '../../../common/decorators/auth-provenance.decorator';
-import { ApiKeyService } from '../../api-key/api-key.service';
 
 @Injectable()
 export class JwtStrategy extends PassportStrategy(Strategy, 'jwt') {
+  private logger = new Logger('JwtStrategy');
+
   constructor(
     private userRepo: UserRepo,
     private workspaceRepo: WorkspaceRepo,
     private userSessionRepo: UserSessionRepo,
     private sessionActivityService: SessionActivityService,
     private readonly environmentService: EnvironmentService,
-    private readonly apiKeyService: ApiKeyService,
+    private moduleRef: ModuleRef,
   ) {
     super({
       jwtFromRequest: (req: FastifyRequest) => {
@@ -100,17 +102,12 @@ export class JwtStrategy extends PassportStrategy(Strategy, 'jwt') {
   }
 
   private async validateApiKey(req: any, payload: JwtApiKeyPayload) {
-    // The fork ships the core `ApiKeyService` (the EE `ee/api-key` module is
-    // absent). `validate` throws a bare UnauthorizedException on any definite
-    // deny (missing/revoked/expired row, disabled user, kill-switch off) and
-    // propagates infra errors (→ 5xx) rather than masking them as a 401.
-    const result = await this.apiKeyService.validate(payload);
+    const apiKeyService = this.resolveApiKeyService();
+    if (!apiKeyService) {
+      throw new UnauthorizedException('Enterprise API Key module missing');
+    }
 
-    // Stamp the principal kind + key id so the /api-keys management surface can
-    // enforce "a token cannot manage tokens". Done in this branch because it
-    // returns before the shared ACCESS-path stamping below.
-    req.raw.authType = 'api_key';
-    req.raw.apiKeyId = payload.apiKeyId;
+    const result = await apiKeyService.validateApiKey(payload);
 
     // Stamp the agent-edit provenance for the API-KEY path too (#486). Unlike the
     // access-token path above, it CANNOT be resolved before this point: the
@@ -122,10 +119,32 @@ export class JwtStrategy extends PassportStrategy(Strategy, 'jwt') {
     // SERVER-SIDE user (never a client field), so an 'agent' badge is unspoofable
     // — mirroring the access-token path. Passing `null` for the claim means the
     // actor is decided solely by user.isAgent.
-    const provenance = resolveProvenance(result.user, null);
+    const provenance = resolveProvenance((result as any)?.user, null);
     req.raw.actor = provenance.actor;
     req.raw.aiChatId = provenance.aiChatId;
 
     return result;
   }
+
+  /**
+   * Resolve the enterprise ApiKeyService, or `null` when the EE module is not
+   * bundled in this build (community build). Extracted as an overridable seam so
+   * the API-key provenance stamping can be unit-tested without the EE package
+   * present (docmost is OSS + a separate EE bundle; `require` of the EE path
+   * throws here). Any load/resolve error is treated as "module missing".
+   */
+  protected resolveApiKeyService(): {
+    validateApiKey: (payload: JwtApiKeyPayload) => Promise<unknown>;
+  } | null {
+    try {
+      // eslint-disable-next-line @typescript-eslint/no-require-imports
+      const ApiKeyModule = require('./../../../ee/api-key/api-key.service');
+      return this.moduleRef.get(ApiKeyModule.ApiKeyService, { strict: false });
+    } catch (err) {
+      this.logger.debug(
+        'API Key module requested but enterprise module not bundled in this build',
+      );
+      return null;
+    }
+  }
 }

apps/server/src/core/core.module.ts

@@ -6,7 +6,6 @@ import {
 } from '@nestjs/common';
 import { UserModule } from './user/user.module';
 import { AuthModule } from './auth/auth.module';
-import { ApiKeyModule } from './api-key/api-key.module';
 import { WorkspaceModule } from './workspace/workspace.module';
 import { PageModule } from './page/page.module';
 import { AttachmentModule } from './attachment/attachment.module';
@@ -30,7 +29,6 @@ import { ClsMiddleware } from 'nestjs-cls';
   imports: [
     UserModule,
     AuthModule,
-    ApiKeyModule,
     WorkspaceModule,
     PageModule,
     AttachmentModule,

apps/server/src/database/database.module.ts

@@ -32,13 +32,11 @@ import { TemplateRepo } from '@docmost/db/repos/template/template.repo';
 import { AiChatRepo } from '@docmost/db/repos/ai-chat/ai-chat.repo';
 import { AiChatMessageRepo } from '@docmost/db/repos/ai-chat/ai-chat-message.repo';
 import { AiChatRunRepo } from '@docmost/db/repos/ai-chat/ai-chat-run.repo';
-import { AiChatRunStepRepo } from '@docmost/db/repos/ai-chat/ai-chat-run-step.repo';
 import { AiChatPageSnapshotRepo } from '@docmost/db/repos/ai-chat/ai-chat-page-snapshot.repo';
 import { AiProviderCredentialsRepo } from '@docmost/db/repos/ai-chat/ai-provider-credentials.repo';
 import { AiMcpServerRepo } from '@docmost/db/repos/ai-chat/ai-mcp-server.repo';
 import { AiAgentRoleRepo } from '@docmost/db/repos/ai-agent-roles/ai-agent-roles.repo';
 import { PageEmbeddingRepo } from '@docmost/db/repos/ai-chat/page-embedding.repo';
-import { ApiKeyRepo } from '@docmost/db/repos/api-key/api-key.repo';
 import { PageListener } from '@docmost/db/listeners/page.listener';
 import { PostgresJSDialect } from 'kysely-postgres-js';
 import * as postgres from 'postgres';
@@ -126,13 +124,11 @@ import { firstSqlToken } from '../integrations/metrics/metrics.constants';
     AiChatRepo,
     AiChatMessageRepo,
     AiChatRunRepo,
-    AiChatRunStepRepo,
     AiChatPageSnapshotRepo,
     AiProviderCredentialsRepo,
     AiMcpServerRepo,
     AiAgentRoleRepo,
     PageEmbeddingRepo,
-    ApiKeyRepo,
     PageListener,
   ],
   exports: [
@@ -163,13 +159,11 @@ import { firstSqlToken } from '../integrations/metrics/metrics.constants';
     AiChatRepo,
     AiChatMessageRepo,
     AiChatRunRepo,
-    AiChatRunStepRepo,
     AiChatPageSnapshotRepo,
     AiProviderCredentialsRepo,
     AiMcpServerRepo,
     AiAgentRoleRepo,
     PageEmbeddingRepo,
-    ApiKeyRepo,
   ],
 })
 export class DatabaseModule implements OnApplicationBootstrap {

apps/server/src/database/migrations/20260708T120000-ai-chat-run-steps.ts

@@ -1,70 +0,0 @@
-import { type Kysely, sql } from 'kysely';
-
-/**
- * `ai_chat_run_steps` — append-only per-step persistence for an assistant turn
- * (#492 wave C). Each finished agent step's UI `parts` (its text part + a part
- * per tool call, WITH the tool output) is INSERTed as its own lightweight row the
- * moment the step ends, instead of REWRITING the whole assistant row's growing
- * `metadata.parts` jsonb on every `onStepFinish`.
- *
- * WHY a separate table + INSERT (not a jsonb `||` append on the message row): a
- * Postgres jsonb UPDATE rewrites the ENTIRE TOASTed row version under MVCC, so
- * re-persisting a growing `metadata.parts` on every step is O(n²) write volume
- * (a 50-step run with ~100 KB tool outputs wrote hundreds of MB of WAL / dead
- * tuples per turn, hammering autovacuum). `||` would only shave the network
- * payload — the WAL/TOAST rewrite harm remains. An INSERT into a per-step table
- * writes ONLY that step's bytes, so the per-turn write volume is O(Σ steps).
- *
- * The full `metadata.parts` on the message row is assembled ONCE at finalize (the
- * terminal completed/error/aborted write). Mid-run, a resuming client's seed is
- * reconstructed by concatenating these step rows in `step_index` order — which
- * reproduces exactly what the old per-step full-row rewrite persisted. Records
- * written the OLD way (full `metadata.parts` on the row, no step rows) still
- * reconstruct from the row unchanged; the two eras are distinguished by whether
- * the row already carries non-empty `metadata.parts` (see reconstructRunParts /
- * assembleStepParts in ai-chat.service.ts).
- *
- * ON DELETE CASCADE on `message_id`: the step rows are a derived projection of the
- * assistant message; they must vanish with it (or with its workspace).
- */
-export async function up(db: Kysely<any>): Promise<void> {
-  await db.schema
-    .createTable('ai_chat_run_steps')
-    .ifNotExists()
-    .addColumn('id', 'uuid', (col) =>
-      col.primaryKey().defaultTo(sql`gen_uuid_v7()`),
-    )
-    // The assistant message row this step belongs to (the #183 projection). The
-    // step rows are a derived, per-step slice of that message, so they cascade.
-    .addColumn('message_id', 'uuid', (col) =>
-      col.references('ai_chat_messages.id').onDelete('cascade').notNull(),
-    )
-    .addColumn('workspace_id', 'uuid', (col) =>
-      col.references('workspaces.id').onDelete('cascade').notNull(),
-    )
-    // 0-based index of the finished step within the turn. Ordering key for
-    // reconstruction; unique per message (idempotent step re-persist).
-    .addColumn('step_index', 'integer', (col) => col.notNull())
-    // The step's UI parts (text part + a `tool-*` part per call, WITH output).
-    // Concatenated in step order to rebuild the turn's `metadata.parts`.
-    .addColumn('parts', 'jsonb', (col) => col.notNull())
-    .addColumn('created_at', 'timestamptz', (col) =>
-      col.notNull().defaultTo(sql`now()`),
-    )
-    .execute();
-
-  // Idempotent per-step persist: a retried INSERT of the same (message, step)
-  // is a no-op (the service uses ON CONFLICT DO NOTHING). This also serves the
-  // reconstruction read (WHERE message_id ORDER BY step_index).
-  await db.schema
-    .createIndex('ai_chat_run_steps_message_step_uidx')
-    .ifNotExists()
-    .on('ai_chat_run_steps')
-    .columns(['message_id', 'step_index'])
-    .unique()
-    .execute();
-}
-
-export async function down(db: Kysely<any>): Promise<void> {
-  await db.schema.dropTable('ai_chat_run_steps').ifExists().execute();
-}

apps/server/src/database/repos/ai-chat/ai-chat-run-step.repo.ts

@@ -1,95 +0,0 @@
-import { Injectable } from '@nestjs/common';
-import { InjectKysely } from 'nestjs-kysely';
-import { KyselyDB, KyselyTransaction } from '../../types/kysely.types';
-import { dbOrTx } from '../../utils';
-import { AiChatRunStep } from '@docmost/db/types/entity.types';
-
-/**
- * Append-only per-step persistence for an assistant turn (#492). Each finished
- * agent step's UI `parts` (its text part + a `tool-*` part per call, WITH the
- * tool output) is INSERTed as its own lightweight row the moment the step ends —
- * instead of REWRITING the assistant row's growing `metadata.parts` jsonb on every
- * `onStepFinish` (a Postgres jsonb UPDATE rewrites the whole TOASTed row version
- * under MVCC, so that was O(n²) WAL/dead-tuple churn per turn).
- *
- * The full `metadata.parts` on the message row is assembled ONCE at finalize;
- * mid-run, a resuming client's seed is rebuilt from these rows in `stepIndex`
- * order (see `assembleStepParts` / the reconstruct seam in ai-chat.service.ts).
- * Every method is workspace-scoped as defense-in-depth.
- */
-@Injectable()
-export class AiChatRunStepRepo {
-  constructor(@InjectKysely() private readonly db: KyselyDB) {}
-
-  /**
-   * Append one finished step's parts. Idempotent: a retried persist of the SAME
-   * (message, stepIndex) is a no-op via ON CONFLICT DO NOTHING — the per-step
-   * writes are fired fire-and-forget + serialized, and a duplicate must never
-   * throw into the stream or double the parts. Returns whether a NEW row landed
-   * (false = the step was already persisted).
-   */
-  async insertStep(
-    messageId: string,
-    workspaceId: string,
-    stepIndex: number,
-    parts: unknown,
-    trx?: KyselyTransaction,
-  ): Promise<boolean> {
-    const db = dbOrTx(this.db, trx);
-    const inserted = await db
-      .insertInto('aiChatRunSteps')
-      .values({
-        messageId,
-        workspaceId,
-        stepIndex,
-        // jsonb column: cast through never (same pattern as the message repo).
-        parts: parts as never,
-      })
-      .onConflict((oc) => oc.columns(['messageId', 'stepIndex']).doNothing())
-      .returning('id')
-      .executeTakeFirst();
-    return inserted !== undefined;
-  }
-
-  /** All persisted steps for ONE assistant message, in step order. */
-  async findByMessage(
-    messageId: string,
-    workspaceId: string,
-  ): Promise<AiChatRunStep[]> {
-    return this.db
-      .selectFrom('aiChatRunSteps')
-      .selectAll('aiChatRunSteps')
-      .where('messageId', '=', messageId)
-      .where('workspaceId', '=', workspaceId)
-      .orderBy('stepIndex', 'asc')
-      .execute();
-  }
-
-  /**
-   * All persisted steps for a SET of assistant messages, grouped by messageId
-   * (each group in step order). One query for the batch — the hydration seam
-   * (getMessages / delta / export) calls this only for the rows that actually
-   * need reconstruction (an active new-style row whose `metadata.parts` is still
-   * empty), which is usually none, so this is skipped on the common path.
-   */
-  async findByMessageIds(
-    messageIds: string[],
-    workspaceId: string,
-  ): Promise<Map<string, AiChatRunStep[]>> {
-    const byMessage = new Map<string, AiChatRunStep[]>();
-    if (messageIds.length === 0) return byMessage;
-    const rows = await this.db
-      .selectFrom('aiChatRunSteps')
-      .selectAll('aiChatRunSteps')
-      .where('messageId', 'in', messageIds)
-      .where('workspaceId', '=', workspaceId)
-      .orderBy('stepIndex', 'asc')
-      .execute();
-    for (const row of rows) {
-      const list = byMessage.get(row.messageId);
-      if (list) list.push(row);
-      else byMessage.set(row.messageId, [row]);
-    }
-    return byMessage;
-  }
-}

apps/server/src/database/repos/api-key/api-key.repo.ts

@@ -1,110 +0,0 @@
-import { Injectable } from '@nestjs/common';
-import { InjectKysely } from 'nestjs-kysely';
-import { KyselyDB, KyselyTransaction } from '@docmost/db/types/kysely.types';
-import { DB, Users } from '@docmost/db/types/db';
-import { dbOrTx } from '@docmost/db/utils';
-import { ApiKey, InsertableApiKey } from '@docmost/db/types/entity.types';
-import { jsonObjectFrom } from 'kysely/helpers/postgres';
-import { ExpressionBuilder } from 'kysely';
-
-// Public-facing shape: the api_keys row plus a compact creator attribution used
-// by the admin list (the workspace-wide view attributes each key to its author).
-export type ApiKeyWithCreator = ApiKey & {
-  creator: Pick<Users, 'id' | 'name' | 'email' | 'avatarUrl'> | null;
-};
-
-@Injectable()
-export class ApiKeyRepo {
-  constructor(@InjectKysely() private readonly db: KyselyDB) {}
-
-  // Compact creator attribution embedded in the admin list rows. A nested
-  // subquery (jsonObjectFrom) keeps it one round-trip; the token material is
-  // never stored, so nothing sensitive is joined in.
-  private withCreator(eb: ExpressionBuilder<DB, 'apiKeys'>) {
-    return jsonObjectFrom(
-      eb
-        .selectFrom('users')
-        .select(['users.id', 'users.name', 'users.email', 'users.avatarUrl'])
-        .whereRef('users.id', '=', 'apiKeys.creatorId'),
-    ).as('creator');
-  }
-
-  async findById(
-    id: string,
-    workspaceId: string,
-    trx?: KyselyTransaction,
-  ): Promise<ApiKey | undefined> {
-    const db = dbOrTx(this.db, trx);
-    return db
-      .selectFrom('apiKeys')
-      .selectAll()
-      .where('id', '=', id)
-      .where('workspaceId', '=', workspaceId)
-      // Convention (soft-delete): a revoked key is invisible to reads.
-      .where('deletedAt', 'is', null)
-      .executeTakeFirst();
-  }
-
-  async findByCreator(
-    creatorId: string,
-    workspaceId: string,
-  ): Promise<ApiKeyWithCreator[]> {
-    return this.db
-      .selectFrom('apiKeys')
-      .selectAll('apiKeys')
-      .select((eb) => this.withCreator(eb))
-      .where('creatorId', '=', creatorId)
-      .where('workspaceId', '=', workspaceId)
-      .where('deletedAt', 'is', null)
-      .orderBy('createdAt', 'desc')
-      .execute() as unknown as Promise<ApiKeyWithCreator[]>;
-  }
-
-  async findAllInWorkspace(
-    workspaceId: string,
-  ): Promise<ApiKeyWithCreator[]> {
-    return this.db
-      .selectFrom('apiKeys')
-      .selectAll('apiKeys')
-      .select((eb) => this.withCreator(eb))
-      .where('workspaceId', '=', workspaceId)
-      .where('deletedAt', 'is', null)
-      .orderBy('createdAt', 'desc')
-      .execute() as unknown as Promise<ApiKeyWithCreator[]>;
-  }
-
-  async insert(
-    insertable: InsertableApiKey,
-    trx?: KyselyTransaction,
-  ): Promise<ApiKey> {
-    const db = dbOrTx(this.db, trx);
-    return db
-      .insertInto('apiKeys')
-      .values(insertable)
-      .returningAll()
-      .executeTakeFirst();
-  }
-
-  // Soft-delete = revoke. Terminal: the row stays for forensics but is invisible
-  // to every read above, so validate() denies it immediately (no grace window).
-  async softDelete(id: string, workspaceId: string): Promise<void> {
-    await this.db
-      .updateTable('apiKeys')
-      .set({ deletedAt: new Date(), updatedAt: new Date() })
-      .where('id', '=', id)
-      .where('workspaceId', '=', workspaceId)
-      .where('deletedAt', 'is', null)
-      .execute();
-  }
-
-  // Throttled best-effort forensics stamp. Not on the critical path: callers
-  // fire-and-forget and swallow errors, so it never fails a request.
-  async touchLastUsed(id: string): Promise<void> {
-    await this.db
-      .updateTable('apiKeys')
-      .set({ lastUsedAt: new Date() })
-      .where('id', '=', id)
-      .where('deletedAt', 'is', null)
-      .execute();
-  }
-}

apps/server/src/database/types/db.d.ts

@@ -692,22 +692,6 @@ export interface AiChatRuns {
   updatedAt: Generated<Timestamp>;
 }
 
-// Append-only per-step persistence for an assistant turn (#492). Mirrors
-// migration 20260708T120000-ai-chat-run-steps.ts. Each finished agent step's UI
-// `parts` are INSERTed as their own row (instead of rewriting the message row's
-// growing `metadata.parts` jsonb every step — an O(n²) WAL/TOAST churn). The full
-// `metadata.parts` is assembled once at finalize; mid-run a resuming client's seed
-// is rebuilt by concatenating these rows in `stepIndex` order. Cascades with the
-// assistant message row it projects.
-export interface AiChatRunSteps {
-  id: Generated<string>;
-  messageId: string;
-  workspaceId: string;
-  stepIndex: number;
-  parts: Json;
-  createdAt: Generated<Timestamp>;
-}
-
 // Per-(chat,page) snapshot of the open page's Markdown at the END of the agent's
 // previous turn (#274). Mirrors migration 20260702T120000-ai-chat-page-snapshot.ts.
 // The next turn diffs the CURRENT Markdown against `contentMd` to surface edits a
@@ -745,7 +729,6 @@ export interface DB {
   aiChats: AiChats;
   aiChatMessages: AiChatMessages;
   aiChatRuns: AiChatRuns;
-  aiChatRunSteps: AiChatRunSteps;
   aiChatPageSnapshots: AiChatPageSnapshots;
   apiKeys: ApiKeys;
   attachments: Attachments;

apps/server/src/database/types/entity.types.ts

@@ -4,7 +4,6 @@ import {
   AiChats,
   AiChatMessages,
   AiChatRuns,
-  AiChatRunSteps,
   AiChatPageSnapshots,
   Attachments,
   Comments,
@@ -65,12 +64,6 @@ export type InsertableAiChatMessage = Omit<Insertable<AiChatMessages>, 'tsv'>;
 export type AiChatRun = Selectable<AiChatRuns>;
 export type InsertableAiChatRun = Insertable<AiChatRuns>;
 
-// AI Chat Run Step (#492): append-only per-step parts persistence. Each finished
-// agent step's UI parts are stored as their own row; the full turn's parts are
-// assembled from these (in stepIndex order) for a mid-run resume seed.
-export type AiChatRunStep = Selectable<AiChatRunSteps>;
-export type InsertableAiChatRunStep = Insertable<AiChatRunSteps>;
-
 // AI Chat Page Snapshot (#274): per-(chat,page) Markdown snapshot taken at the
 // end of the agent's previous turn, diffed against the current page next turn to
 // detect human edits made between turns.

apps/server/src/integrations/environment/environment.service.ts

@@ -70,23 +70,6 @@ export class EnvironmentService {
     return this.configService.get<string>('JWT_TOKEN_EXPIRES_IN', '90d');
   }
 
-  // Kill-switch for the agent API-key feature. Default ON (unset -> enabled): a
-  // deploy that never sets the variable must NOT silently kill every agent. The
-  // parse is STRICT — the value is validated by environment.validation to be
-  // exactly 'true' or 'false' (or absent), so a typo like `=0`/`=off`/`=False`
-  // fails at boot rather than being read as "enabled" (which would leave an
-  // operator who yanked the switch during an incident with it still on). When
-  // OFF: validate() denies every api-key token and the issuance endpoints 404.
-  isApiKeysEnabled(): boolean {
-    return this.configService.get<string>('API_KEYS_ENABLED', 'true') !== 'false';
-  }
-
-  // Raw value (or undefined) for the boot log, so the state after each deploy is
-  // verifiable in container logs: `API keys: ENABLED/DISABLED (API_KEYS_ENABLED=...)`.
-  getApiKeysEnabledRaw(): string | undefined {
-    return this.configService.get<string>('API_KEYS_ENABLED');
-  }
-
   getCookieExpiresIn(): Date {
     const expiresInStr = this.getJwtTokenExpiresIn();
     let msUntilExpiry: number;

apps/server/src/integrations/environment/environment.validation.ts

@@ -103,15 +103,6 @@ export class EnvironmentVariables {
   @IsString()
   TYPESENSE_LOCALE: string;
 
-  // Agent API-key kill-switch. Optional (absent -> default ON). STRICT: only the
-  // literals 'true'/'false' are accepted, so `=0`/`=off`/`=False` fail at boot
-  // instead of being silently read as "enabled" — the switch must actually flip
-  // when an operator flips it. See EnvironmentService.isApiKeysEnabled.
-  @IsOptional()
-  @IsIn(['true', 'false'])
-  @IsString()
-  API_KEYS_ENABLED: string;
-
   @IsOptional()
   @ValidateIf((obj) => obj.AI_DRIVER)
   @IsIn(['openai', 'openai-compatible', 'gemini', 'ollama'])

apps/server/src/integrations/import/services/import-attachment.service.drawio-svg.spec.ts

@@ -1,131 +0,0 @@
-// p-limit and @sindresorhus/slugify are ESM-only and not in jest's transform
-// allowlist; both are irrelevant to createDrawioSvg (a pure fs + string method),
-// so they are mocked out to keep the module graph loadable under ts-jest.
-jest.mock('p-limit', () => ({
-  __esModule: true,
-  default: () => (fn: () => unknown) => fn(),
-}));
-jest.mock('@sindresorhus/slugify', () => ({
-  __esModule: true,
-  default: (input: string) => String(input),
-}));
-
-import { promises as fs } from 'fs';
-import * as os from 'os';
-import * as path from 'path';
-import { ImportAttachmentService } from './import-attachment.service';
-
-/**
- * Unit test for ImportAttachmentService.createDrawioSvg (issue #507).
- *
- * The Confluence import wraps a `.drawio` file into a `.drawio.svg` attachment.
- * The `content=` payload MUST be the mxfile XML entity-escaped (draw.io's native
- * form), NOT base64 — draw.io's editor decodes a base64 content= via Latin-1
- * atob, mangling every non-ASCII char into mojibake. createDrawioSvg touches no
- * injected dependency, so the service is built with placeholder deps.
- */
-describe('ImportAttachmentService.createDrawioSvg (#507)', () => {
-  const service = new ImportAttachmentService(
-    {} as any,
-    {} as any,
-    {} as any,
-  );
-  const call = (p: string): Promise<Buffer> =>
-    (service as any).createDrawioSvg(p);
-
-  let tmpDir: string;
-  beforeAll(async () => {
-    tmpDir = await fs.mkdtemp(path.join(os.tmpdir(), 'drawio-507-'));
-  });
-  afterAll(async () => {
-    await fs.rm(tmpDir, { recursive: true, force: true });
-  });
-
-  const writeDrawio = async (name: string, xml: string): Promise<string> => {
-    const p = path.join(tmpDir, name);
-    await fs.writeFile(p, xml, 'utf-8');
-    return p;
-  };
-
-  it('writes content= as entity-encoded XML, not base64', async () => {
-    const drawio =
-      '<mxfile host="Confluence"><diagram name="Схема — ёж">' +
-      '<mxGraphModel><root><mxCell id="0"/>' +
-      '<mxCell id="2" value="Старт-бит" vertex="1" parent="0"/>' +
-      '</root></mxGraphModel></diagram></mxfile>';
-    const p = await writeDrawio('cyrillic.drawio', drawio);
-
-    const svg = (await call(p)).toString('utf-8');
-    const content = /content="([^"]*)"/.exec(svg)?.[1];
-    expect(content).toBeDefined();
-
-    // Entity-encoded XML form, starting with &lt;mxfile — never a base64 blob.
-    expect(content).toMatch(/^&lt;mxfile/);
-    expect(content).toContain('&lt;');
-    // Non-ASCII survives as raw UTF-8, with no Latin-1 mojibake.
-    expect(content).toContain('Старт-бит');
-    expect(content).toContain('Схема — ёж');
-    expect(content).not.toContain('Ð');
-
-    // Decoding the attribute (un-escaping) yields the original drawio file.
-    const decoded = content!
-      .replace(/&lt;/g, '<')
-      .replace(/&gt;/g, '>')
-      .replace(/&quot;/g, '"')
-      .replace(/&amp;/g, '&');
-    expect(decoded).toBe(drawio);
-  });
-
-  it('encodes literal tab/newline/CR as numeric char-refs, not literal control chars (#507 F1)', async () => {
-    // A literal tab/newline/CR inside the mxfile XML would be collapsed to a
-    // single space by XML attribute-value normalization when the draw.io editor
-    // reads content=, silently flattening multi-line labels and tab-bearing
-    // values. They must be emitted as numeric char-refs instead.
-    const drawio =
-      '<mxfile><diagram name="p">' +
-      '<mxGraphModel><root><mxCell id="0"/>' +
-      '<mxCell id="2" value="col1\tcol2" style="html=1;\nshadow=0" vertex="1" parent="0"/>' +
-      '<mxCell id="3" value="Line1\nLine2\rLine3" vertex="1" parent="0"/>' +
-      '</root></mxGraphModel></diagram></mxfile>';
-    const p = await writeDrawio('ctrl.drawio', drawio);
-
-    const svg = (await call(p)).toString('utf-8');
-    const content = /content="([^"]*)"/.exec(svg)?.[1];
-    expect(content).toBeDefined();
-    // No literal control chars survive in the attribute value.
-    expect(content).not.toMatch(/[\t\n\r]/);
-    // They round-trip as numeric char-refs.
-    expect(content).toContain('&#x9;');
-    expect(content).toContain('&#xa;');
-    expect(content).toContain('&#xd;');
-    // Decoding (char-refs back to literal, entities back) recovers the file.
-    const decoded = content!
-      .replace(/&lt;/g, '<')
-      .replace(/&gt;/g, '>')
-      .replace(/&quot;/g, '"')
-      .replace(/&#39;/g, "'")
-      .replace(/&#x9;/gi, '\t')
-      .replace(/&#xa;/gi, '\n')
-      .replace(/&#xd;/gi, '\r')
-      .replace(/&amp;/g, '&');
-    expect(decoded).toBe(drawio);
-  });
-
-  it('escapes XML metacharacters in the drawio payload', async () => {
-    const drawio = '<mxfile><diagram name="a &amp; b">"q" &lt;x&gt;</diagram></mxfile>';
-    const p = await writeDrawio('meta.drawio', drawio);
-
-    const svg = (await call(p)).toString('utf-8');
-    const content = /content="([^"]*)"/.exec(svg)?.[1];
-    expect(content).toBeDefined();
-    // The attribute value must contain no bare `<`, `>` or `"` that would break
-    // out of the content="..." attribute or the SVG element.
-    expect(content).not.toMatch(/[<>"]/);
-    const decoded = content!
-      .replace(/&lt;/g, '<')
-      .replace(/&gt;/g, '>')
-      .replace(/&quot;/g, '"')
-      .replace(/&amp;/g, '&');
-    expect(decoded).toBe(drawio);
-  });
-});

apps/server/src/integrations/import/services/import-attachment.service.ts

@@ -8,7 +8,6 @@ import { createReadStream } from 'node:fs';
 import { promises as fs } from 'fs';
 import { Readable } from 'stream';
 import { getMimeType, sanitizeFileName } from '../../../common/helpers';
-import { htmlEscape } from '../../../common/helpers/html-escaper';
 import { v7 } from 'uuid';
 import { FileTask } from '@docmost/db/types/entity.types';
 import { getAttachmentFolderPath } from '../../../core/attachment/attachment.utils';
@@ -850,12 +849,7 @@ export class ImportAttachmentService {
   ): Promise<Buffer> {
     try {
       const drawioContent = await fs.readFile(drawioPath, 'utf-8');
-      // Write the mxfile XML XML-entity-escaped (draw.io's native content= form),
-      // NOT base64. draw.io's editor decodes a base64 content= via Latin-1 atob
-      // (no UTF-8 step), turning every non-ASCII char (Cyrillic, ё, —) into
-      // mojibake; the entity-encoded form is decoded by the DOM as UTF-8 and
-      // opens intact. Docmost's own decoder reads both forms.
-      const drawioEscaped = this.xmlEscapeContent(drawioContent);
+      const drawioBase64 = Buffer.from(drawioContent).toString('base64');
 
       let imageElement = '';
       // If we have a PNG, include it in the SVG
@@ -881,7 +875,7 @@ export class ImportAttachmentService {
       width="600"
       height="400"
       viewBox="0 0 600 400"
-      content="${drawioEscaped}">${imageElement}</svg>`;
+      content="${drawioBase64}">${imageElement}</svg>`;
 
       return Buffer.from(svgContent, 'utf-8');
     } catch (error) {
@@ -890,24 +884,6 @@ export class ImportAttachmentService {
     }
   }
 
-  /**
-   * Escape a string so it is safe as the value of a double-quoted XML attribute
-   * (the `content=` payload of a `.drawio.svg`). The shared `htmlEscape` covers
-   * `& < > " '` (a strict superset of what this attribute needs; the extra `'`
-   * escape is harmless in a `"`-delimited value). On top of that, the numeric
-   * char-refs for tab/newline/CR are required: a literal tab/newline/CR inside
-   * an attribute value is collapsed to a single space by XML attribute-value
-   * normalization on DOM read (both our decoder and the real draw.io editor),
-   * silently flattening multi-line labels and tab-bearing values. Char-refs
-   * survive that normalization (#507).
-   */
-  private xmlEscapeContent(s: string): string {
-    return htmlEscape(s)
-      .replace(/\t/g, '&#x9;')
-      .replace(/\n/g, '&#xa;')
-      .replace(/\r/g, '&#xd;');
-  }
-
   private async uploadWithRetry(opts: {
     abs: string;
     storageFilePath: string;

apps/server/src/integrations/mcp/mcp-auth.helpers.ts

@@ -304,102 +304,37 @@ export function clientIp(req: ClientIpRequest): string {
   return 'unknown';
 }
 
-// The decoded payload shared by the /mcp Bearer allowlist. Carries the `type`
-// discriminator and the API-key `apiKeyId`, on top of the access-token fields.
-export interface McpBearerPayload {
-  type?: JwtType;
+// Minimal structural shape of the TokenService.verifyJwt method we depend on,
+// so this module never imports the concrete TokenService (heavy graph).
+export interface AccessJwtVerifier {
+  verifyJwt: (
+    token: string,
+    type: JwtType,
+  ) => Promise<{
+    sub?: string;
+    email?: string;
+    workspaceId?: string;
+    sessionId?: string;
+  }>;
+}
+
+/**
+ * Bind a TokenService-like verifier into a one-arg `verifyJwt(token)` that
+ * ALWAYS enforces `JwtType.ACCESS`. This is the single place where the /mcp
+ * Bearer path pins the token type: a Bearer access token must be verified AS an
+ * access token (not refresh/exchange/collab/etc.), so the type literal is fixed
+ * here rather than at the call site. McpService.verifyMcpBearer delegates to
+ * this, keeping the `JwtType.ACCESS` choice testable without the heavy graph.
+ */
+export function bindAccessJwtVerifier(
+  tokenService: AccessJwtVerifier,
+): (token: string) => Promise<{
   sub?: string;
   email?: string;
   workspaceId?: string;
   sessionId?: string;
-  apiKeyId?: string;
-}
-
-// Minimal structural shape of the TokenService.verifyJwtOneOf method.
-export interface OneOfJwtVerifier {
-  verifyJwtOneOf: (
-    token: string,
-    allowed: JwtType[],
-  ) => Promise<McpBearerPayload>;
-}
-
-/**
- * Bind a TokenService-like verifier into a one-arg `verifyJwtOneOf(token)` that
- * pins the /mcp Bearer ALLOWLIST to exactly {ACCESS, API_KEY}. This is the single
- * place the /mcp Bearer path pins the token type: the /mcp Bearer slot
- * legitimately accepts either an ACCESS token (a
- * human's session token) OR an API_KEY token (an agent's key), but NOTHING else
- * (collab/exchange/attachment/etc. are rejected with the generic type error).
- * The allowlist is fixed here rather than at the call site, and the signature is
- * verified exactly once (see verifyMcpBearer).
- */
-export function bindMcpBearerVerifier(
-  tokenService: OneOfJwtVerifier,
-): (token: string) => Promise<McpBearerPayload> {
-  return (token: string) =>
-    tokenService.verifyJwtOneOf(token, [JwtType.ACCESS, JwtType.API_KEY]);
-}
-
-// Deps for the /mcp Bearer router. `verifyJwtOneOf` is the one-arg verifier bound
-// above (allowlist {ACCESS, API_KEY}); the ACCESS-specific revocation/disabled
-// deps mirror BearerVerifyDeps; `validateApiKey` is the SHARED api-key row-check.
-export interface McpBearerDeps
-  extends Omit<BearerVerifyDeps, 'verifyJwt'> {
-  verifyJwtOneOf: (token: string) => Promise<McpBearerPayload>;
-  // Row-check for an API_KEY principal — the SAME validator REST uses. Throws
-  // UnauthorizedException on a definite deny; PROPAGATES an infra error (→ 5xx),
-  // never masking it as a 401. Not a login attempt: the Basic limiter is not
-  // involved on this path.
-  validateApiKey: (payload: McpBearerPayload) => Promise<unknown>;
-}
-
-/**
- * Verify a /mcp Bearer token that may be an ACCESS token OR an API_KEY token, and
- * route by type. The signature is verified EXACTLY ONCE (verifyJwtOneOf); the
- * result is reused so the ACCESS branch does not re-verify.
- *
- *   - API_KEY -> bind to THIS instance's workspace FIRST (a token for another
- *     workspace is rejected), THEN run the shared `validateApiKey` row-check.
- *     No session/limiter involvement (an API key is not a login).
- *   - ACCESS  -> the unchanged `verifyBearerAccess` (session-active + not-disabled
- *     checks), fed a closure over the already-verified payload so "verify once"
- *     and "helper unchanged" coexist.
- *
- * Throws UnauthorizedException on any auth failure (uniform generic message — no
- * enumeration of why); propagates an infra error from `validateApiKey` as itself.
- */
-export async function verifyMcpBearer(
-  token: string,
-  deps: McpBearerDeps,
-): Promise<{ sub?: string; email?: string }> {
-  const generic = 'Invalid or expired token';
-  const payload = await deps.verifyJwtOneOf(token);
-
-  if (payload.type === JwtType.API_KEY) {
-    if (!payload.sub || !payload.workspaceId) {
-      throw new UnauthorizedException(generic);
-    }
-    // Instance-binding (mirrors verifyBearerAccess): reject an API_KEY token
-    // minted for a different workspace before touching the DB.
-    if (
-      deps.expectedWorkspaceId &&
-      payload.workspaceId !== deps.expectedWorkspaceId
-    ) {
-      throw new UnauthorizedException(generic);
-    }
-    // Shared row-check. A definite deny throws Unauthorized; an infra error
-    // propagates (→ 5xx), which the caller must NOT convert to a 401.
-    await deps.validateApiKey(payload);
-    return { sub: payload.sub };
-  }
-
-  // ACCESS: reuse verifyBearerAccess WITHOUT re-verifying the signature.
-  return verifyBearerAccess(token, {
-    verifyJwt: async () => payload,
-    expectedWorkspaceId: deps.expectedWorkspaceId,
-    findUser: deps.findUser,
-    findActiveSession: deps.findActiveSession,
-  });
+}> {
+  return (token: string) => tokenService.verifyJwt(token, JwtType.ACCESS);
 }
 
 // Minimal shapes for the Bearer revocation/disabled check. Kept structural so
@@ -793,24 +728,18 @@ export async function resolveMcpSessionConfig(
     };
   }
 
-  // --- 2) fallback A: Bearer JWT (user-supplied ACCESS or agent API_KEY) ---
+  // --- 2) fallback A: Bearer access-JWT (user-supplied token) ---
   const bearer = extractBearer(authHeader);
   if (bearer) {
     let payload: { sub?: string; email?: string };
     try {
       payload = await deps.verifyAccessJwt(bearer);
     } catch (err) {
-      // Anti-enumeration (Bearer leg): EVERY auth failure surfaces the SAME
-      // generic 401 — expired/revoked/wrong-type/unknown are indistinguishable
-      // to the caller (its reaction is identical either way). But an UNEXPECTED
-      // (infra) error is NOT an auth verdict: rethrow it AS ITSELF so the surface
-      // maps it to 5xx (mapAuthResultToResponse), never masking a DB/Redis
-      // outage as a bad token. verifyMcpBearer throws UnauthorizedException on a
-      // definite deny and lets an infra error from validateApiKey propagate.
-      if (err instanceof UnauthorizedException) {
-        throw new UnauthorizedException('Invalid or expired token');
-      }
-      throw err;
+      const message =
+        err instanceof Error && err.message
+          ? err.message
+          : 'Invalid or expired token';
+      throw new UnauthorizedException(message);
     }
     return {
       config: { apiUrl, getToken: async () => bearer },

apps/server/src/integrations/mcp/mcp-basic-login-gate.spec.ts

@@ -115,7 +115,6 @@ function makeService(opts: {
     undefined as never, // userRepo
     undefined as never, // userSessionRepo
     moduleRef as never, // moduleRef (read by the MFA branch)
-    undefined as never, // apiKeyService (unused by the login-gate path)
     undefined as never, // sandboxStore (unused by the login-gate path)
   );
   // Stop the constructor's unref'd sweep timer leaking across tests.

apps/server/src/integrations/mcp/mcp.module.ts

@@ -4,16 +4,13 @@ import { McpService } from './mcp.service';
 import { DatabaseModule } from '@docmost/db/database.module';
 import { AuthModule } from '../../core/auth/auth.module';
 import { TokenModule } from '../../core/auth/token.module';
-import { ApiKeyModule } from '../../core/api-key/api-key.module';
 
 // Community MCP feature: the server itself serves the Model Context Protocol
 // over HTTP at /mcp. DatabaseModule (global) provides WorkspaceRepo. AuthModule
 // supplies AuthService (per-user HTTP-Basic login validation) and TokenModule
-// supplies TokenService (Bearer JWT verification for the token path). ApiKeyModule
-// supplies ApiKeyService (the shared api-key row-check for the API_KEY Bearer
-// branch, so an agent authenticates with a key instead of the bcrypt Basic path).
+// supplies TokenService (Bearer access-JWT verification for the token fallback).
 @Module({
-  imports: [DatabaseModule, AuthModule, TokenModule, ApiKeyModule],
+  imports: [DatabaseModule, AuthModule, TokenModule],
   controllers: [McpController],
   providers: [McpService],
 })

apps/server/src/integrations/mcp/mcp.service.spec.ts

@@ -6,10 +6,9 @@ import {
   isCredentialsFailure,
   isInitializeRequestBody,
   verifyBearerAccess,
-  verifyMcpBearer,
-  bindMcpBearerVerifier,
   sharedTokenMatches,
   clientIp,
+  bindAccessJwtVerifier,
   extractBearer,
   decideBasicGate,
   mapAuthResultToResponse,
@@ -525,28 +524,13 @@ describe('resolveMcpSessionConfig', () => {
     expect(resolved.identity).toBe('bearer:user-9');
   });
 
-  it('Bearer invalid -> UNIFORM generic 401 (anti-enumeration, reason not leaked)', async () => {
-    // #501: the Bearer leg no longer surfaces the specific reason. Whether the
-    // token is expired, revoked, wrong-type or unknown, the caller sees ONE bare
-    // 'Invalid or expired token' — an agent's reaction is identical, and a leaked
-    // class would be an enumeration oracle.
+  it('Bearer invalid -> specific 401 from verifyAccessJwt', async () => {
     const verifyAccessJwt = jest
       .fn()
       .mockRejectedValue(new UnauthorizedException('jwt expired'));
     await expect(
       resolveMcpSessionConfig('Bearer expired', makeDeps({ verifyAccessJwt })),
-    ).rejects.toThrow('Invalid or expired token');
-  });
-
-  it('Bearer INFRA error -> propagates (NOT masked as 401)', async () => {
-    // A non-UnauthorizedException (e.g. a DB outage in the api-key row-check) is
-    // not an auth verdict: it must propagate so the surface maps it to 5xx.
-    const verifyAccessJwt = jest
-      .fn()
-      .mockRejectedValue(new Error('connection terminated'));
-    await expect(
-      resolveMcpSessionConfig('Bearer x', makeDeps({ verifyAccessJwt })),
-    ).rejects.toThrow('connection terminated');
+    ).rejects.toThrow('jwt expired');
   });
 
   it('no creds + env service account configured -> service-account config', async () => {
@@ -1017,104 +1001,48 @@ describe('clientIp (XFF-fallback precedence, item 5)', () => {
   });
 });
 
-describe('bindMcpBearerVerifier pins the {ACCESS, API_KEY} allowlist (#501)', () => {
-  it('calls verifyJwtOneOf with exactly [ACCESS, API_KEY]', async () => {
-    const verifyJwtOneOf = jest
+describe('bindAccessJwtVerifier enforces JwtType.ACCESS (item 3)', () => {
+  it('calls TokenService.verifyJwt with JwtType.ACCESS as the second argument', async () => {
+    // Mock TokenService: assert the type literal is pinned to ACCESS so swapping
+    // to REFRESH (or omitting the type) breaks this test.
+    const verifyJwt = jest
       .fn()
-      .mockResolvedValue({ type: JwtType.API_KEY, sub: 'u-1' });
-    await bindMcpBearerVerifier({ verifyJwtOneOf })('the.jwt');
-    expect(verifyJwtOneOf).toHaveBeenCalledWith('the.jwt', [
-      JwtType.ACCESS,
-      JwtType.API_KEY,
-    ]);
-    // Pin the concrete enum values too.
-    expect(verifyJwtOneOf.mock.calls[0][1]).toEqual(['access', 'api_key']);
-  });
-});
+      .mockResolvedValue({ sub: 'user-1', workspaceId: 'ws-1' });
+    const verify = bindAccessJwtVerifier({ verifyJwt });
 
-describe('verifyMcpBearer routes by token type (#501)', () => {
-  const accessDeps = (over: any = {}) => ({
-    verifyJwtOneOf: jest.fn(),
-    expectedWorkspaceId: 'ws-1',
-    findUser: jest.fn().mockResolvedValue({ deactivatedAt: null }),
-    findActiveSession: jest
-      .fn()
-      .mockResolvedValue({ userId: 'u-1', workspaceId: 'ws-1' }),
-    validateApiKey: jest.fn(),
-    ...over,
+    await verify('the.access.jwt');
+
+    expect(verifyJwt).toHaveBeenCalledTimes(1);
+    expect(verifyJwt).toHaveBeenCalledWith('the.access.jwt', JwtType.ACCESS);
+    // Pin the real enum value too, so renaming/repointing the enum member is caught.
+    expect(verifyJwt.mock.calls[0][1]).toBe('access');
   });
 
-  it('API_KEY -> row-checks via validateApiKey and does NOT touch session/limiter', async () => {
-    const deps = accessDeps({
-      verifyJwtOneOf: jest.fn().mockResolvedValue({
-        type: JwtType.API_KEY,
-        sub: 'svc-1',
-        workspaceId: 'ws-1',
-        apiKeyId: 'k-1',
-      }),
-      validateApiKey: jest.fn().mockResolvedValue({ user: { id: 'svc-1' } }),
-    });
-    const res = await verifyMcpBearer('tok', deps);
-    expect(deps.validateApiKey).toHaveBeenCalledTimes(1);
-    // No session lookup on the api-key path (not a login).
-    expect(deps.findActiveSession).not.toHaveBeenCalled();
-    expect(res).toEqual({ sub: 'svc-1' });
+  it('passes through the verified payload', async () => {
+    const payload = { sub: 'user-9', email: 'u@e.com', workspaceId: 'ws-1' };
+    const verifyJwt = jest.fn().mockResolvedValue(payload);
+    await expect(
+      bindAccessJwtVerifier({ verifyJwt })('t'),
+    ).resolves.toBe(payload);
   });
 
-  it('API_KEY for ANOTHER workspace -> rejected before the row-check', async () => {
-    const deps = accessDeps({
-      verifyJwtOneOf: jest.fn().mockResolvedValue({
-        type: JwtType.API_KEY,
-        sub: 'svc-1',
-        workspaceId: 'ws-OTHER',
-        apiKeyId: 'k-1',
-      }),
-      validateApiKey: jest.fn(),
-    });
-    await expect(verifyMcpBearer('tok', deps)).rejects.toBeInstanceOf(
-      UnauthorizedException,
-    );
-    expect(deps.validateApiKey).not.toHaveBeenCalled();
-  });
-
-  it('API_KEY infra error from validateApiKey PROPAGATES (not masked)', async () => {
-    const boom = new Error('db down');
-    const deps = accessDeps({
-      verifyJwtOneOf: jest.fn().mockResolvedValue({
-        type: JwtType.API_KEY,
-        sub: 'svc-1',
-        workspaceId: 'ws-1',
-        apiKeyId: 'k-1',
-      }),
-      validateApiKey: jest.fn().mockRejectedValue(boom),
-    });
-    await expect(verifyMcpBearer('tok', deps)).rejects.toBe(boom);
-  });
-
-  it('ACCESS -> runs the session/disabled checks and does NOT call validateApiKey', async () => {
-    const deps = accessDeps({
-      verifyJwtOneOf: jest.fn().mockResolvedValue({
-        type: JwtType.ACCESS,
-        sub: 'u-1',
-        workspaceId: 'ws-1',
-        sessionId: 'sess-1',
-        email: 'u@e.com',
-      }),
-    });
-    const res = await verifyMcpBearer('tok', deps);
-    expect(deps.findActiveSession).toHaveBeenCalledWith('sess-1');
-    expect(deps.validateApiKey).not.toHaveBeenCalled();
-    expect(res).toEqual({ sub: 'u-1', email: 'u@e.com' });
-  });
-
-  it('verifies the signature exactly ONCE (single verifyJwtOneOf, no re-verify)', async () => {
-    const verifyJwtOneOf = jest.fn().mockResolvedValue({
-      type: JwtType.ACCESS,
-      sub: 'u-1',
+  // The Bearer revocation/disabled checks (verifyBearerAccess) are covered above;
+  // this binds the ACCESS-type enforcement that verifyMcpBearer wires in.
+  it('feeds verifyBearerAccess so the whole Bearer chain enforces ACCESS', async () => {
+    const verifyJwt = jest.fn().mockResolvedValue({
+      sub: 'user-1',
       workspaceId: 'ws-1',
+      sessionId: 'sess-1',
     });
-    await verifyMcpBearer('tok', accessDeps({ verifyJwtOneOf }));
-    expect(verifyJwtOneOf).toHaveBeenCalledTimes(1);
+    const res = await verifyBearerAccess('t', {
+      verifyJwt: bindAccessJwtVerifier({ verifyJwt }),
+      findUser: jest.fn().mockResolvedValue({ deactivatedAt: null }),
+      findActiveSession: jest
+        .fn()
+        .mockResolvedValue({ userId: 'user-1', workspaceId: 'ws-1' }),
+    });
+    expect(verifyJwt).toHaveBeenCalledWith('t', JwtType.ACCESS);
+    expect(res).toEqual({ sub: 'user-1', email: undefined });
   });
 });
 
@@ -1270,7 +1198,6 @@ describe('McpService.onModuleDestroy — CollabSession teardown (#486)', () => {
       {} as any,
       {} as any,
       {} as any,
-      {} as any,
     );
   }
 

apps/server/src/integrations/mcp/mcp.service.ts

@@ -14,17 +14,16 @@ import { UserSessionRepo } from '@docmost/db/repos/session/user-session.repo';
 import { AuthService } from '../../core/auth/services/auth.service';
 import { TokenService } from '../../core/auth/services/token.service';
 import { validateSsoEnforcement } from '../../core/auth/auth.util';
-import { JwtApiKeyPayload } from '../../core/auth/dto/jwt-payload';
+import { JwtPayload } from '../../core/auth/dto/jwt-payload';
 import { Workspace } from '@docmost/db/types/entity.types';
-import { ApiKeyService } from '../../core/api-key/api-key.service';
 import {
   FailedLoginLimiter,
   resolveMcpSessionConfig,
-  verifyMcpBearer,
+  verifyBearerAccess,
   isInitializeRequestBody,
   sharedTokenMatches,
   clientIp,
-  bindMcpBearerVerifier,
+  bindAccessJwtVerifier,
   decideBasicGate,
   mapAuthResultToResponse,
   DocmostMcpConfig,
@@ -106,10 +105,6 @@ export class McpService implements OnModuleDestroy {
     private readonly userRepo: UserRepo,
     private readonly userSessionRepo: UserSessionRepo,
     private readonly moduleRef: ModuleRef,
-    // Shared api-key row-check for the /mcp API_KEY Bearer branch (same validator
-    // REST uses). Also lets an agent authenticate to /mcp with an api key instead
-    // of the bcrypt Basic path, so parallel reads stop starving the limiter.
-    private readonly apiKeyService: ApiKeyService,
     // Shared singleton in-RAM blob store backing the stash tool.
     private readonly sandboxStore: SandboxStore,
   ) {
@@ -199,41 +194,37 @@ export class McpService implements OnModuleDestroy {
     }
   }
 
-  // Bearer verification for the /mcp token path. The Bearer slot accepts EITHER
-  // an ACCESS token (a human session token) OR an API_KEY token (an agent's key)
-  // — the allowlist is pinned in bindMcpBearerVerifier. An ACCESS token is
-  // checked exactly as JwtStrategy does (signature/exp/type + session-active +
-  // not-disabled), so the MCP path is not weaker than the cookie/header path. An
-  // API_KEY token is HMAC-verified (microseconds) then row-checked via the shared
-  // ApiKeyService.validate — NOT a login attempt, so the Basic bcrypt path and
-  // its anti-brute-force limiter are never touched (the parallel-reads fix).
+  // Bearer access-JWT verification for the /mcp token fallback. verifyJwt only
+  // checks signature/exp/type, but a logged-out (revoked) or disabled user can
+  // still hold an unexpired access JWT. JwtStrategy additionally checks the
+  // session is active and the user is not disabled; we mirror those exact checks
+  // here so the MCP Bearer path is not weaker than the normal cookie/header path.
   private async verifyMcpBearer(
     token: string,
   ): Promise<{ sub?: string; email?: string }> {
-    // Resolve THIS instance's workspace so the router can bind the token's
-    // `workspaceId` claim to it (mirrors JwtStrategy). The community build is
-    // single-workspace (findFirst), so this is the default workspace and the
-    // check is a no-op here; it only rejects a foreign-workspace token in a
-    // multi-workspace deployment. Undefined (no workspace configured) means no
-    // check — the credentials path would already have failed with no workspace.
+    // Resolve THIS instance's workspace so verifyBearerAccess can bind the
+    // token's `workspaceId` claim to it (mirrors JwtStrategy). The community
+    // build is single-workspace (findFirst), so this is the default workspace
+    // and the check is a no-op here; it only rejects a foreign-workspace token
+    // in a multi-workspace deployment. Undefined (no workspace configured) means
+    // no check — the credentials path would already have failed with no
+    // workspace, and an undefined here keeps the helper a no-op rather than
+    // rejecting every token.
     const instanceWorkspace = await this.workspaceRepo.findFirst();
-    // The type-routing + revocation/disabled decision logic lives in the
-    // framework-free verifyMcpBearer helper (unit-testable without the heavy auth
-    // graph); this method only wires in the concrete TokenService + repos + the
-    // shared api-key validator.
-    return verifyMcpBearer(token, {
-      // The {ACCESS, API_KEY} allowlist enforcement lives in bindMcpBearerVerifier
-      // (a pure, testable seam) so the type set cannot silently drift.
-      verifyJwtOneOf: bindMcpBearerVerifier(this.tokenService),
+    // The revocation/disabled decision logic lives in the framework-free
+    // verifyBearerAccess helper (unit-testable without the heavy auth graph);
+    // this method only wires in the concrete TokenService + repos.
+    return verifyBearerAccess(token, {
+      // The JwtType.ACCESS enforcement lives in bindAccessJwtVerifier (a pure,
+      // testable seam) so the type literal cannot silently drift to REFRESH.
+      verifyJwt: bindAccessJwtVerifier(this.tokenService) as (
+        t: string,
+      ) => Promise<JwtPayload>,
       expectedWorkspaceId: instanceWorkspace?.id,
       findUser: (sub, workspaceId) =>
         this.userRepo.findById(sub, workspaceId),
       findActiveSession: (sessionId) =>
         this.userSessionRepo.findActiveById(sessionId),
-      // Shared with REST: a definite deny throws Unauthorized, an infra error
-      // propagates (→ 5xx). The /mcp bearer catch must preserve that distinction.
-      validateApiKey: (payload) =>
-        this.apiKeyService.validate(payload as JwtApiKeyPayload),
     });
   }
 

apps/server/test/integration/ai-chat-append-persist-service.int-spec.ts

@@ -1,412 +0,0 @@
-import * as http from 'node:http';
-import { Kysely } from 'kysely';
-import { tool } from 'ai';
-import { z } from 'zod';
-import { MockLanguageModelV3, convertArrayToReadableStream } from 'ai/test';
-import { AiChatRepo } from '@docmost/db/repos/ai-chat/ai-chat.repo';
-import { AiChatMessageRepo } from '@docmost/db/repos/ai-chat/ai-chat-message.repo';
-import { AiChatRunStepRepo } from '@docmost/db/repos/ai-chat/ai-chat-run-step.repo';
-import {
-  AiChatService,
-  assembleStepParts,
-  assistantParts,
-  rowHasInlineParts,
-  stepMarkerMetadata,
-} from 'src/core/ai-chat/ai-chat.service';
-import {
-  getTestDb,
-  destroyTestDb,
-  createWorkspace,
-  createUser,
-  createChat,
-  createMessage,
-} from './db';
-
-/**
- * #492 append-persist — the REAL onStep WRITE path (F2) and the model-REPLAY
- * hydration path (F1), driven through `AiChatService.stream` against a LIVE
- * Postgres with a REAL `AiChatRunStepRepo` INJECTED. The existing append-persist
- * int-specs hand-roll the insert+marker cycle via the repos directly and build
- * the service with `aiChatRunStepRepo: undefined` (only the legacy-fallback branch
- * is covered), so an off-by-one on `stepsPersisted-1`, a wrong `capturedSteps`
- * slice, or a broken marker payload would pass all of them. These tests exercise
- * the actual `updateStreaming` append-persist branch end to end.
- *
- * The seam is the injected `model` (a seeded `MockLanguageModelV3` from `ai/test`)
- * plus a REAL Node `ServerResponse` as the hijacked socket — mirrors
- * ai-chat-stream.int-spec.ts.
- */
-
-const sleep = (ms: number) => new Promise((r) => setTimeout(r, ms));
-
-async function waitFor(
-  cond: () => Promise<boolean> | boolean,
-  { timeoutMs = 15_000, stepMs = 25 } = {},
-): Promise<void> {
-  const start = Date.now();
-  while (Date.now() - start < timeoutMs) {
-    if (await cond()) return;
-    await sleep(stepMs);
-  }
-  throw new Error('waitFor: condition not met within timeout');
-}
-
-// A real Node ServerResponse wired to a live socket (identical helper to the
-// stream int-spec) so the SDK's pipe/heartbeat writes behave as in prod.
-function makeRealResponse(): Promise<{
-  res: http.ServerResponse;
-  cleanup: () => Promise<void>;
-}> {
-  return new Promise((resolve) => {
-    const server = http.createServer((_req, res) => {
-      resolve({
-        res,
-        cleanup: () =>
-          new Promise<void>((done) => {
-            try {
-              if (!res.writableEnded) res.end();
-            } catch {
-              /* socket already gone */
-            }
-            server.close(() => done());
-          }),
-      });
-    });
-    server.listen(0, () => {
-      const port = (server.address() as any).port;
-      const creq = http.request({ port, method: 'GET' }, (cres) => {
-        cres.resume();
-      });
-      creq.on('error', () => undefined);
-      creq.end();
-    });
-  });
-}
-
-// Stream parts for a normal, successful single-step turn.
-function successStream() {
-  return convertArrayToReadableStream([
-    { type: 'stream-start', warnings: [] },
-    { type: 'text-start', id: 't1' },
-    { type: 'text-delta', id: 't1', delta: 'Hello' },
-    { type: 'text-delta', id: 't1', delta: ' there' },
-    { type: 'text-end', id: 't1' },
-    {
-      type: 'finish',
-      finishReason: 'stop',
-      usage: { inputTokens: 10, outputTokens: 5, totalTokens: 15 },
-    },
-  ] as any);
-}
-
-// A THREE-step turn: steps 0 and 1 each emit text + an `echo` tool call (the SDK
-// runs the tool and continues); step 2 answers and stops. Three steps is
-// deliberate: the LAST finished step's append-persist write races the terminal
-// finalize (which writes the full inline parts anyway, so a lost last-step row is
-// by design), but the NON-final steps 0 and 1 always drain to the steps table
-// before finalize — so those are what the test asserts on deterministically.
-function threeStepModel(): MockLanguageModelV3 {
-  let step = 0;
-  const toolStep = (i: number) => ({
-    stream: convertArrayToReadableStream([
-      { type: 'stream-start', warnings: [] },
-      { type: 'text-start', id: `s${i}` },
-      { type: 'text-delta', id: `s${i}`, delta: `step ${i} ` },
-      { type: 'text-end', id: `s${i}` },
-      {
-        type: 'tool-call',
-        toolCallId: `c${i}`,
-        toolName: 'echo',
-        input: JSON.stringify({ msg: `m${i}` }),
-      },
-      {
-        type: 'finish',
-        finishReason: 'tool-calls',
-        usage: { inputTokens: 5, outputTokens: 3, totalTokens: 8 },
-      },
-    ] as any),
-  });
-  return new MockLanguageModelV3({
-    doStream: async () => {
-      const n = step++;
-      // Realistic inter-step latency. A real model spends seconds per step, so the
-      // fire-and-forget per-step write chain drains to the steps table BETWEEN
-      // steps; the mock otherwise collapses all steps into microseconds and the
-      // terminal finalize wins the race before any but the first step persists.
-      if (n > 0) await sleep(200);
-      if (n < 2) return toolStep(n);
-      return {
-        stream: convertArrayToReadableStream([
-          { type: 'stream-start', warnings: [] },
-          { type: 'text-start', id: 's2' },
-          { type: 'text-delta', id: 's2', delta: 'final answer' },
-          { type: 'text-end', id: 's2' },
-          {
-            type: 'finish',
-            finishReason: 'stop',
-            usage: { inputTokens: 6, outputTokens: 4, totalTokens: 10 },
-          },
-        ] as any),
-      };
-    },
-  } as any);
-}
-
-describe('#492 append-persist service paths [integration]', () => {
-  let db: Kysely<any>;
-  let aiChatRepo: AiChatRepo;
-  let msgRepo: AiChatMessageRepo;
-  let stepRepo: AiChatRunStepRepo;
-  let workspaceId: string;
-  let userId: string;
-
-  let closeCalls: number;
-  const mcpClients = {
-    toolsFor: async () => ({
-      tools: {},
-      clients: [
-        {
-          close: async () => {
-            closeCalls += 1;
-          },
-        },
-      ],
-      outcomes: [],
-      instructions: [],
-    }),
-  };
-
-  // Build the service WITH a REAL AiChatRunStepRepo injected (the property under
-  // test) — unlike the legacy-fallback harness that passes it as undefined.
-  const echoTool = tool({
-    description: 'echo the message back',
-    inputSchema: z.object({ msg: z.string() }),
-    execute: async ({ msg }) => ({ echoed: msg }),
-  });
-
-  function buildService(): AiChatService {
-    return new AiChatService(
-      { getChatModel: async () => null } as any,
-      aiChatRepo,
-      msgRepo,
-      {} as any, // aiChatPageSnapshotRepo
-      { resolve: async () => null } as any, // aiSettings
-      { forUser: async () => ({ echo: echoTool }) } as any, // tools
-      mcpClients as any,
-      {} as any, // aiAgentRoleRepo
-      {} as any, // pageRepo
-      {} as any, // pageAccess
-      {
-        isAiChatDeferredToolsEnabled: () => false,
-        isAiChatFinalStepLockdownEnabled: () => false,
-      } as any, // environment (deferred OFF -> all tools active every step)
-      undefined, // streamRegistry
-      undefined, // aiChatRunService
-      stepRepo, // #492 aiChatRunStepRepo — the append-persist backend
-    );
-  }
-
-  function userUiMessage(text: string) {
-    return {
-      id: `u-${Math.random()}`,
-      role: 'user',
-      parts: [{ type: 'text', text }],
-    };
-  }
-
-  async function runStream(opts: {
-    model: MockLanguageModelV3;
-    chatId: string;
-    body: any;
-  }): Promise<void> {
-    closeCalls = 0;
-    const service = buildService();
-    const { res, cleanup } = await makeRealResponse();
-    try {
-      await service.stream({
-        user: { id: userId, workspaceId } as any,
-        workspace: { id: workspaceId, name: 'WS' } as any,
-        sessionId: 'sess-1',
-        body: opts.body,
-        res: { raw: res } as any,
-        signal: new AbortController().signal,
-        model: opts.model as any,
-        role: null,
-      } as any);
-      await waitFor(async () => {
-        const rows = await msgRepo.findAllByChat(opts.chatId, workspaceId);
-        return rows.some(
-          (r) =>
-            r.role === 'assistant' &&
-            ['completed', 'error', 'aborted'].includes(r.status as string),
-        );
-      });
-      await waitFor(() => closeCalls > 0, { timeoutMs: 5_000 });
-    } finally {
-      await cleanup();
-    }
-  }
-
-  beforeAll(async () => {
-    db = getTestDb();
-    aiChatRepo = new AiChatRepo(db as any);
-    msgRepo = new AiChatMessageRepo(db as any);
-    stepRepo = new AiChatRunStepRepo(db as any);
-    workspaceId = (await createWorkspace(db)).id;
-    userId = (await createUser(db, workspaceId)).id;
-  });
-
-  afterAll(async () => {
-    await destroyTestDb();
-  });
-
-  // --- F2: the real onStep append-persist WRITE branch -----------------------
-  it('drives steps through the real onStep path: per-step rows + marker match a single-row flush', async () => {
-    const chatId = (await createChat(db, { workspaceId, creatorId: userId })).id;
-    const model = threeStepModel();
-
-    // Capture the mid-run step-marker UPDATEs the append-persist branch writes on
-    // the assistant row (a { parts: [], toolTraceVersion, stepsPersisted } patch).
-    const updateSpy = jest.spyOn(msgRepo, 'update');
-    try {
-      await runStream({
-        model,
-        chatId,
-        body: { chatId, messages: [userUiMessage('call the tool then answer')] },
-      });
-
-      const rows = await msgRepo.findAllByChat(chatId, workspaceId);
-      const assistant = rows.find((r) => r.role === 'assistant')!;
-      expect(assistant).toBeDefined();
-      expect(assistant.status).toBe('completed');
-      // The turn finalizes with the FULL inline parts assembled by a single-row
-      // flush (assistantParts over every step) — the baseline the per-step slices
-      // must reproduce.
-      expect(rowHasInlineParts(assistant)).toBe(true);
-      const finalParts = (assistant.metadata as { parts: any[] }).parts;
-
-      // The two NON-final finished steps each landed their own row, in stepIndex
-      // order. (The fire-and-forget write chain drains before the next step, so
-      // poll until both are on disk; the LAST step's write may lose the finalize
-      // race, which is by design — its parts are already in `finalParts`.)
-      await waitFor(async () => {
-        const s = await stepRepo.findByMessage(assistant.id, workspaceId);
-        return s.length >= 2;
-      });
-      const steps = await stepRepo.findByMessage(assistant.id, workspaceId);
-      expect(steps[0].stepIndex).toBe(0);
-      expect(steps[1].stepIndex).toBe(1);
-
-      // Each per-step row carries a NON-trivial slice: this step's text part + its
-      // paired tool part (guards a mutation that persists empty/whole-turn parts).
-      const s0 = steps[0].parts as any[];
-      expect(s0).toContainEqual({ type: 'text', text: 'step 0 ' });
-      expect(s0.some((p) => p.type === 'tool-echo')).toBe(true);
-
-      // The per-step slices are EXACTLY the corresponding prefix of the single-row
-      // flush: assembleStepParts([step0, step1]) === finalParts[0 .. len0+len1].
-      // This is what an off-by-one on `stepsPersisted-1` (a wrong `capturedSteps`
-      // slice) or a shifted stepIndex breaks — the prefix no longer aligns.
-      const prefixLen =
-        (steps[0].parts as any[]).length + (steps[1].parts as any[]).length;
-      expect(assembleStepParts([steps[0], steps[1]] as any)).toEqual(
-        finalParts.slice(0, prefixLen),
-      );
-
-      // The mid-run step markers advanced 1 -> 2 -> ... (the resume frontier), each
-      // a shape-stable empty-parts marker equal to a single-row flush's marker.
-      const markerCounts = updateSpy.mock.calls
-        .map((c) => (c[2] as any)?.metadata)
-        .filter(
-          (m) =>
-            m &&
-            Array.isArray(m.parts) &&
-            m.parts.length === 0 &&
-            typeof m.stepsPersisted === 'number',
-        )
-        .map((m) => m.stepsPersisted);
-      // Monotonic from 1, covering at least the two non-final steps.
-      expect(markerCounts.slice(0, 2)).toEqual([1, 2]);
-      expect(
-        updateSpy.mock.calls
-          .map((c) => (c[2] as any)?.metadata)
-          .find((m) => m && m.stepsPersisted === 2),
-      ).toEqual(stepMarkerMetadata(2));
-    } finally {
-      updateSpy.mockRestore();
-    }
-  }, 60_000);
-
-  // --- F1: model-REPLAY hydrates a hard-crashed mid-run turn from the steps table
-  it('replays a hard-crashed mid-run turn WITH its partial steps hydrated from the steps table', async () => {
-    const chatId = (await createChat(db, { workspaceId, creatorId: userId })).id;
-
-    // Prior turn: a genuine user question...
-    await createMessage(db, {
-      workspaceId,
-      chatId,
-      userId,
-      role: 'user',
-      content: 'What is in the design doc?',
-      createdAt: new Date(Date.now() - 3000),
-    });
-    // ...and an assistant row that a HARD crash (SIGKILL/OOM) left mid-run: only a
-    // step marker on the row (metadata.parts:[] , content:''), NO terminal
-    // callback ever fired, so its real parts live ONLY in ai_chat_run_steps.
-    const crashed = await createMessage(db, {
-      workspaceId,
-      chatId,
-      role: 'assistant',
-      content: '',
-      status: 'aborted',
-      metadata: stepMarkerMetadata(1),
-      createdAt: new Date(Date.now() - 2000),
-    });
-    // The durable partial step: some reasoning text + a completed getPage tool
-    // call (input + output), exactly what #183 step-granular durability preserves.
-    await stepRepo.insertStep(
-      crashed.id,
-      workspaceId,
-      0,
-      assistantParts(
-        [
-          {
-            text: 'HYDRATED_PARTIAL_STEP the doc says',
-            toolCalls: [
-              { toolCallId: 'g1', toolName: 'getPage', input: { id: 'p1' } },
-            ],
-            toolResults: [
-              {
-                toolCallId: 'g1',
-                toolName: 'getPage',
-                output: { id: 'p1', body: 'PARTIAL_TOOL_OUTPUT budget section' },
-              },
-            ],
-          } as any,
-        ],
-        '',
-      ),
-    );
-
-    // The NEXT turn: the model just answers. The service must REPLAY the crashed
-    // assistant turn with its partial parts hydrated from the steps table.
-    const model = new MockLanguageModelV3({
-      doStream: async () => ({ stream: successStream() }),
-    } as any);
-    await runStream({
-      model,
-      chatId,
-      body: { chatId, messages: [userUiMessage('Continue please')] },
-    });
-
-    expect(model.doStreamCalls.length).toBeGreaterThan(0);
-    const prompt = JSON.stringify(model.doStreamCalls[0].prompt);
-    // The partial step's TEXT reached the model context (it would be an empty text
-    // part without hydration — rowToUiMessage falls back to `content:''`).
-    expect(prompt).toContain('HYDRATED_PARTIAL_STEP');
-    // The partial TOOL RESULT survived too (durable in the steps table, replayed).
-    expect(prompt).toContain('PARTIAL_TOOL_OUTPUT');
-    // The genuine prior user turn is present as well (sanity: real history replay).
-    expect(prompt).toContain('What is in the design doc?');
-  }, 60_000);
-});

apps/server/test/integration/ai-chat-append-persist-wal.int-spec.ts

@@ -1,173 +0,0 @@
-import { randomBytes } from 'crypto';
-import { Kysely, sql } from 'kysely';
-import { AiChatRunStepRepo } from '@docmost/db/repos/ai-chat/ai-chat-run-step.repo';
-import { AiChatMessageRepo } from '@docmost/db/repos/ai-chat/ai-chat-message.repo';
-import {
-  assistantParts,
-  flushAssistant,
-  stepMarkerMetadata,
-} from '../../src/core/ai-chat/ai-chat.service';
-import {
-  getTestDb,
-  destroyTestDb,
-  createWorkspace,
-  createUser,
-  createChat,
-} from './db';
-
-/**
- * #492 append-persist — WRITE-VOLUME regression on a LIVE Postgres, measured via
- * the `pg_current_wal_lsn()` delta around a realistic multi-step run driven through
- * the REAL repos (not a mock — a mock cannot observe MVCC/TOAST rewrite volume, the
- * whole point). Proves the core claim:
- *
- *   NEW (per-step INSERT into ai_chat_run_steps + a CHEAP step-marker UPDATE on the
- *   message row) writes O(Σ steps) of WAL — each step writes only its own bytes.
- *
- *   OLD (the pre-#492 full-row rewrite: re-persist the GROWING metadata.parts on
- *   every onStepFinish) writes O(n²) — step k rewrites the whole TOASTed jsonb of
- *   all k prior outputs.
- *
- * The OLD path here IS the reverted behavior, so this doubles as the mutation
- * check: swapping the new path back to `flushAssistant` full-row UPDATEs reddens
- * the assertion (OLD is many times larger).
- */
-type Step = {
-  text: string;
-  toolCalls: Array<{ toolCallId: string; toolName: string; input: unknown }>;
-  toolResults: Array<{ toolCallId: string; toolName: string; output: unknown }>;
-};
-
-// ~100 KB INCOMPRESSIBLE output per step (a page read). Random base64 so TOAST
-// cannot compress it away and hide the real write volume.
-function makeStep(i: number, outputBytes = 100_000): Step {
-  const body = randomBytes(Math.ceil(outputBytes * 0.75)).toString('base64');
-  return {
-    text: `step ${i} reasoning`,
-    toolCalls: [
-      { toolCallId: `c${i}`, toolName: 'getPage', input: { id: `p${i}` } },
-    ],
-    toolResults: [
-      {
-        toolCallId: `c${i}`,
-        toolName: 'getPage',
-        output: { id: `p${i}`, title: `Page ${i}`, body },
-      },
-    ],
-  };
-}
-
-async function walDelta(
-  db: Kysely<any>,
-  fn: () => Promise<void>,
-): Promise<number> {
-  const before = (
-    await sql<{ l: string }>`select pg_current_wal_lsn() as l`.execute(db)
-  ).rows[0].l;
-  await fn();
-  // NOTE: no pg_switch_wal() — a segment switch pads the LSN to the next 16 MB
-  // boundary and would swamp the delta. The raw LSN advances by the WAL bytes.
-  const after = (
-    await sql<{ l: string }>`select pg_current_wal_lsn() as l`.execute(db)
-  ).rows[0].l;
-  return Number(
-    (
-      await sql<{
-        d: string;
-      }>`select pg_wal_lsn_diff(${after}::pg_lsn, ${before}::pg_lsn) as d`.execute(
-        db,
-      )
-    ).rows[0].d,
-  );
-}
-
-describe('#492 append-persist write volume (pg_current_wal_lsn delta) [integration]', () => {
-  let db: Kysely<any>;
-  let stepRepo: AiChatRunStepRepo;
-  let msgRepo: AiChatMessageRepo;
-  let workspaceId: string;
-  let userId: string;
-  let chatId: string;
-
-  beforeAll(async () => {
-    db = getTestDb();
-    stepRepo = new AiChatRunStepRepo(db as any);
-    msgRepo = new AiChatMessageRepo(db as any);
-    workspaceId = (await createWorkspace(db)).id;
-    userId = (await createUser(db, workspaceId)).id;
-    chatId = (await createChat(db, { workspaceId, creatorId: userId })).id;
-  });
-
-  afterAll(async () => {
-    await destroyTestDb();
-  });
-
-  const seedRow = () =>
-    msgRepo.insert({
-      chatId,
-      workspaceId,
-      userId,
-      role: 'assistant',
-      content: '',
-      status: 'streaming',
-      metadata: stepMarkerMetadata(0) as never,
-    });
-
-  const STEPS = 40;
-
-  it('NEW per-step INSERT is O(Σ steps); OLD full-row rewrite is O(n²)', async () => {
-    const steps: Step[] = [];
-    for (let i = 0; i < STEPS; i++) steps.push(makeStep(i));
-
-    // NEW: per-step INSERT of THIS step's parts + a cheap marker UPDATE.
-    const newRow = await seedRow();
-    const newWal = await walDelta(db, async () => {
-      for (let i = 0; i < STEPS; i++) {
-        await stepRepo.insertStep(
-          newRow.id,
-          workspaceId,
-          i,
-          assistantParts([steps[i]], ''),
-        );
-        await msgRepo.update(
-          newRow.id,
-          workspaceId,
-          { metadata: stepMarkerMetadata(i + 1) },
-          { onlyIfStreaming: true },
-        );
-      }
-    });
-
-    // OLD (the pre-#492 revert): re-persist the GROWING metadata.parts on the
-    // message row on every step.
-    const oldRow = await seedRow();
-    const oldWal = await walDelta(db, async () => {
-      const acc: Step[] = [];
-      for (let i = 0; i < STEPS; i++) {
-        acc.push(steps[i]);
-        await msgRepo.update(
-          oldRow.id,
-          workspaceId,
-          flushAssistant(acc as never, '', 'streaming'),
-          { onlyIfStreaming: true },
-        );
-      }
-    });
-
-    // eslint-disable-next-line no-console
-    console.log(
-      `[#492 WAL] ${STEPS} steps ×100KB: new=${(newWal / 1e6).toFixed(1)}MB ` +
-        `old=${(oldWal / 1e6).toFixed(1)}MB (${(oldWal / newWal).toFixed(
-          1,
-        )}x smaller)`,
-    );
-
-    // O(Σ steps): ~STEPS × (100KB output + marker) of WAL. 40 × ~100KB parts plus
-    // 40 tiny markers is a few tens of MB at most — bounded, linear in step count.
-    expect(newWal).toBeLessThan(30_000_000);
-    // O(n²): step k rewrites ~k × 100KB. Σ over 40 steps ≈ 80+ MB — far larger.
-    expect(oldWal).toBeGreaterThan(30_000_000);
-    // The load-bearing claim: the new path writes a small FRACTION of the old.
-    expect(newWal).toBeLessThan(oldWal * 0.35);
-  }, 120_000);
-});

apps/server/test/integration/ai-chat-hydrate-controller.int-spec.ts

@@ -1,169 +0,0 @@
-import { Kysely } from 'kysely';
-import { AiChatController } from 'src/core/ai-chat/ai-chat.controller';
-import {
-  assembleStepParts,
-  assistantParts,
-  stepMarkerMetadata,
-} from 'src/core/ai-chat/ai-chat.service';
-import { AiChatRepo } from '@docmost/db/repos/ai-chat/ai-chat.repo';
-import { AiChatMessageRepo } from '@docmost/db/repos/ai-chat/ai-chat-message.repo';
-import { AiChatRunStepRepo } from '@docmost/db/repos/ai-chat/ai-chat-run-step.repo';
-import type { User, Workspace } from '@docmost/db/types/entity.types';
-import {
-  getTestDb,
-  destroyTestDb,
-  createWorkspace,
-  createUser,
-  createChat,
-  createMessage,
-} from './db';
-
-/**
- * #492 controller hydration (crash-before-finalize RESUME) on a LIVE Postgres.
- * `AiChatController.withReconstructedParts` is wired into getMessages/delta/export/
- * run, but `aiChatRunStepRepo` is OPTIONAL and every controller unit spec passes it
- * as `undefined`, so the hydration branch early-returns and NEVER executes in those
- * tests. This drives the real read path — a mid-run streaming row (marker only,
- * empty inline parts) PLUS its `ai_chat_run_steps` rows — through getMessages WITH
- * the repo present, exercising the `role==='assistant' && !rowHasInlineParts`
- * needy predicate, the workspace-scoped batch step fetch, and the endpoint binding.
- */
-describe('#492 controller hydration read path [integration]', () => {
-  let db: Kysely<any>;
-  let aiChatRepo: AiChatRepo;
-  let msgRepo: AiChatMessageRepo;
-  let stepRepo: AiChatRunStepRepo;
-  let workspaceId: string;
-  let otherWorkspaceId: string;
-  let userId: string;
-
-  // Build the controller WITH a real AiChatRunStepRepo injected (position 9), the
-  // seam the unit specs leave undefined. Only the read-path deps are real.
-  function buildController(): AiChatController {
-    return new AiChatController(
-      {} as any, // aiChatService
-      {} as any, // aiChatRunService
-      aiChatRepo,
-      msgRepo,
-      {} as any, // aiTranscription
-      {} as any, // pageRepo
-      undefined, // streamRegistry
-      undefined, // environment
-      stepRepo, // #492 aiChatRunStepRepo
-    );
-  }
-
-  beforeAll(async () => {
-    db = getTestDb();
-    aiChatRepo = new AiChatRepo(db as any);
-    msgRepo = new AiChatMessageRepo(db as any);
-    stepRepo = new AiChatRunStepRepo(db as any);
-    workspaceId = (await createWorkspace(db)).id;
-    otherWorkspaceId = (await createWorkspace(db)).id;
-    userId = (await createUser(db, workspaceId)).id;
-  });
-
-  afterAll(async () => {
-    await destroyTestDb();
-  });
-
-  it('getMessages reconstructs a mid-run row from the steps table (finished rows untouched)', async () => {
-    const chatId = (
-      await createChat(db, { workspaceId, creatorId: userId })
-    ).id;
-    const user = { id: userId } as User;
-    const workspace = { id: workspaceId } as Workspace;
-
-    // A prior FINISHED assistant row that already carries inline parts — the needy
-    // predicate must SKIP it (no step fetch), returned untouched.
-    const finishedParts = assistantParts(
-      [{ text: 'done earlier', toolCalls: [], toolResults: [] } as any],
-      '',
-    );
-    await createMessage(db, {
-      workspaceId,
-      chatId,
-      role: 'assistant',
-      content: 'done earlier',
-      status: 'completed',
-      metadata: { parts: finishedParts, toolTraceVersion: 2, stepsPersisted: 1 },
-      createdAt: new Date(Date.now() - 3000),
-    });
-
-    // The mid-run row a crash-before-finalize left behind: a step marker only
-    // (parts:[] , content:''), status 'streaming'. Its real parts live ONLY in the
-    // steps table.
-    const midRun = await createMessage(db, {
-      workspaceId,
-      chatId,
-      role: 'assistant',
-      content: '',
-      status: 'streaming',
-      metadata: stepMarkerMetadata(2),
-      createdAt: new Date(Date.now() - 1000),
-    });
-
-    const step0 = assistantParts(
-      [
-        {
-          text: 'reasoning about the page',
-          toolCalls: [
-            { toolCallId: 'g1', toolName: 'getPage', input: { id: 'p1' } },
-          ],
-          toolResults: [
-            { toolCallId: 'g1', toolName: 'getPage', output: { id: 'p1', body: 'B' } },
-          ],
-        } as any,
-      ],
-      '',
-    );
-    const step1 = assistantParts(
-      [{ text: 'partial synthesis so far', toolCalls: [], toolResults: [] } as any],
-      '',
-    );
-    await stepRepo.insertStep(midRun.id, workspaceId, 0, step0);
-    await stepRepo.insertStep(midRun.id, workspaceId, 1, step1);
-
-    // Workspace-scoping guard: a step row for the SAME message id under a DIFFERENT
-    // workspace must NEVER leak into this workspace's reconstruction.
-    await stepRepo.insertStep(midRun.id, otherWorkspaceId, 99, [
-      { type: 'text', text: 'FOREIGN_WORKSPACE_LEAK' },
-    ]);
-
-    const res = await buildController().getMessages(
-      { chatId } as any,
-      { limit: 50 } as any,
-      user,
-      workspace,
-    );
-
-    const items = res.items as any[];
-    const finished = items.find((r) => r.status === 'completed');
-    const reconstructed = items.find((r) => r.id === midRun.id);
-
-    // The finished row passed through with its inline parts unchanged.
-    expect(finished.metadata.parts).toEqual(finishedParts);
-
-    // The mid-run row's parts were reconstructed from the two step rows, in order,
-    // exactly as assembleStepParts concatenates them — the client seed sees the
-    // persisted progress with no change to itself.
-    const expected = assembleStepParts([
-      { stepIndex: 0, parts: step0 },
-      { stepIndex: 1, parts: step1 },
-    ] as any);
-    expect(reconstructed.metadata.parts).toEqual(expected);
-    // The foreign-workspace step row did NOT leak in.
-    expect(JSON.stringify(reconstructed.metadata.parts)).not.toContain(
-      'FOREIGN_WORKSPACE_LEAK',
-    );
-    // Sanity: reconstruction produced real content (text + the paired tool part +
-    // the second step's text), not an empty fallback.
-    expect(reconstructed.metadata.parts).toContainEqual({
-      type: 'text',
-      text: 'reasoning about the page',
-    });
-    expect(
-      (reconstructed.metadata.parts as any[]).some((p) => p.type === 'tool-getPage'),
-    ).toBe(true);
-  }, 60_000);
-});

apps/server/test/integration/ai-chat-run-step-repo.int-spec.ts

@@ -1,163 +0,0 @@
-import { randomBytes } from 'crypto';
-import { Kysely } from 'kysely';
-import { AiChatRunStepRepo } from '@docmost/db/repos/ai-chat/ai-chat-run-step.repo';
-import { AiChatMessageRepo } from '@docmost/db/repos/ai-chat/ai-chat-message.repo';
-import {
-  assistantParts,
-  reconstructRunParts,
-  hydrateAssistantParts,
-  stepMarkerMetadata,
-  rowHasInlineParts,
-} from '../../src/core/ai-chat/ai-chat.service';
-import {
-  getTestDb,
-  destroyTestDb,
-  createWorkspace,
-  createUser,
-  createChat,
-} from './db';
-
-/**
- * #492 append-persist — the reconstruct CONTRACT on a live Postgres. Proves that a
- * turn persisted the NEW way (per-step rows in `ai_chat_run_steps`, only a step
- * marker on the message row) reconstructs to the SAME UI parts as a turn persisted
- * the OLD way (full `metadata.parts` inline on the row, no step rows) — so the
- * era-switch is invisible to attach / delta-poll / export. Real repos + real jsonb
- * roundtrip, not a mock (a mock cannot prove the parts survive the jsonb column
- * byte-identical).
- */
-type Step = {
-  text: string;
-  toolCalls: Array<{ toolCallId: string; toolName: string; input: unknown }>;
-  toolResults: Array<{ toolCallId: string; toolName: string; output: unknown }>;
-};
-
-// A realistic step: some text + a getPage tool call whose ~100 KB body is
-// INCOMPRESSIBLE random base64 (a 'x'.repeat filler would TOAST away and hide the
-// real bytes). Under MAX_TOOL_OUTPUT_BYTES (200 KB) it is stored uncompacted.
-function makeStep(i: number, outputBytes = 4_000): Step {
-  const body = randomBytes(Math.ceil(outputBytes * 0.75)).toString('base64');
-  return {
-    text: `step ${i} text`,
-    toolCalls: [
-      { toolCallId: `c${i}`, toolName: 'getPage', input: { id: `p${i}` } },
-    ],
-    toolResults: [
-      {
-        toolCallId: `c${i}`,
-        toolName: 'getPage',
-        output: { id: `p${i}`, title: `Page ${i}`, body },
-      },
-    ],
-  };
-}
-
-describe('AiChatRunStepRepo + reconstruct contract [integration]', () => {
-  let db: Kysely<any>;
-  let stepRepo: AiChatRunStepRepo;
-  let msgRepo: AiChatMessageRepo;
-  let workspaceId: string;
-  let userId: string;
-  let chatId: string;
-
-  beforeAll(async () => {
-    db = getTestDb();
-    stepRepo = new AiChatRunStepRepo(db as any);
-    msgRepo = new AiChatMessageRepo(db as any);
-    workspaceId = (await createWorkspace(db)).id;
-    userId = (await createUser(db, workspaceId)).id;
-    chatId = (await createChat(db, { workspaceId, creatorId: userId })).id;
-  });
-
-  afterAll(async () => {
-    await destroyTestDb();
-  });
-
-  const seedRow = (metadata: unknown, status: string) =>
-    msgRepo.insert({
-      chatId,
-      workspaceId,
-      userId,
-      role: 'assistant',
-      content: '',
-      status,
-      metadata: metadata as never,
-    });
-
-  it('insertStep is idempotent per (message, stepIndex) and reads back in order', async () => {
-    const row = await seedRow(stepMarkerMetadata(0), 'streaming');
-    const parts0 = assistantParts([makeStep(0)], '');
-    const parts1 = assistantParts([makeStep(1)], '');
-
-    expect(await stepRepo.insertStep(row.id, workspaceId, 0, parts0)).toBe(true);
-    expect(await stepRepo.insertStep(row.id, workspaceId, 1, parts1)).toBe(true);
-    // A retried persist of the SAME step is a no-op (ON CONFLICT DO NOTHING).
-    expect(await stepRepo.insertStep(row.id, workspaceId, 0, parts0)).toBe(
-      false,
-    );
-
-    const steps = await stepRepo.findByMessage(row.id, workspaceId);
-    expect(steps.map((s) => s.stepIndex)).toEqual([0, 1]);
-
-    // Batch fetch groups by message id in step order.
-    const map = await stepRepo.findByMessageIds([row.id], workspaceId);
-    expect(map.get(row.id)!.map((s) => s.stepIndex)).toEqual([0, 1]);
-  });
-
-  it('a NEW-style (step-table) run reconstructs identically to an OLD-style (inline) run', async () => {
-    const steps = [makeStep(10), makeStep(11)];
-    // The inline parts the OLD full-row flush would have written.
-    const fullParts = assistantParts(steps, '');
-
-    // OLD-style record: full parts inline on the row, NO step rows.
-    const oldRow = await seedRow(
-      { parts: fullParts, toolTraceVersion: 2, stepsPersisted: 2 },
-      'completed',
-    );
-
-    // NEW-style record: only a step marker on the row + per-step rows.
-    const newRow = await seedRow(stepMarkerMetadata(2), 'streaming');
-    for (let i = 0; i < steps.length; i++) {
-      await stepRepo.insertStep(
-        newRow.id,
-        workspaceId,
-        i,
-        assistantParts([steps[i]], ''),
-      );
-    }
-
-    // Re-read both from the DB (proves the jsonb roundtrip).
-    const oldFetched = await msgRepo.findById(oldRow.id, workspaceId);
-    const newFetched = await msgRepo.findById(newRow.id, workspaceId);
-    const oldSteps = await stepRepo.findByMessage(oldRow.id, workspaceId);
-    const newSteps = await stepRepo.findByMessage(newRow.id, workspaceId);
-
-    // The discriminator: the old row carries inline parts, the new one does not.
-    expect(rowHasInlineParts(oldFetched!)).toBe(true);
-    expect(rowHasInlineParts(newFetched!)).toBe(false);
-    expect(oldSteps).toHaveLength(0);
-    expect(newSteps).toHaveLength(2);
-
-    const oldRecon = reconstructRunParts(oldFetched!, oldSteps);
-    const newRecon = reconstructRunParts(newFetched!, newSteps);
-
-    // Both reconstruct to the SAME parts + step count — the era is invisible.
-    expect(newRecon.parts).toEqual(fullParts);
-    expect(oldRecon.parts).toEqual(fullParts);
-    expect(newRecon.parts).toEqual(oldRecon.parts);
-    expect(newRecon.stepsPersisted).toBe(2);
-    expect(oldRecon.stepsPersisted).toBe(2);
-
-    // hydrateAssistantParts fills the new row's metadata.parts to match the old
-    // row's inline parts — so a consumer reading `metadata.parts` off the raw row
-    // (the client seed/poll, export) is unchanged across the era.
-    const map = await stepRepo.findByMessageIds([newRow.id], workspaceId);
-    const [hydrated] = hydrateAssistantParts([newFetched!], map);
-    expect((hydrated.metadata as { parts: unknown }).parts).toEqual(fullParts);
-    // A row that already has inline parts passes through untouched (same ref-shape).
-    const [oldPassThrough] = hydrateAssistantParts([oldFetched!], map);
-    expect((oldPassThrough.metadata as { parts: unknown }).parts).toEqual(
-      fullParts,
-    );
-  });
-});

apps/server/test/integration/ai-chat-run-steps-migration.int-spec.ts

@@ -1,63 +0,0 @@
-import { Kysely, sql } from 'kysely';
-import {
-  up,
-  down,
-} from '../../src/database/migrations/20260708T120000-ai-chat-run-steps';
-import { getTestDb, destroyTestDb } from './db';
-
-/**
- * #492 migration up/down roundtrip on a LIVE Postgres. global-setup already
- * migrated docmost_test to latest (so the table exists at start); this drives the
- * migration's own down()/up() and asserts the table presence toggles, then leaves
- * it PRESENT (up) so the shared test DB is intact for any spec that runs after.
- */
-async function tableExists(db: Kysely<any>): Promise<boolean> {
-  const row = (
-    await sql<{ t: string | null }>`select to_regclass('ai_chat_run_steps') as t`.execute(
-      db,
-    )
-  ).rows[0];
-  return row.t !== null;
-}
-
-async function uniqueIndexExists(db: Kysely<any>): Promise<boolean> {
-  const row = (
-    await sql<{
-      t: string | null;
-    }>`select to_regclass('ai_chat_run_steps_message_step_uidx') as t`.execute(db)
-  ).rows[0];
-  return row.t !== null;
-}
-
-describe('20260708 ai_chat_run_steps migration roundtrip [integration]', () => {
-  let db: Kysely<any>;
-
-  beforeAll(() => {
-    db = getTestDb();
-  });
-
-  afterAll(async () => {
-    // Belt-and-suspenders: guarantee the table is present for later specs even if
-    // an assertion threw mid-roundtrip.
-    if (!(await tableExists(db))) await up(db);
-    await destroyTestDb();
-  });
-
-  it('down() drops the table+index and up() recreates them (idempotent)', async () => {
-    // Starts applied (global-setup migrated to latest).
-    expect(await tableExists(db)).toBe(true);
-    expect(await uniqueIndexExists(db)).toBe(true);
-
-    await down(db);
-    expect(await tableExists(db)).toBe(false);
-    expect(await uniqueIndexExists(db)).toBe(false);
-
-    await up(db);
-    expect(await tableExists(db)).toBe(true);
-    expect(await uniqueIndexExists(db)).toBe(true);
-
-    // up() is idempotent (ifNotExists) — a second run is a harmless no-op.
-    await expect(up(db)).resolves.not.toThrow();
-    expect(await tableExists(db)).toBe(true);
-  });
-});

apps/server/test/integration/ai-chat-stream.int-spec.ts

@@ -322,21 +322,20 @@ describe('AiChatService.stream [integration]', () => {
   });
 
   /**
-   * #332 + #490 deferred tool loading, the ON path. Turn 1 starts COLD (CORE +
-   * loadTools only) and activates a deferred tool via loadTools; that activation
-   * is PERSISTED into the chat's metadata.activatedTools (#490) so the NEXT turn
-   * SEEDS from it and the tool is active from the fresh turn's FIRST step — the
-   * model never re-runs loadTools to re-activate the same tool. The unit tests
-   * only exercise pure prepareAgentStep with hand-fed Sets; this pins the real
-   * wiring end-to-end (loadTools.execute -> activatedTools -> persist -> next-turn
-   * seed -> prepareStep -> per-step activeTools) against the real streamText loop.
-   * We drive a MockLanguageModelV3 whose step 1 calls loadTools(['createPage'])
-   * and assert, via the model's recorded per-step CallOptions.tools (the AI SDK
-   * filters the provider tool list by activeTools), that the deferred tool becomes
-   * active on the SAME turn's next step AND, seeded from metadata, on the next
-   * turn's first step.
+   * #332 deferred tool loading, the ON path. The riskiest property is that the
+   * per-turn `activatedTools` Set is created FRESH inside each stream() call, so a
+   * tool a previous turn activated via loadTools is NOT still active when the next
+   * turn starts — the new turn begins "cold" (CORE + loadTools only). The unit
+   * tests only exercise pure prepareAgentStep with hand-fed Sets; this pins the
+   * real wiring end-to-end (loadTools.execute -> activatedTools -> prepareStep ->
+   * per-step activeTools) against the real streamText loop, and proves there is no
+   * cross-turn leak. We drive a MockLanguageModelV3 whose step 1 calls
+   * loadTools(['createPage']) and assert, via the model's recorded per-step
+   * CallOptions.tools (the AI SDK filters the provider tool list by activeTools),
+   * that the deferred tool becomes active on the SAME turn's next step but NOT on a
+   * fresh turn's first step.
    */
-  describe('deferred tool loading ON — cross-turn activation persistence (#332 + #490)', () => {
+  describe('deferred tool loading ON — per-turn activation, no leak (#332)', () => {
     // A stub deferred (non-core) tool the agent can activate. Its execute is never
     // called — the model only needs to SEE it become active — but it must be a
     // valid AI-SDK tool so the SDK includes it in a step's tool list once active.
@@ -452,7 +451,7 @@ describe('AiChatService.stream [integration]', () => {
       } as any);
     }
 
-    it('activates a deferred tool for the SAME turn, and a NEW turn SEEDS it from persisted chat metadata (#490)', async () => {
+    it('activates a deferred tool for the SAME turn, and a NEW turn starts cold (no leak)', async () => {
       const chatId = (await createChat(db, { workspaceId, creatorId: userId })).id;
 
       // --- Turn 1: loadTools(createPage) on step 1, then answer on step 2. ---
@@ -475,7 +474,7 @@ describe('AiChatService.stream [integration]', () => {
       // Step 2 of the SAME turn sees the just-activated deferred tool.
       expect(step2Tools).toContain('createPage');
 
-      // --- Turn 2 on the SAME chat: seeds the persisted activation (#490). ---
+      // --- Turn 2 on the SAME chat: must start cold again. ---
       const model2 = new MockLanguageModelV3({
         doStream: async () => ({ stream: successStream() }),
       } as any);
@@ -486,10 +485,9 @@ describe('AiChatService.stream [integration]', () => {
 
       const nextTurnFirstStep = toolNames(model2.doStreamCalls[0]);
       expect(nextTurnFirstStep).toContain('loadTools');
-      // #490: activation PERSISTS across turns — turn 1 wrote createPage into the
-      // chat's metadata.activatedTools, so the next turn seeds from it and the
-      // deferred tool is active from the FIRST step (no need to re-run loadTools).
-      expect(nextTurnFirstStep).toContain('createPage');
+      // The activated set is per-turn: the prior turn's createPage did NOT leak,
+      // so the fresh turn's first step sees it deferred again.
+      expect(nextTurnFirstStep).not.toContain('createPage');
     });
   });
 });

packages/mcp/src/lib/drawio-xml.ts

@@ -178,11 +178,10 @@ function sliceModel(xml: string): string | null {
 // --- decode chain ----------------------------------------------------------
 
 /**
- * Read the `content=` attribute out of a `.drawio.svg` string. Docmost writes
- * the mxfile XML entity-encoded there (buildDrawioSvg / createDrawioSvg), which
- * is also how draw.io's own SVG export stores it; older attachments stored a
- * base64 payload instead. The DOM decodes entities for us, so the caller only
- * has to distinguish "starts with '<'" (raw XML) from base64.
+ * Read the `content=` attribute out of a `.drawio.svg` string. Docmost stores a
+ * base64 payload there (createDrawioSvg); draw.io's own SVG export may store the
+ * XML entity-encoded instead. The DOM decodes entities for us, so the caller
+ * only has to distinguish "starts with '<'" (raw XML) from base64.
  */
 export function extractContentAttr(svg: string): string {
   const { doc, error } = parseXml(svg);
@@ -196,23 +195,12 @@ export function extractContentAttr(svg: string): string {
   // value itself never contains a double-quote (base64 / entity-encoded XML).
   const m = /content="([^"]*)"/.exec(svg);
   if (m) {
-    // Decode the handful of XML entities a raw regex would leave encoded. The
-    // numeric char-refs for tab/newline/CR MUST be decoded here too: the DOM
-    // path above turns them back into the literal control chars, so this
-    // regex fallback has to agree or the two decode paths diverge (#507).
-    // `&amp;` is decoded last so an escaped `&amp;#x9;` reads back as the
-    // literal text `&#x9;`, not a tab.
+    // Decode the handful of XML entities a raw regex would leave encoded.
     return m[1]
       .replace(/&lt;/g, "<")
       .replace(/&gt;/g, ">")
       .replace(/&quot;/g, '"')
       .replace(/&#39;/g, "'")
-      .replace(/&#x9;/gi, "\t")
-      .replace(/&#9;/g, "\t")
-      .replace(/&#xa;/gi, "\n")
-      .replace(/&#10;/g, "\n")
-      .replace(/&#xd;/gi, "\r")
-      .replace(/&#13;/g, "\r")
       .replace(/&amp;/g, "&");
   }
   throw new Error("drawio: SVG has no content= attribute to decode");
@@ -319,16 +307,9 @@ export function encodeDrawioFile(modelXml: string, title = "Page-1"): string {
 /**
  * Build the `diagram.drawio.svg` attachment. Mirrors the import service's
  * createDrawioSvg contract exactly:
- *   <svg xmlns=… xmlns:xlink=… content="${xmlEscape(drawioFile)}">${inner}</svg>
+ *   <svg xmlns=… xmlns:xlink=… content="${base64(drawioFile)}">${inner}</svg>
  * plus width/height/viewBox from the diagram bounding box and the schematic
  * preview as the visible children (`inner`).
- *
- * The `content=` value is the mxfile XML XML-entity-escaped (draw.io's own
- * native form), NOT base64. draw.io's editor decodes a base64 content= via
- * Latin-1 atob (no UTF-8 step), turning every non-ASCII char (e.g. Cyrillic,
- * ё, —) into mojibake; the entity-encoded form is decoded by the DOM as UTF-8
- * and opens intact. Our decoder (decodeDrawioSvg) reads both forms, so old
- * base64 attachments still round-trip.
  */
 export function buildDrawioSvg(
   modelXml: string,
@@ -337,14 +318,14 @@ export function buildDrawioSvg(
   title = "Page-1",
 ): string {
   const file = encodeDrawioFile(modelXml, title);
-  const content = xmlEscape(file);
+  const base64 = Buffer.from(file, "utf-8").toString("base64");
   const w = Math.max(1, Math.round(bbox.width));
   const h = Math.max(1, Math.round(bbox.height));
   return (
     `<svg xmlns="http://www.w3.org/2000/svg" ` +
     `xmlns:xlink="http://www.w3.org/1999/xlink" ` +
     `width="${w}" height="${h}" viewBox="0 0 ${w} ${h}" ` +
-    `content="${content}">${inner}</svg>`
+    `content="${base64}">${inner}</svg>`
   );
 }
 
@@ -353,15 +334,7 @@ function xmlEscape(s: string): string {
     .replace(/&/g, "&amp;")
     .replace(/</g, "&lt;")
     .replace(/>/g, "&gt;")
-    .replace(/"/g, "&quot;")
-    // A literal tab/newline/CR inside an attribute value is collapsed to a
-    // single space by XML attribute-value normalization on DOM read (both jsdom
-    // here and the real draw.io editor), silently flattening multi-line labels
-    // and tab-bearing values. Numeric char-refs survive that normalization, so
-    // emit them the way draw.io's own native export does (#507).
-    .replace(/\t/g, "&#x9;")
-    .replace(/\n/g, "&#xa;")
-    .replace(/\r/g, "&#xd;");
+    .replace(/"/g, "&quot;");
 }
 
 // --- normalization + hash --------------------------------------------------

packages/mcp/test/unit/drawio-xml.test.mjs

@@ -51,14 +51,6 @@ const hasRule = (issues, rule, cellId) =>
     (i) => i.rule === rule && (cellId === undefined || i.cellId === cellId),
   );
 
-// Reverse the attribute-value XML escaping used in the `content=` payload.
-const unescapeAttr = (s) =>
-  s
-    .replace(/&lt;/g, "<")
-    .replace(/&gt;/g, ">")
-    .replace(/&quot;/g, '"')
-    .replace(/&amp;/g, "&");
-
 // --- style parsing ---------------------------------------------------------
 
 test("parseStyle: base stylename + key=value pairs", () => {
@@ -343,20 +335,16 @@ test("encode/build: a title with < > \" & round-trips without corrupting the SVG
 
   // The outer content="..." attribute must not be broken by the title: the raw
   // title metacharacters never appear literally in the SVG markup (they are
-  // entity-escaped inside content=, doubly so where the file XML already escaped
-  // them inside name="...").
+  // base64-encoded inside content=, and escaped inside the file XML).
   const contentMatch = /content="([^"]*)"/.exec(svg);
   assert.ok(contentMatch, "SVG has a single well-formed content= attribute");
 
   // The diagram model still decodes losslessly despite the exotic title.
   assert.equal(decodeDrawioSvg(svg), model);
 
-  // content= is now entity-encoded XML (draw.io's native form), never base64.
-  assert.match(contentMatch[1], /^&lt;mxfile/);
-
   // The file XML is well-formed: the title lives in name="..." as escaped
   // entities, so unescaping recovers the original title byte-for-byte.
-  const fileXml = unescapeAttr(contentMatch[1]);
+  const fileXml = Buffer.from(contentMatch[1], "base64").toString("utf-8");
   const nameMatch = /<diagram id="[^"]*" name="([^"]*)">/.exec(fileXml);
   assert.ok(nameMatch, "the diagram name attribute is intact and quote-safe");
   const decodedTitle = nameMatch[1]
@@ -370,112 +358,3 @@ test("encode/build: a title with < > \" & round-trips without corrupting the SVG
   const file = encodeDrawioFile(model, title);
   assert.match(file, /name="A &lt; B &gt; C &quot; D &amp; E">/);
 });
-
-// --- #507: content= is entity-encoded XML, never base64 --------------------
-
-// A model whose cell values carry Cyrillic, ё and an em dash — exactly the
-// characters draw.io's Latin-1 atob mangles when content= is base64.
-const CYRILLIC_MODEL =
-  "<mxGraphModel><root>" +
-  '<mxCell id="0"/><mxCell id="1" parent="0"/>' +
-  '<mxCell id="2" value="Старт-бит — ёж" style="rounded=1;html=1;" vertex="1" parent="1">' +
-  '<mxGeometry x="100" y="100" width="120" height="60" as="geometry"/></mxCell>' +
-  "</root></mxGraphModel>";
-
-test("#507: buildDrawioSvg writes content= as entity-encoded XML (not base64)", () => {
-  const model = normalizeXml(CYRILLIC_MODEL);
-  const svg = buildDrawioSvg(model, "<g/>", { width: 200, height: 120 }, "Диаграмма");
-
-  const contentMatch = /content="([^"]*)"/.exec(svg);
-  assert.ok(contentMatch, "SVG has a single well-formed content= attribute");
-  const content = contentMatch[1];
-
-  // The content= value is the entity-encoded mxfile XML — starts with `&lt;mxfile`.
-  assert.match(content, /^&lt;mxfile/, "content= is entity-encoded mxfile XML");
-  // It must NOT be a base64 blob: base64 has no XML entities and no literal `<`.
-  assert.ok(content.includes("&lt;"), "content= carries XML entities, not base64");
-
-  // Cyrillic / ё / — survive verbatim in the attribute (raw UTF-8, not atob-mangled).
-  assert.ok(content.includes("Старт-бит — ёж"), "non-ASCII value is raw UTF-8 in content=");
-  assert.ok(content.includes("Диаграмма"), "non-ASCII title is raw UTF-8 in content=");
-  // The mojibake that base64+atob would have produced must be absent.
-  assert.ok(!content.includes("Ð"), "no Latin-1 mojibake in content=");
-});
-
-test("#507: Cyrillic model round-trips byte-stable through buildDrawioSvg -> decodeDrawioSvg", () => {
-  const model = normalizeXml(CYRILLIC_MODEL);
-  const svg = buildDrawioSvg(model, "<g/>", { width: 200, height: 120 }, "Заголовок — ё");
-  assert.equal(decodeDrawioSvg(svg), model);
-});
-
-test("#507 back-compat: an OLD base64-form .drawio.svg still decodes losslessly", () => {
-  const model = normalizeXml(CYRILLIC_MODEL);
-  // Reproduce the pre-fix write path: encodeDrawioFile -> base64 in content=.
-  const file = encodeDrawioFile(model, "Старая диаграмма");
-  const base64 = Buffer.from(file, "utf-8").toString("base64");
-  const svg =
-    `<svg xmlns="http://www.w3.org/2000/svg" content="${base64}"><g/></svg>`;
-  // No XML entities, purely base64 alphabet — this is the legacy form.
-  assert.ok(!base64.includes("<") && !base64.includes("&"));
-  assert.equal(decodeDrawioSvg(svg), model);
-});
-
-test("#507 negative: non-ASCII in a cell value AND in the title round-trip clean", () => {
-  const model = normalizeXml(CYRILLIC_MODEL);
-  const title = "Тест — ёмкость № 5";
-  const svg = buildDrawioSvg(model, "<g/>", { width: 200, height: 120 }, title);
-
-  // Model recovered byte-for-byte.
-  assert.equal(decodeDrawioSvg(svg), model);
-
-  // Title lands in the <diagram name="..."> of the decoded file XML, intact.
-  const content = /content="([^"]*)"/.exec(svg)[1];
-  const fileXml = unescapeAttr(content);
-  const nameMatch = /<diagram id="[^"]*" name="([^"]*)">/.exec(fileXml);
-  assert.ok(nameMatch, "diagram name attribute present");
-  assert.equal(unescapeAttr(nameMatch[1]), title);
-});
-
-// --- #507 F1: literal tab/newline/CR in an attribute value survive the
-// content= round-trip. The whole mxfile XML lives in one content="..." attr;
-// XML attribute-value normalization collapses a LITERAL tab/newline/CR to a
-// single space on DOM read, so the escape must emit numeric char-refs instead.
-const CTRL_MODEL =
-  "<mxGraphModel><root>" +
-  '<mxCell id="0"/><mxCell id="1" parent="0"/>' +
-  '<mxCell id="2" value="col1\tcol2" style="html=1;\nshadow=0" vertex="1" parent="1">' +
-  '<mxGeometry x="10" y="10" width="120" height="60" as="geometry"/></mxCell>' +
-  '<mxCell id="3" value="Line1\nLine2\rLine3" vertex="1" parent="1">' +
-  '<mxGeometry x="10" y="100" width="120" height="60" as="geometry"/></mxCell>' +
-  "</root></mxGraphModel>";
-
-test("#507 F1: literal tab/newline/CR in a value round-trip byte-stable (DOM decode)", () => {
-  const svg = buildDrawioSvg(CTRL_MODEL, "<g/>", { width: 200, height: 200 });
-  const content = /content="([^"]*)"/.exec(svg)[1];
-  // The tab/newline/CR must be emitted as numeric char-refs, never as literal
-  // control chars (which DOM attribute-value normalization would eat).
-  assert.ok(
-    !/[\t\n\r]/.test(content),
-    "no literal tab/newline/CR survive in the content= attribute",
-  );
-  assert.ok(
-    content.includes("&#x9;") &&
-      content.includes("&#xa;") &&
-      content.includes("&#xd;"),
-    "tab/newline/CR are emitted as numeric char-refs",
-  );
-  // Full DOM decode recovers the model byte-for-byte, control chars intact.
-  assert.equal(decodeDrawioSvg(svg), CTRL_MODEL);
-});
-
-test("#507 F1: regex fallback decodes tab/newline/CR char-refs (agrees with DOM path)", () => {
-  const svg = buildDrawioSvg(CTRL_MODEL, "<g/>", { width: 200, height: 200 });
-  const content = /content="([^"]*)"/.exec(svg)[1];
-  // A bare `&` makes the SVG wrapper malformed, forcing extractContentAttr onto
-  // its regex fallback branch. That branch must decode the tab/newline/CR
-  // char-refs exactly like the DOM path, or the two decoders diverge.
-  const malformedSvg = `<svg content="${content}">&</svg>`;
-  assert.equal(decodeDrawioSvg(malformedSvg), CTRL_MODEL);
-  // The well-formed (DOM) path yields the identical result.
-  assert.equal(decodeDrawioSvg(svg), CTRL_MODEL);
-});