Compare commits

...

36 Commits

Author SHA1 Message Date
agent_coder 8503ff1f3d fix(ai-chat): hydrate crashed mid-run steps for model replay + cover write/read seams (#492)
F1: the server model-replay loaded history via findAllByChat().map(rowToUiMessage)
WITHOUT hydrating parts from ai_chat_run_steps. A HARD crash mid-run (SIGKILL/OOM)
fires no terminal callback, so the assistant row stays parts:[] and its partial
tool-calls/results/text (durable in the steps table) dropped out of the model's
next-turn context. Hydrate needy assistant rows (role==='assistant' &&
!rowHasInlineParts) via findByMessageIds + hydrateAssistantParts before the replay
map — mirroring the controller's withReconstructedParts exactly — guarded on the
optional repo. Fix the now-false interrupt-resume comment.

F2: add a service int-spec that drives the REAL onStep append-persist WRITE branch
through AiChatService.stream with a real AiChatRunStepRepo injected, asserting the
per-step rows' stepIndex + parts slice and the step-marker metadata match a
single-row flush (catches an stepsPersisted-1 off-by-one).

F3: add a controller int-spec that drives withReconstructedParts through getMessages
WITH the repo present (a mid-run marker-only row + its step rows), asserting the
reconstructed metadata.parts and workspace-scoping.

F4: remove the dead countByMessage (zero prod callers; reconstructRunParts derives
stepsPersisted inline) + its now-unused sql import and the redundant test assertion.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-12 07:50:50 +03:00
agent_coder ae3dfd8de6 Merge remote-tracking branch 'gitea/develop' into feat/492-incremental-render 2026-07-12 05:34:07 +03:00
agent_vscode 03eafa6c68 Merge remote-tracking branch 'gitea/develop' into develop 2026-07-12 05:16:51 +03:00
agent_vscode a42f1ead48 test(ai-chat): align #332 deferred-tool test with #490 cross-turn persistence
The #332 integration test (ai-chat-stream.int-spec.ts) asserted a deferred
tool activated in one turn does NOT leak into the next ("cold start per
turn"). #490 (f5bbfdb2) deliberately reversed that: it persists the
activation set into chat metadata.activatedTools and seeds the next turn
from it, so the model need not re-run loadTools. The test only surfaced now
because the token-estimate CI fix let the integration step run again.

Adopt #490 persistence as the intended behavior:
- flip the turn-2 assertion to expect createPage IS active on the fresh
  turn's first step (seeded from metadata); keep turn-1 cold-start assertions.
- rewrite the test docstring, describe/it titles and comments accordingly.
- fix two stale "not persisted / per-turn" comments in ai-chat.service.ts
  (prepareAgentStep + the streaming-loop activation block); no logic change.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-12 05:16:47 +03:00
vvzvlad b7a3ec227d Merge pull request 'fix(drawio): content= — entity-XML вместо base64 (кириллица в редакторе) (#507)' (#521) from fix/507-drawio-cyrillic into develop
Reviewed-on: #521
2026-07-12 05:08:57 +03:00
agent_coder 846341d7d4 fix(drawio): encode literal tab/newline/CR in content= as numeric char-refs (#507 review)
Review follow-up to the base64→entity-XML content= switch. The whole mxfile XML
now lives in one content="..." attribute; XML attribute-value normalization
collapses a LITERAL tab/newline/CR to a single space on DOM read (jsdom and the
real draw.io editor alike), silently flattening multi-line labels and
tab-bearing values that the old base64 form stored verbatim.

Finding 1 (data-loss): both encode paths — buildDrawioSvg's xmlEscape (mcp) and
the import service's escape — now append &#x9;/&#xa;/&#xd; after the four
&<>" replaces (numeric char-refs survive normalization, as draw.io's own export
does). The extractContentAttr regex fallback now decodes those char-refs (hex
case-insensitive plus decimal &#9;/&#10;/&#13;) so it agrees with the DOM path;
&amp; stays decoded last so an escaped &amp;#x9; reads back as literal text.

Finding 2 (dedup): the server's private xmlEscapeAttr is replaced by the shared
htmlEscape helper (& < > " ' — a strict superset, the extra ' is harmless in a
"-delimited value) wrapped in xmlEscapeContent, which adds the three control-char
char-refs on top (htmlEscape does not escape them).

Finding 3 (docs): narrow the CHANGELOG healing claim — only a diagram still
holding its original correct-UTF-8 base64 (not yet opened/autosaved) is
recoverable; one already opened in the editor persisted mojibake at rest and its
text is lost.

Tests: new mcp round-trip test with literal tab/newline/CR in a value (DOM path,
byte-stable) plus a fallback-branch test forcing a malformed wrapper so both
decode paths are proven to agree; new server spec asserting char-ref encoding.
Mutation-checked: dropping the encode replaces reddens both new mcp tests;
dropping only the fallback decode reddens just the fallback test.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-12 05:05:52 +03:00
agent_coder 32e10ca6d3 fix(drawio): content= пишется entity-XML, не base64 — кириллица не мойбейкает в редакторе (#507)
Реviewer-filed баг: draw.io-редактор декодит base64 content= как Latin-1
(atob-семантика) → кириллица разваливается в мойбейк, автосейв редактора
персистит порчу и убивает превью. Нативная форма draw.io — entity-encoded
mxfile-XML (content="&lt;mxfile…"), DOM-декодится как UTF-8.

Фикс двух write-путей: buildDrawioSvg (mcp, все create/update) и
createDrawioSvg (server Confluence-импорт) теперь XML-эскейпят content=
вместо base64. Декодер уже различает startsWith("&lt;") vs base64 — обе формы
читаются, старые base64-файлы открываются (back-compat), byte-stable
round-trip. CHANGELOG + заметка про лечение старых диаграмм (drawioGet→Update).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-12 05:05:52 +03:00
vvzvlad 74d212cfd3 Merge pull request 'feat(auth): серверная API-key аутентификация агентов на /mcp — фаза 1 (#501)' (#526) from feat/501-mcp-apikey-auth into develop
Reviewed-on: #526
2026-07-12 05:02:51 +03:00
agent_coder c96fafc4ad perf(ai-chat): append-персист шагов — per-step INSERT вместо переписи строки (#492)
Раньше каждый onStepFinish переписывал ВСЮ строку ассистента (растущий
metadata.parts jsonb со всеми выводами инструментов) → O(n²) объёма записи
на прогон: под MVCC/TOAST апдейт jsonb переписывает всю версию строки, так
что шаг k пишет ~k×вывод. Прогон из 50 шагов по ~100 КБ = сотни МБ WAL и
мёртвых кортежей за ход, что молотит autovacuum. (#490 убрал только ВТОРУЮ
копию в tool_calls; сам metadata.parts всё ещё рос и переписывался.)

Теперь каждый завершённый шаг ДОПИСЫВАЕТСЯ отдельной строкой в лёгкую
таблицу ai_chat_run_steps (только парты этого шага), а строка сообщения
получает дешёвый маркер (stepsPersisted + toolTraceVersion, без растущего
блоба parts). Полный metadata.parts собирается ОДИН раз на финализации.
НЕ jsonb-append (||): апдейт всё равно переписывает всю TOAST-версию —
экономится только сетевой payload, а WAL/мёртвые кортежи остаются; поэтому
именно ОТДЕЛЬНАЯ таблица + INSERT.

Три обязательные интеграции:
- reconstructRunParts(row, stepRows) → { parts, stepsPersisted }: единый
  шов переключения бэкенда. Читает парты из СТРОКИ, если она уже несёт
  inline-parts (старые записи + ЛЮБАЯ финализированная), иначе из ТАБЛИЦЫ
  ШАГОВ (mid-run запись #492). Дискриминатор — наличие непустого
  metadata.parts (флаг схемы не нужен). Потребители (attach-seed,
  delta-poll, export, reconnect) прогоняют строки через hydrateAssistantParts
  на границе чтения — их контракт/вывод не меняется, старые и новые записи
  восстанавливаются идентично.
- сигнал ротации кольца реестра #491 (confirmPersistedStep) теперь стреляет
  на подтверждённый INSERT шага, под тем же контрактом (updateStreaming
  возвращает stepsPersisted / null).
- era-marker toolTraceVersion (#490) больше не ставится полной переписью —
  ставится в маркере шага и на финализации (flushAssistant), остаётся
  консистентным.

Полная обратная совместимость: прогон, записанный по-старому (полная строка,
без строк шагов), восстанавливается/attach/export идентично. При отсутствии
репозитория шагов (позиционные тест-конструкции) — фолбэк на прежний
полнострочный flush (без регрессии, только без выигрыша WAL).

Тесты (реальный pg, int-lane):
- WAL-дельта (pg_current_wal_lsn) на прогоне 40×100КБ: new=4.3МБ vs
  old=90.3МБ (20.8x) — O(Σ шагов) против O(n²); старый путь в тесте И есть
  ревертнутое поведение (мутация-проверка).
- reconstruct-контракт: new-style (таблица шагов) и old-style (inline) прогоны
  восстанавливаются в идентичные parts; hydrate заполняет строку.
- миграция up/down roundtrip.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-12 05:02:25 +03:00
vvzvlad 9a6389133b Merge pull request 'fix(page-tree/realtime): insertByPosition теряет непоказанных детей (#525)' (#531) from fix/525-insert-unloaded into develop
Reviewed-on: #531
2026-07-12 05:01:38 +03:00
agent_vscode 0b8497e496 Merge remote-tracking branch 'gitea/develop' into develop 2026-07-12 04:53:26 +03:00
agent_vscode 433252bdb1 fix(ci): build @docmost/token-estimate before its non-nx consumers (#490)
The shared @docmost/token-estimate package (main: ./dist/index.js, dist/
gitignored) was never built in the CI jobs that bypass nx, so develop CI
went red:
- test.yml `test` (pnpm -r test): client vitest failed to resolve the import.
- develop.yml `e2e-server` (direct jest e2e): server history-budget.ts hit
  TS2307 Cannot find module '@docmost/token-estimate'.
The nx-driven jobs (build, e2e-mcp) stayed green via dependsOn: ^build.

- test.yml: add "Build token-estimate" step to the `test` job.
- develop.yml: add "Build token-estimate" step to the `e2e-server` job.
- Dockerfile: ship packages/token-estimate/dist + package.json into the
  installer stage — apps/server depends on it (workspace:*) and imports it
  at runtime, so the prod image would otherwise crash with
  ERR_MODULE_NOT_FOUND (masked so far because publish/smoke never ran).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-12 04:53:22 +03:00
agent_coder 2f8c5d9a98 perf(client): инкрементальный рендер стрима ответа (#492)
Путь ответа ассистента теперь рендерится инкрементально, по образцу
StreamingPlainText из ветки reasoning (#492, волна C эпика #497).

Раньше MarkdownPart прогонял ВЕСЬ накопленный ответ через канонический
конвейер (markdownToProseMirrorSync → PMNode.fromJSON → DOMSerializer →
DOMPurify) на КАЖДОМ throttled-тике (~20 Гц). На синтетическом потоке
~100 КБ это 394 вызова renderChatMarkdown — O(числа тиков), причём каждый
вызов заново парсит всю растущую строку.

Теперь:
- StreamingMarkdownText делит текст на блоки по безопасному срезу
  (splitPlainChunks — та же append-only-инвариантность, что у reasoning):
  СТАБИЛИЗИРОВАННЫЕ блоки идут через канонический конвейер и мемоизируются
  (каждый блок парсится РОВНО ОДИН раз), живой ХВОСТ — дешёвый plain-text
  (React-escaped, без парсера/санитайзера/innerHTML) до стабилизации.
- На финализации (флип state → done или конец хода) — ОДИН полный
  канонический рендер всего текста: побайтовая визуальная паритетность с
  прежним выводом (включая <li><p>-обёртки схемы и scoped-CSS из #498).
- Гейт liveness тот же, что у ReasoningBlock: streaming =
  turnStreaming && part.state === "streaming".

Также цикл рендера частей переведён на ИСЧЕРПЫВАЮЩИЙ switch по видам частей
с never-проверкой в default (вместо прежнего WARNING-комментария): новый
закрытый вид части в UIMessagePart теперь ошибка компиляции.

Тесты:
- perf-smoke: на ~100 КБ потоке число вызовов renderChatMarkdown ≤ blocks+2
  и ≪ ticks (O(блоков), не O(тиков)); мутация (снять memo с MarkdownChunk)
  краснит ассерт (78631 вызовов).
- visual-regression: финальный рендер побайтово равен renderChatMarkdown
  всего текста (в т.ч. <li><p>), инкрементальный вид сходится к нему на
  финализации; учтён neutralizeInternalLinks.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-12 04:35:06 +03:00
vvzvlad 6f81067f4d Merge pull request 'feat(#370 PR-1): ядро версий страниц — kind + ручной Save/idle/boundary триггеры' (#374) from feat/370-page-versioning into develop
Reviewed-on: #374
2026-07-12 04:33:39 +03:00
agent_coder 55e8f61b3c Merge remote-tracking branch 'gitea/develop' into feat/501-mcp-apikey-auth 2026-07-12 04:33:17 +03:00
agent_vscode ab133fa0d0 Merge remote-tracking branch 'gitea/develop' into develop 2026-07-12 04:24:56 +03:00
agent_coder d42ca8dc57 test(auth): cover collab-token api-key arm-seam; docs + dead-code cleanup (#501)
Hardening follow-ups from PR #526 review (no defect fixes):

- DO1: add auth.controller.spec tests for the collab-token handler, the
  arm-seam of the anti-laundering defense. Assert getCollabToken is invoked
  with { apiKeyId } only when the SIGNED req.raw marks an api_key principal,
  and undefined otherwise (session, missing apiKeyId, or a spoofed body).
  Verified non-vacuous: nulling the ternary reddens the ARMED test.
- DO2: document the API_KEYS_ENABLED kill-switch in .env.example next to the
  other feature flags (default ON; strict true/false; OFF denies api-key auth
  and 404s the management endpoints).
- DO3: remove dead bindAccessJwtVerifier (+ its now-orphaned AccessJwtVerifier
  interface and its dedicated spec block); prod switched to
  bindMcpBearerVerifier. verifyBearerAccess is retained (still used).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-12 04:19:52 +03:00
vvzvlad 25a31e6c0d Merge pull request 'perf(ai-chat): resume-стек #491 — восстановление в develop (осиротел при stacked-мерже #518)' (#540) from feat/491-resume-stack into develop
Reviewed-on: #540
2026-07-12 04:16:54 +03:00
agent_coder 3550bfa411 fix(ai-chat): ре-ревью #518 — CI-фиделити delta, attach-дубль, чистки (#491)
Ребейз на обновлённый feat/490 (PR #510, Option-A): стек #491 перенесён с
6e872f2d на 6e42752a; run-fsm.ts/chat-thread.tsx приехали из develop-версии
(W1/S4: RUN_ALREADY_ACTIVE несёт activeRunId + supersede-CAS адопция) — мои
#491-правки легли сверху без конфликтов, W1 НЕ откачен (run-fsm.ts == develop
байт-в-байт; RUN_ALREADY_ACTIVE диспатчит activeRunId).

1. [CI-fidelity] Дельта-курсор спек падал в CI unit-лейне: `.spec.ts` дефолтил
   на НЕмигрированный docmost → 5/6 ERROR relation does not exist (skip-гард ловит
   только connection-fail). Переименован в `*.int-spec.ts` (исключается из unit-
   regex `.spec.ts$`, гоняется в test:int, чей global-setup мигрирует docmost_test)
   + DSN по умолчанию → docmost_test. Теперь 6/6 реально исполняются в CI-верном
   окружении; overlap-мутация роняет RACE-тесты (не вакуумен).

2. [regression #137/#161] attach: отсутствие `n` схлопывалось в frontier 0 →
   finished-неротированный ран (coverageFloor 0) отдавал ВЕСЬ tail вместо 204;
   парамслесс/легаси-вкладка допишет полный replay → дубль. Различаем ОТСУТСТВИЕ
   `n` (null — не tail-aware) от `n=0` (tail-aware): контроллер шлёт null при
   missing/invalid; registry.attach(n: number|null) 204-ит finished-ран при
   n===null (старый `finished && !expectLive` гейт), n=0 по-прежнему отдаёт хвост.
   Тесты (registry + controller) + mutation-verify: нейтрализация гейта роняет их.

4. [conventions] Ring-кап env-var переименован RUN_STREAM_MAX_BUFFER_BYTES →
   AI_CHAT_RUN_STREAM_MAX_BUFFER_BYTES (префикс как у сиблингов) + запись в
   .env.example (дефолт 4MB, 0/invalid→дефолт, subscriber-cap=2×).

5. [docs] run-fsm.spec.md: item4 переписан («реализовано в #491 — дельта несёт
   run:{id,status}|null; клиент run-поле ещё не потребляет») + добавлена строка
   перехода POLL_IDLE_CAP stopping→idle (Review #4, редьюсер это делает).

6. [simplification] Удалена мёртвая цепочка reconstructRunParts /
   reconstructPartsFromRow (ноль прод-вызовов) + опц. messageRepo-инъекция в
   AiChatRunService + спек-блоки; вернётся с первым реальным вызывателем. Маркер
   metadata.stepsPersisted (реально используемый) сохранён.

DROP-пункты ревьюера (осиротевший import, DELTA_POLL_MAX_ROWS) не трогаю.
Прогон: server tsc 0, ai-chat unit 202, delta-int 6/6 (int-lane), int attach 6/6,
client vitest 403, tsc client 0 ai-chat, mcp 834/0. FSM-инварианты #488 сохранены.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-12 04:09:00 +03:00
agent_coder a0ecf21cb5 fix(ai-chat): ре-ревью #491 — дубль на getRun-fail + epoch RUN_FACT + doc (#491)
Ре-ревью нашло регрессию класса #137/#161 на пути отказа getRun при tail-only
re-seed:

- НАХОДКА 1 (MEDIUM/HIGH): в onFinish-disconnect ветка .catch (отказ getRun) на
  локальном дропе входила в reconnect-ladder БЕЗ ре-сида и БЕЗ фильтра «живой»
  частичной строки. anchorRef оставался устаревшим (mount-инициализатор), живая
  частичная строка со шагами N..M-1 — последней; через ~1с реконнект строил
  ?anchor=&n=N_mount, и при живом ране с покрытием от N_mount (flaky-сеть: SSE и
  getRun упали, сеть поднялась за 1с) сервер отдавал кадры ≥N_mount → SDK дописывал
  их к строке, где они УЖЕ есть → дублирование (клиентского дедупа реплея против
  parts нет). Фикс: восстановлена структурная гарантия удалённого resumeStream-
  фильтра — на отказе getRun (и на no-persisted-row, и на no-cid) живая частичная
  строка удаляется из стора по id + anchorRef=null → реконнект реплеит со start в
  ЧИСТЫЙ стор (полная пересборка) либо 204→poll. Нет пути, где attach tail-applies
  на строку с уже присутствующими шагами. Тест: getRun-reject на локальном
  дисконнекте → живая строка отфильтрована + URL без параметров (mutation-verify:
  без фикса тест краснеет — фильтр не срабатывает).

- НАХОДКА 2 (LOW): RUN_FACT в enterReconnect теперь epoch-штампуется (epoch:
  stampEpoch), как везде (postRun): getRun-rtt расширяет окно onFinish→dispatch,
  конкурентный SEND_LOCAL во время rtt теперь дропает устаревший RUN_FACT по I1,
  а не перетирает runFact.runId нового хода.

- НАХОДКА 3 (LOW, doc): run-fsm.spec.md обновлён — stripRef/strippedRowRef →
  anchorRef {id, stepsPersisted}, tail-only + re-seed-from-persist.

FSM run-fsm.ts не тронут; инварианты #488 (epoch/honor-in-stopping/ownership-reset/
disconnect-first/render-gate) сохранены. Клиент ai-chat vitest 399 зелёный, tsc 0
ai-chat-ошибок; сервер delta(6, реально исполняется)/registry/step-marker/attach +
integration attach — зелёное.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-12 04:09:00 +03:00
agent_coder d81781aa27 feat(ai-chat): клиентская активация tail-only resume + delta-поллинг (#491)
Переводит клиент на серверный tail-only контракт resume (задел commit 3),
не трогая FSM run-fsm.ts — меняется только рантайм-обвязка в chat-thread.tsx.

A. Убран STRIP-механизм. Seed теперь содержит ВСЕ персистнутые строки без
   изъятия хвоста (стриминговый хвост — это шаги 0..N-1, к которым SDK-
   продолжение дописывает tail). stripRef/strippedRowRef заменены на anchorRef
   { id, stepsPersisted } — персистнутая assistant-строка, питающая
   ?anchor=<id>&n=<stepsPersisted>. Восстановления stripped-строки на
   204/NONE/starved удалены (строку никто не изымал — нечего восстанавливать);
   invalidateQueries + диспатчи FSM сохранены. Блок anchor-mismatch в reconcile
   сверяется по id из свежей персист-истории, а не по «живой» строке.

B. Вход в attaching/reconnecting — ВСЕГДА через re-seed из персиста; «живой»
   стор НИКОГДА не база для tail-apply. На локальном FINISH_DISCONNECT (и на
   live-follow повторном дропе observer-а) сначала getRun(chatId) → замена
   «живой» частичной строки персистнутой по id (mergeById) + установка anchor,
   и лишь ПОСЛЕ этого диспатч RUN_FACT + FINISH_DISCONNECT (который планирует
   реконнект). Так attach не может продублировать частичный шаг N. Фильтр
   «живой» строки в resumeStream-эффекте убран (его заменяет re-seed). Инварианты
   FSM (I1 epoch-штамп, I4 honor-in-stopping, DISCONNECT-FIRST, сброс ownership на
   терминалах, render-gate) сохранены.

C. URL attach: ?anchor=<id>&n=<stepsPersisted> при наличии якоря, без expect.

D. Degraded-поллинг переведён с полного рефетча всех страниц на дельту:
   useAiChatMessagesQuery больше не поллит (seed один раз), а окно при
   degradedPoll раз в 2.5с зовёт getAiChatMessagesDelta(chatId, cursor) и
   идемпотентно по id мёржит строки в тот же infinite-query кэш через новый
   чистый хелпер mergeDeltaRowsIntoPages. Арминг/разарм (onResumeFallback) и
   idle-cap не тронуты.

Хелперы: seedRows удалён; добавлены stepsPersistedOf и mergeDeltaRowsIntoPages
(+ юнит-тесты на идемпотентность). Тесты chat-thread обновлены под новый URL,
seed-без-стрипа, re-seed-из-персиста на дисконнекте (mutation-verify: падают
без re-seed и при n мимо персиста) и 204→poll-без-restore. Весь ai-chat vitest
зелёный (398), tsc без новых ошибок.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-12 04:09:00 +03:00
agent_coder daca5ce8d6 fix(ai-chat): целостность delta-spec (молчаливый скип) + 2 hardening (#491)
Внутреннее ре-ревью: DB-backed delta-spec молча СКИПАЛСЯ — нулевое покрытие
инварианта DB-clock курсора.

- Импорт `import postgres from 'postgres'` (default) при tsconfig
  module:commonjs без esModuleInterop компилился в `postgres_1.default(...)`,
  а CJS-`postgres` не имеет `.default` → TypeError в beforeAll → пустой catch →
  reachable=false → все 6 тестов уходили в console.warn('SKIP') и return БЕЗ
  ассертов (suite оставался бы зелёным даже при регрессии на new Date()).
  Фикс: `import * as postgres from 'postgres'` (как в рабочем int-харнессе).
- Хардненг харнесса: реальная ошибка программирования в beforeAll больше НЕ
  маскируется под «DB unreachable» — скип легитимен только для сетевого отказа
  (ECONNREFUSED и т.п.), иначе rethrow → suite падает громко.
- Два DB-clock теста использовали jest.useFakeTimers() целиком, что замораживало
  внутренние таймеры postgres.js → awaited DB round-trip зависал на 5s-капе (и
  вешал afterAll). Фейкаем ТОЛЬКО Date (doNotFake всех таймеров) — запрос
  резолвится, а инвариант «стамп от часов БД, не app-clock» по-прежнему доказан
  (скос процесс-часов в 2099 → стамп остаётся на времени БД). Теперь все 6
  тестов РЕАЛЬНО исполняются и зелёные против живого Postgres.

Два дешёвых hardening из ревью:
- registry coverageFloor: пустая ветка возвращает max(currentStamp,
  persistedFloor) — инвариант «клиент с n=persistedFloor всегда покрыт»
  структурный, а не тайминг-зависимый.
- GetChatDeltaDto.cursor: @IsString → @IsISO8601 — битый курсор отсекается 400
  на уровне DTO, а не 500 на `::timestamptz`.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-12 04:09:00 +03:00
agent_coder daeeb1f3f2 refactor(ai-chat): registry — step-aligned retention + tail-only attach (#491)
Реестр ран-стримов больше не буферизует до 32МБ сырых SSE-кадров на активный
ран и не выливает весь буфер в сокет синхронно при attach — это давало OOM на
1ГБ-контейнере при нескольких марафонских ранах. Теперь кольцо ограничено
(env-настраиваемо, по умолчанию 4МБ) и держится в границах за счёт ротации по
шагам.

Серверная часть (суть коммита):

- Штамповка кадров по шагам в ingestFrame. Штамп кадра = число `finish-step`
  кадров ДО него (с 0); сам finish-step несёт текущее значение, затем счётчик
  инкрементится. Так штамп совпадает с `metadata.stepsPersisted`: клиент с N
  персистнутыми шагами имеет 0..N-1 в сиде и просит хвост `stamp >= N`. Границу
  ловим дешёвым startsWith по `data: {"type":"finish-step"` — форма кадра
  проверена эмпирически против ai@6.0.207 (одна часть на кадр, type всегда
  первый ключ; кавычки в text-delta экранированы, ложных срабатываний нет).

- Кольцо ротируется ТОЛЬКО на подтверждённом персисте шага N
  (`confirmPersistedStep`), сбрасывая кадры `stamp < N` (эти шаги уже на диске и
  придут в свежем сиде). `updateStreaming` теперь СИГНАЛИЗИРУЕТ исход (число
  персистнутых шагов или null), и ротация вызывается лишь при не-null возврате —
  провал персиста ничего не ротирует, кольцо покрывает БОЛЬШЕ (анти-инверсия:
  наивная ротация в .then() после НЕзаписанного шага дырявила бы гарантию).

- Переполнение кольца сверх байтового капа вытесняет старейшие кадры; вытеснение
  ещё-не-персистнутого кадра открывает GAP. Гэп НЕ липкий: floor покрытия
  считается из кольца, поздний персист, проротировав дырявые шаги, его чистит.

- attach(chatId, anchor, n): маркер шага N приходит ТОЛЬКО от клиента (сервер не
  читает строку — N из устаревшего сида дал бы тихую дыру в один шаг). Покрытие
  ОК ⟺ coverageFloor <= n; иначе 204 → клиент рефетчит (больший N) и
  переподключается. Хвост = синтетический `start`-кадр (ран-факт runId/chatId) +
  кадры `stamp >= n`. Инвариант 6 (нет кросс-ран реплея) сохранён через anchor;
  инвариант 4 (снапшот+регистрация в один синхронный тик) сохранён. N-срез
  применяется во ВСЕХ ветках, включая finished-retained: finished + N=N_final →
  пустой хвост + finish-кадр, клиент закрывает стрим.

- Контроллер пишет хвост чанками с учётом drain (writeTailRespectingDrain), а не
  синхронным залпом (вторая половина OOM). Кап подписчика — производное 2× кап
  кольца, обе величины env-резолвятся на инстансе.

Клиент: в тип строки добавлен `metadata.stepsPersisted` (источник N). PIN-SPEC
трип-вайр на ai@6.0.207: `readUIMessageStream({ message })` продолжает последнее
сообщение, `start`-кадр не сбрасывает parts, текст не пересекает finish-step —
на этом держится продолжение при attach; апгрейд ai теперь падает громко.

Тесты (observable-property против РЕАЛЬНОГО реестра/БД): детектор границы на
реальной форме кадра, N-срез (в т.ч. посреди шага), ротация только на
подтверждённом персисте, «персист провалился но кольцо влезло → attach успешен /
провал + переполнение → 204», «устаревший N → 204 → после рефетча успех», очистка
гэпа поздним персистом, finished-retained + N_final, memory-bound (5 параллельных
марафонов сверх 32МБ, каждое кольцо ≤ кап). Обновлены registry/controller specs и
DB-backed интеграционный attach-spec под новую сигнатуру/семантику.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-12 04:09:00 +03:00
agent_coder a919c79cb2 feat(ai-chat): персист step-маркера (#491)
Персист отставал от live-стрима на целый текущий шаг, а run.stepCount как
источник границы шага НЕгоден: recordStep — fire-and-forget, не атомарен с
записью parts (рассинхрон сид↔маркер). Пишем маркер ТЕМ ЖЕ flush'ем, что и
parts:

- flushAssistant стампует metadata.stepsPersisted = число ЗАВЕРШЁННЫХ шагов,
  чьи parts лежат в ЭТОЙ строке (оба выводятся из finished → маркер не может
  разойтись с persisted parts). In-progress хвост (частичный шаг при
  error/abort или mid-stream flush) НЕ считается. Это step-alignment якорь, на
  котором строится resume-стек (ротация кольца по подтверждённому шагу N —
  коммит 3; attach режет хвост по «шаг > N»).
- Контракт AiChatRunService.reconstructRunParts(runId) → { parts,
  stepsPersisted } — единственный интерфейс чтения ЖИВОГО рана (run →
  assistantMessageId → строка → чистый reconstructPartsFromRow). null при
  отсутствии рана / связанной строки / удалённой строки; маркер 0 у pre-#491
  строки — безопасный пол. Потребители: attach (коммит 3), дельта (rows уже
  несут маркер в metadata), экспорт. messageRepo — опциональный 3-й параметр
  конструктора (2-арг тест-конструкции компилируются без изменений).
- /messages и дельта отдают маркер в row автоматически (он внутри metadata).

Тесты: property «stepsPersisted == число завершённых шагов для любого N + parts
согласованы»; частичный хвост не инкрементит маркер; reconstructPartsFromRow
(маркер, дефолт 0, фолбэк на content); reconstructRunParts (резолв + null-кейсы).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-12 04:09:00 +03:00
agent_coder 798a81abfe perf(ai-chat): дельта-эндпоинт поллинга + run-факт в ответе (#491)
Degraded-poll рефетчил ВСЕ страницы infinite-query каждые 2.5 c с полными
parts. Вводим дельта-эндпоинт «строки, изменённые после курсора»:

- POST /ai-chat/messages/delta → { rows, cursor, run: {id,status}|null }.
  Курсор — таймстамп часов БД (now()), клиент эхом возвращает его каждый
  поллинг. Окно перекрытия now()−5s ловит строку, закоммиченную с updatedAt
  чуть раньше момента снятия предыдущего курсора на другом autocommit-
  соединении (одиночные UPDATE, длинных транзакций нет). Окно ГАРАНТИРУЕТ
  повторы → контракт: клиентский merge идемпотентен по id (mergeById).
- Run-факт едет В дельте (отдельный /run-поллинг удвоил бы QPS — отвергнуто).
- Все дельта-релевантные записи (message update/finalizeOwner/reconcile/sweep,
  run update/finalizeIfActive/markStopRequested/sweepRunning) стампуют
  updatedAt через SQL now(), а не app-clock new Date(): единая монотонная ось
  курсора, смешанные источники часов были независимым источником пропусков.
- Клиент: сервис-функция getAiChatMessagesDelta + фиксация контракта
  идемпотентности mergeById тестом (свап degraded-поллинга на дельту — в
  коммите 3, вместе с re-seed-путём attach).

Тесты (observable-property на живом Postgres, не моки):
- дельта-семантика + монотонность курсора;
- RACE «коммит позже с updatedAt раньше курсора» — ловится окном перекрытия
  (наивный updatedAt > cursor пропустил бы);
- окно гарантирует повторы → идемпотентный merge;
- updatedAt стампуется часами БД, а не app-clock (фейк системных часов в 2099
  — стамп остаётся на времени БД);
- контроллер: owner-gate, форма ответа, run-факт только {id,status}.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-12 04:09:00 +03:00
agent_coder 216499d57b Merge remote-tracking branch 'gitea/develop' into feat/370-page-versioning
# Conflicts:
#	CHANGELOG.md
2026-07-12 03:31:04 +03:00
agent_vscode bb5abf29a6 docs(readme): add self-hosted embeddings server guide
Document how to run a local Hugging Face TEI embeddings server for the
AI agent's RAG search, in both README.md and README.ru.md:
- Option A: local container on the same Docker network (no auth)
- Option B: separate host exposed via Traefik + Let's Encrypt (API key,
  rate limit, external curl check)
- settings tables (Workspace settings -> AI -> Embeddings) and notes on
  vector dimension (384), weight caching, version pinning, offline, GPU
2026-07-12 01:21:35 +03:00
agent_coder 0ddeaadeee feat(collab): закрытие api-key→collab отмывания (fail-closed дискриминатор)
api_key-принципал стал полноправным REST-принципалом → мог начеканить 24-ч
COLLAB-токен, который AuthenticationExtension не сверял с api_keys, и ревокация
не доставала websocket-редактирование. Блокировать /collab-token для api_key
нельзя (это правки агентов). Решение fail-closed:

- JwtCollabPayload расширен полем principal ('session'|'api_key') + apiKeyId.
  generateCollabToken штампует его в КАЖДЫЙ токен: 'api_key'+apiKeyId, когда
  чеканит api-key-принципал (внешний MCP-агент), иначе 'session'. Дискриминатор
  ключуется на ОРИГИНЕ (наличии api-key), НЕ на actor:'agent' — внутренний
  session-backed AI-агент остаётся 'session' и на no-check пути.
- /auth/collab-token читает подписью-выведенные req.raw.authType/apiKeyId
  (не тело) и прокидывает origin через getCollabToken → generateCollabToken.
- doAuthenticate: при principal='api_key' на connect прогоняет
  ApiKeyService.validate (отозванный ключ → новых collab-подключений нет; инфра
  пробрасывается). api_key без apiKeyId — reject. Claimless-токен доверяется
  только В ПРЕДЕЛАХ 24-ч grace-окна роллаута (легаси session-токен, api_key его
  начеканить не мог); после окна всякий валидный токен обязан нести
  дискриминатор, поэтому claimless — баг и отвергается (не молча-доверие 24ч).

CollaborationModule импортирует ApiKeyModule.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-12 00:57:15 +03:00
agent_coder bc632f9d56 feat(mcp): приём api_key на /mcp через allowlist-роутер {ACCESS, API_KEY}
/mcp-Bearer больше не пинит тип к ACCESS. Новый framework-free примитив
verifyMcpBearer верифицирует подпись ОДИН раз (bindMcpBearerVerifier пинит
allowlist {ACCESS, API_KEY}) и роутит:

- API_KEY → сперва биндинг к воркспейсу инстанса (чужой воркспейс — отказ), затем
  общий row-check ApiKeyService.validate. HMAC-верификация (микросекунды) вместо
  bcrypt, это НЕ попытка логина — Basic-путь и анти-брутфорс-лимитер не
  затрагиваются (фикс удушения параллельных чтений агентов).
- ACCESS → прежний verifyBearerAccess (session-active + not-disabled), которому
  передаётся замыкание над уже проверенным payload, так что «проверить один раз»
  и «helper без изменений» сосуществуют.

Bearer-нога resolveMcpSessionConfig униформизирована: любой отказ авторизации →
единый generic-401 (анти-enumeration, класс отказа не течёт в тело), а
непредвиденная (инфра) ошибка ПРОБРАСЫВАЕТСЯ (→5xx), не маскируется под 401.
McpService получает ApiKeyService; McpModule импортирует ApiKeyModule.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-12 00:57:15 +03:00
agent_coder d3a049d176 feat(api-key): core-модуль выдачи и валидации api-ключей + kill-switch
Форк не несёт EE-модуль ee/api-key, поэтому api_key-токен сегодня отвергается на
любом /api/*. Вводим core-модуль:

- ApiKeyRepo (kysely): insert/findById/findByCreator/findAllInWorkspace/
  softDelete/touchLastUsed, фильтр deleted_at IS NULL.
- ApiKeyService.create — mint-then-insert: сперва uuid7 id, затем чеканка JWT
  без exp, затем строка (сбой чеканки → строки нет; потерянный ответ →
  осиротевшая строка видна в list и самолечится). Токен отдаётся один раз, в
  таблице не хранится.
- ApiKeyService.validate — единый валидатор для jwt.strategy И /mcp. Владелец
  истины о сроке/отзыве — строка api_keys, не JWT. Любой определённый негатив
  (нет строки / expired / disabled / workspace mismatch / kill-switch off) →
  единый bare-401 (анти-enumeration); инфра-ошибка пробрасывается (→5xx), не
  маскируется. Без кэша — немедленная ревокация. last_used_at троттлится 1ч,
  fire-and-forget.
- REST create/list/revoke: create — self + UserThrottlerGuard; list/revoke —
  admin (CASL Manage на API) видит/отзывает все ключи воркспейса, член — только
  свои. api_key-принципал отвергается на всех трёх (токен не управляет
  токенами — GitHub-PAT-семантика). Аудит API_KEY_CREATED/DELETED + структурный
  лог с apiKeyId и IP.
- jwt.strategy: прямой вызов ApiKeyService.validate вместо динамического
  EE-require; штампует req.raw.authType='api_key' и apiKeyId для гварда выше.
- Kill-switch API_KEYS_ENABLED: дефолт ON (деплой без переменной не убивает
  агентов), строгий парс @IsIn(['true','false']) (=0/=off падают на boot, а не
  молча значат «включено»), boot-лог состояния. off → validate deny и
  эндпоинты 404.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-12 00:57:15 +03:00
agent_coder d738780370 feat(auth): verifyJwtOneOf + чеканка api-key токенов без exp
TokenService получает примитив type-routing`verifyJwtOneOf(token, allowed[])`:
подпись проверяется один раз, тип обязан входить в явный allowlist, иначе — тот
же generic-throw, что и у verifyJwt. Нужен для /mcp-Bearer, который принимает и
ACCESS, и API_KEY на одном слоте, но без переиспользуемого footgun'а
«verify-and-return-any-type».

generateApiToken теперь чеканит api-key токен БЕЗ клейма exp. Единственный
источник истины о сроке/отзыве ключа — строка api_keys (проверяется на каждом
запросе), не JWT. Общий jwtService зарегистрирован с глобальным
signOptions.expiresIn ('90d'), который мержится в любой sign() — даже
sign(payload, {}) — а {expiresIn: undefined} бросает; поэтому «бессрочный» ключ
через общий signer молча получал бы exp=now+90d. Чеканка идёт через выделенный
signer без expiresIn (issuer 'Docmost' для паритета клеймов). Живой баг
подтверждён эмпирически и покрыт тестом.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-12 00:25:16 +03:00
agent_coder daf728676f fix(#370): ревью r4 — F8-близнец в handleSaveVersion + утечка idleBurstStart
WARNING 1 [stability]: handleSaveVersion — write-path-близнец F8-бага. Деструктивный
popContributors (Redis SPOP, не откатывается с PG-tx) внутри executeTx; commit-abort
реджектит СНАРУЖИ колбэка, inner-catch не срабатывает → потеря атрибуции. Фикс:
poppedForRestore/versionedPageId объявлены ДО executeTx + внешний try/catch
восстанавливает (идемпотентный addContributors) на любом tx/commit-abort throw;
inner-catch обнуляет трекер после своего восстановления. Ровно одно восстановление
в каждой ветке. Зеркалит history.processor.ts (F8).

WARNING 2 [stability]: idleBurstStart Map текла + промахивалась по ключу. Рекей
page.id->documentName (как сиблинги) + cleanup в afterUnloadDocument; хаускипинг:
remove idle-job по page.id (реальный jobId), delete маркер по documentName —
page.<slugId> (#260) больше не промахивается.

Simplification: history-list переиспользует historyKindMeta().version.
+3 теста (F8-twin restore, slugId housekeeping, Map cleanup), mutation-verified.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-12 00:13:46 +03:00
agent_coder 1d8e3444f4 fix(#370): ребейз на develop + правки ревью agent_vscode (раунд 3)
Ребейз ветки на текущий gitea/develop (устраняет mergeable:false).
Конфликты разрешены вручную:
- editor-atoms.ts: сохранён type-only импорт Editor (сплит-код с develop),
  добавлен import type HocuspocusProvider из #370.
- collaboration/constants.ts: оставлен EMBED_DEBOUNCE_MS (embed-дебаунс develop),
  убраны более не используемые HISTORY_* (их единственный потребитель — старая
  эвристика computeHistoryJob — удалён в #370), добавлены idle-константы и
  PageHistoryKind.
- persistence.extension.ts: 3-way смёржены метрики/#348/#402 develop поверх
  idle-конвейера #370; computeHistoryJob остаётся idle-версией без остатков.

Миграция переименована 20260705T120000 -> 20260707T120000 (класс #361):
таймстамп был занят perf-indexes на develop; новый строго позже свежайшей
на develop (20260706T120000-search-lookup-trgm). Содержимое не менялось,
внешних ссылок на имя файла нет (Kysely находит по директории).

Документирование (findings 1-2):
- remove-vs-active гонка в idiome remove()->add() enqueuePageHistory: окно,
  где отложенная job уходит в active между remove (проглатывается на active) и
  add (BullMQ отбрасывает add с существующим jobId); ограничено и
  самовосстанавливается (следующий store перевзводит), кроме худшего случая —
  гонящий store был ПОСЛЕДНИМ в сессии: хвостовые правки без trailing-снапшота
  до следующей правки. Явно указано, почему нельзя «унифицировать» с соседним
  embed-дебаунсом (стабильный jobId, без remove).
- допущение single-process у Map idleBurstStart: в памяти процесса, рестарт
  collab теряет метки начала всплеска -> непрерывный всплеск через рестарт может
  ждать до 2x cap. Ограничено и безопасно.

Тест (finding 3): интеграционный тест idle-конвейера против реального BullMQ
(короткие интервалы через jest.mock констант): непрерывный всплеск в неск. cap
-> периодические idle-снапшоты не реже cap и не по одному на store;
прерывистый всплеск -> ровно один trailing-снапшот. Жёсткий teardown
(force-close + settle), чтобы фоновый BullMQ не влиял на соседние suite.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-11 23:47:51 +03:00
agent_coder c50f5b66bb fix(#370): thread trx into addPageWatchers (F7 self-deadlock) + restore contributors on commit-failure (F8) + assert the lock (F9) (review round 2)
The round-1 F3 fix (wrapping the processor's find+save in a locked tx) itself
introduced two regressions:

F7 [CRITICAL] addPageWatchers ran WITHOUT trx inside the tx holding FOR UPDATE on
pages[pageId]. The watcher insert's FK check takes FOR KEY SHARE on the same row,
but on a DIFFERENT pool connection — a true self-deadlock (our tx connection sits
idle-in-transaction awaiting the JS await, the insert connection blocks on the
lock). Now passes trx (addPageWatchers already accepts it and routes it through
insertMany), so the FK lock is taken on the connection that already holds FOR
UPDATE — no self-conflict.

F8 [WARNING] popContributors is a destructive Redis SPOP; the inner catch only
restores on a throw INSIDE the callback. A COMMIT failure throws OUTSIDE it,
rolling the snapshot back while the pop is gone → a retry writes an unattributed
version. Now tracks the popped set and restores it in an outer catch (idempotent
SADD), leaving BullMQ to retry with attribution intact.

F9 [WARNING] The spec asserted saveHistory args with a loosened objectContaining
that stopped verifying trx, and never pinned withLock/trx on findById or the trx
on addPageWatchers — which is exactly why F7 slipped. Restored the exact
saveHistory(trx) assertion and added findById({withLock,trx}) + addPageWatchers
trx assertions (the latter would have caught F7), plus a commit-failure test.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-11 23:47:51 +03:00
agent_coder 51d44b6061 fix(#370): ES2021-safe spec, source-specific idle ceiling, processor lock, tested index mapping (review round 1)
F1 [BLOCKER] persistence-store.spec used Array.prototype.at(-1) (ES2022) but the
server targets ES2021, so server tsc failed (TS2550) and ts-jest could not
compile the suite — 22 core manual-save/idle/boundary tests silently did not run
in CI. Replaced with [length - 1] index access.

F2 [WARNING] The idle burst-reset used a hardcoded IDLE_MAX_WAIT_USER for both
tiers, but computeHistoryJob's ceiling is source-specific. On a continuously
agent-edited page the burst marker stayed stale for 5..10m, forcing delay=0 on
every store and writing one idle row per store — the exact per-store bloat the
debounce prevents. The reset now uses the same source-specific max-wait.

F3 [WARNING] The processor did an unlocked findPageLastHistory -> saveHistory,
which TOCTOU-races a concurrent manual-save (that runs under a page-row lock),
producing two page_history rows with identical content (one idle, one manual) and
defeating promote-not-dup. The snapshot decision is now wrapped in executeTx with
the same page-row lock, so the second writer observes the first's committed row
and the isDeepStrictEqual gate collapses the duplicate.

F4 [WARNING] The risky client filtered-index -> full-list mapping had no tests.
Extracted it to a pure resolvePrevSnapshotId(fullItems, id) helper (diff/restore
baseline against the true previous snapshot in the FULL list, never the previous
visible version) and unit-tested it; removed the now-vestigial index threading.

F5/F6 [low] Renamed the misleading ceiling test + fixed its comment; added a
CHANGELOG entry for the user-facing versioning feature.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-11 23:47:51 +03:00
agent_coder 6247585b66 feat(#370): page-history intentionality tiers — kind column + intentional/idle/boundary triggers (PR-1 core)
PR-1 'core' of #370: introduces page_history.kind ('manual'|'agent'|'idle'|
'boundary'; legacy null = autosave) and rebuilds the snapshot triggers around a
three-tier intentionality model. Draft durability (pages/ydoc hocuspocus
autosave) is unchanged; only the frequency and labelling of history points change.

- Migration 20260705T120000: page_history.kind nullable varchar(20), no default.
- Manual Save: one stateless 'save-version' path for human AND agent; kind is
  derived SERVER-SIDE from the signed context.actor (never the payload), readOnly
  connections rejected, the fresh ydoc runs through the existing store path (no
  REST race), then broadcasts version.saved.
- Idle-flush: trailing debounce (one BullMQ job per page, remove-then-readd) with
  IDLE_INTERVAL_USER=60m / AGENT=15m AND a max-wait ceiling
  (IDLE_MAX_WAIT_USER=10m / AGENT=5m) so a continuous editing session can't starve
  the autosnapshot (review round-1 WARNING).
- Boundary: generalized from the user→agent special-case to ANY lastUpdatedSource
  transition (user↔agent↔git), same isDeepStrictEqual gate — covers git-sync free.
- Removed the agent delay=0 fast path and the old HISTORY_FAST_* constants; the
  agent joins the common idle pipeline.
- Promote-not-dup: a manual save on unchanged content promotes the latest
  autosave's kind in place (or no-ops if already manual) instead of duplicating a
  heavy content row.
- Client: mod+S hotkey + menu button (hidden when readOnly), history-panel kind
  badges, dimmed autosaves, a 'versions only' filter (indices map to the full
  list so diff/restore still target the true previous snapshot), live refresh on
  version.saved.

Internal review: APPROVE-with-suggestions; the round-1 WARNING (idle starvation)
is fixed here via the max-wait ceiling, and the generalized-boundary + ceiling
behaviours are pinned with new tests (115 collab/repo specs green, server tsc 0).

Deferred to later PRs: shares.published_mode (PR-2), the save_page_version MCP
tool + role prompts (PR-3), actor='git' wiring into #359 (PR-4).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-11 23:47:51 +03:00
98 changed files with 7950 additions and 953 deletions
+20
View File
@@ -251,6 +251,16 @@ 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
@@ -302,6 +312,16 @@ MCP_DOCMOST_PASSWORD=
# enabled for a workspace, and the same single-instance constraint applies (the
# registry is process-local).
# AI_CHAT_RESUMABLE_STREAM=false
#
# Per-run replay ring cap (#491), in BYTES, for the resumable-stream registry
# above. The registry buffers the run's recent SSE tail so a reopened tab can
# attach and continue from the step it already persisted; the ring is bounded and
# rotates on every confirmed step-persist. This caps the un-persisted tail between
# rotations — an overflow evicts the oldest frames and a late attach falls back to
# 204 -> degraded poll, so correctness never depends on the size. Default 4194304
# (4MB); a 0/invalid value falls back to the default. The per-subscriber backpressure
# cap is derived as 2x this value. Only meaningful with AI_CHAT_RESUMABLE_STREAM on.
# AI_CHAT_RUN_STREAM_MAX_BUFFER_BYTES=4194304
# --- Run lifecycle tunables (#487) ---
# These govern the universal run machinery (every turn is now a first-class run,
+7
View File
@@ -226,6 +226,13 @@ 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
+8
View File
@@ -159,6 +159,14 @@ 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
+26
View File
@@ -135,6 +135,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
new resync path re-reads the live anchor so the suggestion applies against the
current text, and orphaned anchors (whose marked run was deleted) are
reconciled rather than left blocking. (#496)
- **Save intentional page versions.** Press `Cmd/Ctrl+S` (or use the page menu)
to save a named version of a page. The history panel now distinguishes
intentional versions (a "Saved" / "Agent version" badge) from automatic
snapshots, dims autosaves, and offers an "Only versions" filter. Automatic
snapshots switched from a fixed interval to a trailing idle-flush with a
max-wait ceiling, and a boundary snapshot is pinned whenever the editing source
changes (e.g. a person's edits followed by the AI agent). (#370)
- **Place several images side by side in a row.** A new "Inline (side by
side)" alignment mode in the image bubble menu renders consecutive inline
@@ -385,6 +392,25 @@ 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
+8
View File
@@ -59,6 +59,14 @@ 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/
+131
View File
@@ -206,6 +206,137 @@ start the new migrations apply on top of your existing schema (`CREATE EXTENSION
existing pages are indexed on their next edit. pgvector is still required for the migration to
apply at all.
## Local embeddings server
The AI agent's semantic (RAG) search needs an **embeddings model**. Instead of paying a cloud
provider (e.g. OpenAI `text-embedding-3-*`) to embed every page, you can run a small open-weights
model yourself with Hugging Face
[Text Embeddings Inference](https://github.com/huggingface/text-embeddings-inference) (TEI), which
serves an OpenAI-compatible `/v1/embeddings` endpoint. `intfloat/multilingual-e5-small` is a good
default: multilingual, 384-dim, and comfortable on CPU (~1–2 GB RAM, 1–2 vCPU). Point Gitmost at it
under **Workspace settings → AI → Embeddings**.
### Option A — local (same Docker network as Gitmost)
Run TEI as a container on the network Gitmost is already on. The port is never published, so the
endpoint stays internal and needs no authentication.
```yaml
services:
embeddings:
image: ghcr.io/huggingface/text-embeddings-inference:cpu-1.9 # pin version; use a cuda-* tag for GPU
container_name: embeddings
restart: unless-stopped
networks:
- gitmost_net # same network Gitmost is on
command:
- "--model-id"
- "intfloat/multilingual-e5-small"
- "--auto-truncate" # clamp over-long inputs instead of returning 413
volumes:
- tei-models:/data # weights are downloaded once and cached here
networks:
gitmost_net:
external: true # the network Gitmost already uses
volumes:
tei-models:
```
Gitmost settings (**Workspace settings → AI → Embeddings**):
| Field | Value |
|-------------------|-----------------------------------|
| Model | `intfloat/multilingual-e5-small` |
| Base URL | `http://embeddings:80/v1/` |
| Embedding API key | — (leave empty) |
> `embeddings` is the container name — Gitmost resolves it over DNS inside the Docker network.
> The port is not published, so the endpoint is reachable only by containers on that network and
> no authorization is required.
### Option B — separate host (public via Traefik + Let's Encrypt)
This assumes the host already runs Traefik with an ACME resolver (the example below uses
`letsEncrypt`, the `websecure` entrypoint and a shared `docker_main_net` network). Replace the
domain / network / resolver with your own.
**DNS:** add an A record `embeddings.example.com` → the IP of your Traefik host (same
challenge / port 80 as the rest of your sites).
```yaml
services:
embeddings:
image: ghcr.io/huggingface/text-embeddings-inference:cpu-1.9 # pin version; cuda-* tag for GPU
container_name: embeddings
restart: unless-stopped
networks:
- docker_main_net # the network Traefik is attached to
command:
- "--model-id"
- "intfloat/multilingual-e5-small"
- "--auto-truncate"
- "--api-key"
- "sk-emb-REPLACE_WITH_YOUR_KEY"
volumes:
- tei-models:/data
labels:
traefik.enable: "true"
traefik.http.routers.embeddings.rule: "Host(`embeddings.example.com`)"
traefik.http.routers.embeddings.entrypoints: "websecure"
traefik.http.routers.embeddings.tls: "true"
traefik.http.routers.embeddings.tls.certresolver: "letsEncrypt"
traefik.http.routers.embeddings.service: "embeddings"
traefik.http.services.embeddings.loadbalancer.server.port: "80"
# TEI enforces the Bearer key itself; Traefik only rate-limits to protect the CPU
traefik.http.routers.embeddings.middlewares: "embeddings-rl"
traefik.http.middlewares.embeddings-rl.ratelimit.average: "20"
traefik.http.middlewares.embeddings-rl.ratelimit.burst: "40"
traefik.http.middlewares.embeddings-rl.ratelimit.period: "1s"
networks:
docker_main_net:
external: true
volumes:
tei-models:
```
Gitmost settings (**Workspace settings → AI → Embeddings**):
| Field | Value |
|-------------------|---------------------------------------|
| Model | `intfloat/multilingual-e5-small` |
| Base URL | `https://embeddings.example.com/v1/` |
| Embedding API key | your `sk-emb-…` |
Check it from outside:
```bash
curl -s https://embeddings.example.com/v1/embeddings \
-H "Authorization: Bearer sk-emb-REPLACE_WITH_YOUR_KEY" \
-H "Content-Type: application/json" \
-d '{"model":"intfloat/multilingual-e5-small","input":"query: hello"}' \
| python3 -c 'import sys,json;print("dims:",len(json.load(sys.stdin)["data"][0]["embedding"]))'
# -> dims: 384
```
### Embeddings server notes
- **Vector dimension is 384.** If this Gitmost was previously embedded with a different model
(e.g. `text-embedding-3-large` = 3072-dim), the old pgvector rows won't match the new dimension —
clear the existing embeddings / re-index before switching. Gitmost only compares vectors of the
same dimension, so mixed-dimension rows are silently ignored rather than searched.
- **First start downloads the weights** (hundreds of MB) from `huggingface.co` into the
`tei-models` volume; every start after that reads from the volume.
- **Pin the version.** Pin the image, and optionally the model: add `--revision <commit-sha>` to
`command` (the sha is on the model's page on Hugging Face).
- **Air-gapped / no egress:** seed the `tei-models` volume ahead of time and add
`environment: [HF_HUB_OFFLINE=1]`.
- **GPU:** use the cuda tag of the same release (e.g.
`ghcr.io/huggingface/text-embeddings-inference:cuda-1.9`) and start the container with `gpus: all`.
## Features
- Real-time collaboration
+131
View File
@@ -193,6 +193,137 @@ dump/restore, существующий каталог данных переис
> неизменным и бэкапьте вместе с базой данных.
## Локальный сервер эмбеддингов
Семантическому (RAG) поиску AI-агента нужна **модель эмбеддингов**. Вместо оплаты облачного
провайдера (например, OpenAI `text-embedding-3-*`) за эмбеддинг каждой страницы можно запустить
небольшую open-weights модель у себя через Hugging Face
[Text Embeddings Inference](https://github.com/huggingface/text-embeddings-inference) (TEI) — он
отдаёт OpenAI-совместимый эндпоинт `/v1/embeddings`. Хороший дефолт — `intfloat/multilingual-e5-small`:
многоязычная, 384-мерная, комфортно работает на CPU (~1–2 ГБ RAM, 1–2 vCPU). Пропишите её в
**Настройки воркспейса → AI → Эмбеддинги**.
### Вариант A — локально (та же Docker-сеть, что и Gitmost)
Запустите TEI контейнером в той же сети, где уже работает Gitmost. Порт наружу не публикуется,
поэтому эндпоинт остаётся внутренним и не требует авторизации.
```yaml
services:
embeddings:
image: ghcr.io/huggingface/text-embeddings-inference:cpu-1.9 # pin version; use a cuda-* tag for GPU
container_name: embeddings
restart: unless-stopped
networks:
- gitmost_net # same network Gitmost is on
command:
- "--model-id"
- "intfloat/multilingual-e5-small"
- "--auto-truncate" # clamp over-long inputs instead of returning 413
volumes:
- tei-models:/data # weights are downloaded once and cached here
networks:
gitmost_net:
external: true # the network Gitmost already uses
volumes:
tei-models:
```
Настройки Gitmost (**Настройки воркспейса → AI → Эмбеддинги**):
| Поле | Значение |
|-------------------|-----------------------------------|
| Model | `intfloat/multilingual-e5-small` |
| Base URL | `http://embeddings:80/v1/` |
| Embedding API key | — (оставить пустым) |
> `embeddings` — имя контейнера, Gitmost резолвит его по DNS внутри Docker-сети.
> Наружу порт не публикуется, эндпоинт доступен только контейнерам этой сети, поэтому
> авторизация не нужна.
### Вариант B — на отдельном хосте (наружу через Traefik + Let's Encrypt)
Предполагается, что на хосте уже есть Traefik с ACME-резолвером (в примере ниже — `letsEncrypt`,
entrypoint `websecure`, общая сеть `docker_main_net`). Замените домен / сеть / резолвер на свои.
**DNS:** заведите A-запись `embeddings.example.com` → IP хоста с Traefik (тот же challenge / порт 80,
что и у остальных сайтов).
```yaml
services:
embeddings:
image: ghcr.io/huggingface/text-embeddings-inference:cpu-1.9 # pin version; cuda-* tag for GPU
container_name: embeddings
restart: unless-stopped
networks:
- docker_main_net # the network Traefik is attached to
command:
- "--model-id"
- "intfloat/multilingual-e5-small"
- "--auto-truncate"
- "--api-key"
- "sk-emb-REPLACE_WITH_YOUR_KEY"
volumes:
- tei-models:/data
labels:
traefik.enable: "true"
traefik.http.routers.embeddings.rule: "Host(`embeddings.example.com`)"
traefik.http.routers.embeddings.entrypoints: "websecure"
traefik.http.routers.embeddings.tls: "true"
traefik.http.routers.embeddings.tls.certresolver: "letsEncrypt"
traefik.http.routers.embeddings.service: "embeddings"
traefik.http.services.embeddings.loadbalancer.server.port: "80"
# TEI enforces the Bearer key itself; Traefik only rate-limits to protect the CPU
traefik.http.routers.embeddings.middlewares: "embeddings-rl"
traefik.http.middlewares.embeddings-rl.ratelimit.average: "20"
traefik.http.middlewares.embeddings-rl.ratelimit.burst: "40"
traefik.http.middlewares.embeddings-rl.ratelimit.period: "1s"
networks:
docker_main_net:
external: true
volumes:
tei-models:
```
Настройки Gitmost (**Настройки воркспейса → AI → Эмбеддинги**):
| Поле | Значение |
|-------------------|---------------------------------------|
| Model | `intfloat/multilingual-e5-small` |
| Base URL | `https://embeddings.example.com/v1/` |
| Embedding API key | ваш `sk-emb-…` |
Проверка снаружи:
```bash
curl -s https://embeddings.example.com/v1/embeddings \
-H "Authorization: Bearer sk-emb-REPLACE_WITH_YOUR_KEY" \
-H "Content-Type: application/json" \
-d '{"model":"intfloat/multilingual-e5-small","input":"query: hello"}' \
| python3 -c 'import sys,json;print("dims:",len(json.load(sys.stdin)["data"][0]["embedding"]))'
# -> dims: 384
```
### Заметки про сервер эмбеддингов
- **Размерность вектора — 384.** Если раньше этот Gitmost эмбеддился другой моделью
(например, `text-embedding-3-large` = 3072-dim), старые строки в pgvector не совпадут по
размерности — очистите существующие эмбеддинги / переиндексируйте перед переключением. Gitmost
сравнивает только вектора одной размерности, поэтому строки другой размерности не участвуют в
поиске, а не ломают его.
- **Первый старт тянет веса** (сотни МБ) с `huggingface.co` в том `tei-models`; дальше — из тома.
- **Пин версии.** Пиньте образ, а при желании и модель: добавьте в `command` `--revision <commit-sha>`
(sha берётся со страницы модели на Hugging Face).
- **Без egress (air-gapped):** засейте том `tei-models` заранее и добавьте
`environment: [HF_HUB_OFFLINE=1]`.
- **GPU:** возьмите cuda-тег того же релиза (например,
`ghcr.io/huggingface/text-embeddings-inference:cuda-1.9`) и запустите контейнер с `gpus: all`.
## Возможности
- Совместная работа в реальном времени
@@ -1418,5 +1418,14 @@
"The commented text changed since this suggestion was made; it was not applied.": "The commented text changed since this suggestion was made; it was not applied.",
"Dismiss": "Dismiss",
"Suggestion dismissed": "Suggestion dismissed",
"Failed to dismiss suggestion": "Failed to dismiss suggestion"
"Failed to dismiss suggestion": "Failed to dismiss suggestion",
"Save version": "Save version",
"Ctrl+S": "Ctrl+S",
"Version saved": "Version saved",
"Already saved as the latest version": "Already saved as the latest version",
"Agent version": "Agent version",
"Boundary": "Boundary",
"Autosave": "Autosave",
"Only versions": "Only versions",
"No saved versions yet.": "No saved versions yet."
}
@@ -1433,5 +1433,14 @@
"The commented text changed since this suggestion was made; it was not applied.": "Прокомментированный текст изменился после создания предложения; оно не было применено.",
"Dismiss": "Не применять",
"Suggestion dismissed": "Предложение отклонено",
"Failed to dismiss suggestion": "Не удалось отклонить предложение"
"Failed to dismiss suggestion": "Не удалось отклонить предложение",
"Save version": "Сохранить версию",
"Ctrl+S": "Ctrl+S",
"Version saved": "Версия сохранена",
"Already saved as the latest version": "Уже сохранено как последняя версия",
"Agent version": "Версия агента",
"Boundary": "Граница",
"Autosave": "Автосейв",
"Only versions": "Только версии",
"No saved versions yet.": "Пока нет сохранённых версий."
}
@@ -58,8 +58,11 @@ import ConversationList from "@/features/ai-chat/components/conversation-list.ts
import ChatThread from "@/features/ai-chat/components/chat-thread.tsx";
import {
exportAiChat,
getAiChatMessagesDelta,
stopRun,
} from "@/features/ai-chat/services/ai-chat-service.ts";
import { mergeDeltaRowsIntoPages } from "@/features/ai-chat/utils/resume-helpers.ts";
import type { IAiChatMessageRow } from "@/features/ai-chat/types/ai-chat.types.ts";
import { useChatSession } from "@/features/ai-chat/hooks/use-chat-session.ts";
import {
shouldCollapseOnOutsidePointer,
@@ -269,17 +272,64 @@ export default function AiChatWindow() {
const { data: messageRows, isLoading: messagesLoading } =
useAiChatMessagesQuery(
activeChatId ?? undefined,
// DELIBERATELY DUMB: poll every 2.5s WHILE ARMED, otherwise off. NO error
// checks (TanStack resets fetchFailureCount each fetch; the poll must survive
// a server restart), NO tail checks, NO cap here — the settled/stalled/idle-cap
// semantics all live in ChatThread's FSM, which disarms via onResumeFallback.
() => (degradedPoll === true ? 2500 : false),
// #344: gate on windowOpen too — no message history is fetched (and no
// degraded poll runs) while the window is closed; it loads when the window
// opens with an active chat.
// #491: the full infinite-query no longer POLLS. It seeds the thread ONCE; the
// degraded fallback now runs a DELTA poller (below) that augments THIS cache
// idempotently, instead of refetching every page (with full parts) every 2.5s.
false,
// #344: gate on windowOpen too — no message history is fetched while the window
// is closed; it loads when the window opens with an active chat.
windowOpen,
);
// #491 degraded DELTA poll. While armed (degradedPoll) and the window is open on a
// chat, poll POST /ai-chat/messages/delta every 2.5s: it returns only the rows
// CHANGED since the previous cursor (+ the run fact) in ONE round-trip. We merge
// those rows into the SAME infinite-query cache the thread reads (idempotently by
// id — the delta's overlap window re-delivers rows), so the thread's reconcile
// effect follows the detached run to its terminal row from a fraction of the wire
// cost. The run-fact settle stays the thread FSM's job (row-status reconcile), so
// we do NOT double-poll /run here. Cursor resets when the chat changes / disarms.
const deltaCursorRef = useRef<string | undefined>(undefined);
useEffect(() => {
deltaCursorRef.current = undefined;
}, [activeChatId, degradedPoll]);
useEffect(() => {
if (!degradedPoll || !windowOpen || !activeChatId) return;
const chatId = activeChatId;
let cancelled = false;
const tick = async (): Promise<void> => {
try {
const res = await getAiChatMessagesDelta(chatId, deltaCursorRef.current);
if (cancelled) return;
deltaCursorRef.current = res.cursor;
if (res.rows.length > 0) {
queryClient.setQueryData(
AI_CHAT_MESSAGES_RQ_KEY(chatId),
(
old:
| {
pages: { items: IAiChatMessageRow[]; meta: unknown }[];
pageParams: unknown[];
}
| undefined,
) =>
old
? { ...old, pages: mergeDeltaRowsIntoPages(old.pages, res.rows) }
: old,
);
}
} catch {
// Transient failure (e.g. a server restart mid-run): swallow and retry on
// the next tick — the poll must survive a bounce, like the old dumb refetch.
}
};
const id = setInterval(() => void tick(), 2500);
return () => {
cancelled = true;
clearInterval(id);
};
}, [degradedPoll, windowOpen, activeChatId, queryClient]);
// #184 reconnect-and-live-follow. Whether detached agent runs are enabled for
// this workspace. When the feature is off no runs are ever created, so the
// resume attempt would only ever 204; gating ChatThread's resume on it avoids a
@@ -172,9 +172,18 @@ function resetState() {
h.state.getRun.mockResolvedValue({ run: null, message: null });
}
// #491: the streaming tail carries a persisted step frontier (metadata.stepsPersisted),
// which the tail-only attach reads as `n` in `?anchor=<id>&n=<n>`. Seeded WHOLE now.
const streamingTail = () => [
row("u1", "user", undefined, "hi"),
row("a1", "assistant", "streaming", "partial"),
{
id: "a1",
role: "assistant",
content: "partial",
status: "streaming",
createdAt: "2026-01-01T00:00:00Z",
metadata: { stepsPersisted: 2 },
} as IAiChatMessageRow,
];
const settledTail = () => [
row("u1", "user", undefined, "hi"),
@@ -335,20 +344,24 @@ describe("ChatThread — send now", () => {
expect(screen.getAllByLabelText("Remove queued message")).toHaveLength(1);
});
it("Stop then a REAL network-drop finish exits to idle (honor-in-stopping), NOT a false reconnect", () => {
it("Stop then a REAL network-drop finish exits to idle (honor-in-stopping), NOT a false reconnect", async () => {
// Regression for the disconnect-first reorder: on the STOP path, even a drop-
// form finish { isError:true, isDisconnect:true } arriving in `stopping` must be
// HONORED (reducer) and exit to idle — it must NOT enter the reconnect ladder.
startLocalStreamWithRun(); // live local stream, autonomous
fireEvent.click(screen.getByLabelText("Stop")); // STOP_REQUESTED -> stopping
h.state.error = { message: "Failed to fetch" };
act(() => {
// #491: the disconnect re-seeds from persist (async getRun) before dispatching
// FINISH_DISCONNECT, which the reducer HONORS in `stopping` -> idle. Flush it.
await act(async () => {
h.state.onFinish?.({
message: { id: "a1", role: "assistant", parts: [] },
isAbort: false,
isDisconnect: true,
isError: true,
});
await Promise.resolve();
await Promise.resolve();
});
expect(screen.queryByText(/reconnecting/i)).toBeNull();
});
@@ -803,19 +816,24 @@ describe("ChatThread — resume (attach) machinery", () => {
expect(h.state.resumeStream).not.toHaveBeenCalled();
});
it("strips the streaming tail from the seed, keeps a user tail whole", () => {
it("#491 tail-only: seeds the streaming tail WHOLE (no strip), keeps a user tail whole", () => {
renderThread({ autonomousRunsEnabled: true, initialRows: streamingTail() });
expect(h.state.seededMessages).toHaveLength(1);
// MUTATION-VERIFY: re-introduce the seed-strip and this goes red — the streaming
// tail (steps 0..N-1) MUST be seeded so the SDK continuation appends the tail to
// the RIGHT message. Both rows (user + assistant) are seeded.
expect(h.state.seededMessages).toHaveLength(2);
cleanup();
resetState();
renderThread({ autonomousRunsEnabled: true, initialRows: userTail() });
expect(h.state.seededMessages).toHaveLength(1);
});
it("builds the attach URL with expect=live&anchor only for a stripped streaming tail", () => {
it("#491 tail-only: builds the attach URL with ?anchor=&n= from the persisted step frontier", () => {
renderThread({ autonomousRunsEnabled: true, initialRows: streamingTail() });
// n=2 comes from a1's metadata.stepsPersisted (MUTATION-VERIFY: hardcode n=0 and
// this fails). No `expect=live` param anymore.
expect(h.state.transport!.prepareReconnectToStreamRequest!().api).toBe(
"/api/ai-chat/runs/c1/stream?expect=live&anchor=a1",
"/api/ai-chat/runs/c1/stream?anchor=a1&n=2",
);
cleanup();
resetState();
@@ -839,39 +857,41 @@ describe("ChatThread — resume (attach) machinery", () => {
});
}
it("204 on a streaming tail: restore + invalidate + onResumeFallback(true)", async () => {
it("204 on a streaming tail: NO restore (row kept) + invalidate + onResumeFallback(true)", async () => {
const { onResumeFallback, invalidateSpy } = renderThread({
autonomousRunsEnabled: true,
initialRows: streamingTail(),
});
await attachFetch({ status: 204, ok: false });
expect(h.state.setMessages).toHaveBeenCalledTimes(1); // restore
// #491 tail-only: the anchor row was never stripped, so there is NOTHING to
// restore. MUTATION-VERIFY: re-add a restore setMessages here and it goes red.
expect(h.state.setMessages).not.toHaveBeenCalled();
expect(invalidateSpy).toHaveBeenCalledWith({
queryKey: ["ai-chat-messages", "c1"],
});
expect(onResumeFallback).toHaveBeenCalledWith(true);
});
it("F7 restart-survival: a 500 attach failure restores the row AND arms the poll", async () => {
it("F7 restart-survival: a 500 attach failure arms the poll WITHOUT a restore", async () => {
const { onResumeFallback, invalidateSpy } = renderThread({
autonomousRunsEnabled: true,
initialRows: streamingTail(),
});
await attachFetch({ status: 500, ok: false });
expect(h.state.setMessages).toHaveBeenCalledTimes(1);
expect(h.state.setMessages).not.toHaveBeenCalled();
expect(invalidateSpy).toHaveBeenCalledWith({
queryKey: ["ai-chat-messages", "c1"],
});
expect(onResumeFallback).toHaveBeenCalledWith(true);
});
it("F7 restart-survival: a network throw restores the row AND arms the poll", async () => {
it("F7 restart-survival: a network throw arms the poll WITHOUT a restore", async () => {
const { onResumeFallback, invalidateSpy } = renderThread({
autonomousRunsEnabled: true,
initialRows: streamingTail(),
});
await attachFetch(new Error("network down"), true);
expect(h.state.setMessages).toHaveBeenCalledTimes(1);
expect(h.state.setMessages).not.toHaveBeenCalled();
expect(invalidateSpy).toHaveBeenCalledWith({
queryKey: ["ai-chat-messages", "c1"],
});
@@ -931,7 +951,7 @@ describe("ChatThread — resume (attach) machinery", () => {
expect(h.state.sendMessage).not.toHaveBeenCalled();
});
it("an empty resumed message (starved replay) restores the row AND arms the poll", () => {
it("an empty resumed message (starved replay) arms the poll WITHOUT a restore", () => {
h.state.status = "ready";
const { onResumeFallback } = renderThread({
autonomousRunsEnabled: true,
@@ -947,7 +967,9 @@ describe("ChatThread — resume (attach) machinery", () => {
isError: false,
});
});
expect(h.state.setMessages).toHaveBeenCalledTimes(1); // restore
// #491 tail-only: the seeded steps 0..N-1 are still on screen (the SDK
// continuation never wiped them), so there is nothing to restore — just poll.
expect(h.state.setMessages).not.toHaveBeenCalled();
expect(onResumeFallback).toHaveBeenCalledWith(true); // arm
});
@@ -995,24 +1017,41 @@ describe("ChatThread — live reconnect + stalled", () => {
cleanup();
});
// #491: the authoritative PERSISTED assistant row `getRun` projects on a local
// disconnect — the re-seed source. Its metadata.stepsPersisted becomes `n`.
const persistedAnchor = (steps = 3) => ({
run: { id: "run-1", status: "running" },
message: {
id: "a2",
role: "assistant",
content: "persisted 0..N-1",
status: "streaming",
createdAt: "2026-01-01T00:00:00Z",
metadata: { stepsPersisted: steps },
},
});
// A REAL live SSE drop. ai@6.0.207 emits BOTH { isError:true, isDisconnect:true }
// for a network TypeError AND sets useChat `error` — NOT the { isError:false,
// error:null } form the old tests fed. This is the form browser QA hit; with the
// buggy isError-first routing OR without the errorView render-gate these tests go
// red (a real drop surfaces the terminal error banner, masking the reconnect
// ladder). MUTATION-VERIFY of disconnect-first + the errorView phase-gate.
function disconnect(message: unknown = liveMsg) {
// for a network TypeError AND sets useChat `error`. #491: an autonomous local drop
// now RE-SEEDS from persist (async getRun) BEFORE entering the reconnect ladder, so
// this helper is async and flushes the getRun microtask before returning.
async function disconnect(message: unknown = liveMsg) {
h.state.error = { message: "Failed to fetch" }; // the SDK sets error on the drop
act(() => {
await act(async () => {
h.state.onFinish?.({
message,
isAbort: false,
isDisconnect: true,
isError: true,
});
// Flush the getRun().then re-seed + the deferred FINISH_DISCONNECT dispatch.
await Promise.resolve();
await Promise.resolve();
});
}
function renderLive() {
// The persisted-anchor read the local disconnect performs to re-seed from persist.
h.state.getRun.mockResolvedValue(persistedAnchor());
const view = renderThread({
autonomousRunsEnabled: true,
initialRows: settledTail(),
@@ -1032,35 +1071,80 @@ describe("ChatThread — live reconnect + stalled", () => {
});
}
it("a live disconnect starts a backoff reconnect (banner + resumeStream after backoff)", () => {
it("#491: a live disconnect RE-SEEDS from persist, then backs off to reconnect with ?anchor=&n=", async () => {
renderLive();
disconnect();
await disconnect();
// The re-seed read the authoritative persisted row and replaced the live partial.
// MUTATION-VERIFY: skip the getRun re-seed (send `n` off the live message) and the
// n below no longer matches the PERSISTED stepsPersisted.
expect(h.state.getRun).toHaveBeenCalledWith("c1");
expect(h.state.setMessages).toHaveBeenCalled(); // re-seeded the store from persist
expect(screen.getByText(/reconnecting/i)).toBeTruthy();
expect(h.state.resumeStream).not.toHaveBeenCalled();
advanceToAttempt(1);
expect(h.state.resumeStream).toHaveBeenCalledTimes(1);
// n=3 is the PERSISTED row's stepsPersisted (from getRun), NOT the live store.
expect(h.state.transport!.prepareReconnectToStreamRequest!().api).toBe(
"/api/ai-chat/runs/c1/stream?expect=live&anchor=a2",
"/api/ai-chat/runs/c1/stream?anchor=a2&n=3",
);
});
it("#488 (browser QA): the reconnect banner is SHOWN, not masked by the residual useChat error", () => {
it("#491 regression (#137/#161 dup): getRun REJECT on a live disconnect drops the live partial + nulls the anchor", async () => {
// The re-seed source (getRun) FAILS — a flaky-network blip (SSE + getRun both
// fail, network recovers in ~1s). The OLD .catch just re-entered the ladder with
// NO re-seed and NO filter, so the reconnect could tail-apply the registry's
// frames onto the live partial that ALREADY has those steps -> duplicated text.
renderLive();
h.state.getRun.mockReset();
h.state.getRun.mockRejectedValue(new Error("network"));
await disconnect(); // live partial = liveMsg (id "a2")
expect(h.state.getRun).toHaveBeenCalledWith("c1");
// THE GUARANTEE: on the getRun failure the live partial (a2) is FILTERED from the
// store, so the reconnect can never tail-apply already-present steps onto it.
// MUTATION-VERIFY: revert the .catch fix (enterReconnect only, no filter) and no
// setMessages call removes a2 -> this reddens.
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);
expect(screen.getByText(/reconnecting/i)).toBeTruthy();
advanceToAttempt(1);
expect(h.state.resumeStream).toHaveBeenCalledTimes(1);
// Anchor was nulled -> replay-from-start (no params) / 204 -> poll; never a stale
// ?anchor=&n= over the live partial.
expect(h.state.transport!.prepareReconnectToStreamRequest!().api).toBe(
"/api/ai-chat/runs/c1/stream",
);
});
it("#488 (browser QA): the reconnect banner is SHOWN, not masked by the residual useChat error", async () => {
// The drop sets useChat `error` (real SDK), and the terminal errorView describes
// it ("Lost connection to the server"). The FSM phase-gate must let the
// `reconnecting` banner WIN over that residual error. MUTATION-VERIFY: revert the
// errorView phase-gate (show errorView whenever error is set) and the terminal
// banner masks "reconnecting…" -> red.
renderLive();
disconnect();
await disconnect();
expect(h.state.error).not.toBeNull(); // the SDK error IS set during recovery
expect(screen.getByText(/reconnecting/i)).toBeTruthy();
// The terminal "Lost connection… reload" banner must NOT be showing.
expect(screen.queryByText(/reload and try again/i)).toBeNull();
});
it("#488 commit 2: a disconnect BEFORE the first assistant frame reconnects with NO anchor", () => {
it("#488 commit 2: a disconnect BEFORE the first assistant frame reconnects with NO anchor", async () => {
renderLive();
disconnect(null); // no assistant message yet (pre-first-frame break)
// No persisted assistant row for a pre-first-frame break -> no anchor.
h.state.getRun.mockResolvedValue({ run: null, message: null });
await disconnect(null); // no assistant message yet (pre-first-frame break)
expect(screen.getByText(/reconnecting/i)).toBeTruthy();
expect(
screen.queryByText("Connection lost — the answer was interrupted."),
@@ -1074,7 +1158,7 @@ describe("ChatThread — live reconnect + stalled", () => {
it("a live re-attach (2xx) clears the reconnect banner", async () => {
renderLive();
disconnect();
await disconnect();
advanceToAttempt(1);
await reconnect({ status: 200, ok: true });
expect(screen.queryByText(/reconnecting/i)).toBeNull();
@@ -1082,7 +1166,7 @@ describe("ChatThread — live reconnect + stalled", () => {
it("a 204 arms the degraded poll and backs off to the next attempt", async () => {
const { onResumeFallback } = renderLive();
disconnect();
await disconnect();
advanceToAttempt(1);
expect(h.state.resumeStream).toHaveBeenCalledTimes(1);
await reconnect({ status: 204, ok: false });
@@ -1094,7 +1178,7 @@ describe("ChatThread — live reconnect + stalled", () => {
it("exhausts the attempt limit into a manual Retry, which restarts the sequence", async () => {
renderLive();
disconnect();
await disconnect();
for (let n = 1; n <= 5; n++) {
advanceToAttempt(n);
expect(h.state.resumeStream).toHaveBeenCalledTimes(n);
@@ -1112,22 +1196,23 @@ describe("ChatThread — live reconnect + stalled", () => {
it("#488 commit 3: two breaks in a row produce two reconnect cycles", async () => {
renderLive();
// First break -> reconnect -> re-attach live.
disconnect();
await disconnect();
advanceToAttempt(1);
expect(h.state.resumeStream).toHaveBeenCalledTimes(1);
await reconnect({ status: 200, ok: true });
expect(screen.queryByText(/reconnecting/i)).toBeNull();
// The re-attached observer stream drops AGAIN -> a SECOND reconnect cycle
// (the old one-shot !wasResumed gate sent this to silent poll).
disconnect();
// The re-attached observer (live-follow) stream drops AGAIN -> a SECOND reconnect
// cycle. #491: this too re-seeds from persist before re-attaching (never tail-
// applies over the live-follow partial).
await disconnect();
expect(screen.getByText(/reconnecting/i)).toBeTruthy();
advanceToAttempt(1);
expect(h.state.resumeStream).toHaveBeenCalledTimes(2);
});
it("does NOT reconnect when autonomous runs are disabled", () => {
it("does NOT reconnect when autonomous runs are disabled", async () => {
renderThread({ autonomousRunsEnabled: false, initialRows: settledTail() });
disconnect();
await disconnect();
expect(screen.queryByText(/reconnecting/i)).toBeNull();
expect(
screen.getByText("Connection lost — the answer was interrupted."),
@@ -1138,7 +1223,7 @@ describe("ChatThread — live reconnect + stalled", () => {
it("#488 commit 4a: the poll idle cap surfaces a stalled banner + Retry (not silent)", async () => {
renderLive();
disconnect();
await disconnect();
advanceToAttempt(1);
await reconnect({ status: 204, ok: false }); // arms the poll (reconnecting)
// No activity for the whole idle cap -> stalled.
@@ -42,7 +42,7 @@ import { assistantMessageHasVisibleContent } from "@/features/ai-chat/utils/mess
import {
isStreamingTail,
isSettledAssistantTail,
seedRows,
stepsPersistedOf,
mergeById,
} from "@/features/ai-chat/utils/resume-helpers.ts";
import { getRun } from "@/features/ai-chat/services/ai-chat-service.ts";
@@ -266,25 +266,33 @@ export default function ChatThread({
// is NOT one of the lifecycle flags the FSM replaced.
const mountedRef = useRef(true);
// attachStrategy DATA (behind the resumeStream effect; #491 swaps it to tail-only
// WITHOUT touching the FSM). The controller is effect-owned (aborted in cleanup,
// I5). `stripRef`/`strippedRowRef` are the current full-replay+strip anchor.
// attachStrategy DATA (behind the resumeStream effect; #491 tail-only, WITHOUT
// touching the FSM). The controller is effect-owned (aborted in cleanup, I5).
// `anchorRef` is the PERSISTED assistant row that pins the run (server invariant
// 6) and its persisted step frontier N: it feeds `?anchor=<id>&n=<stepsPersisted>`
// so the tail-only attach returns frames for steps >= N (the seed carries 0..N-1).
// It is NOT a "stripped" row — the seed keeps every row (tail-only replaces the
// old full-replay+strip). Null when there is no streaming/active tail to resume.
const attachAbortRef = useRef<AbortController | null>(null);
const stripRef = useRef(chatId !== null && isStreamingTail(initialRows ?? []));
const strippedRowRef = useRef<IAiChatMessageRow | null>(
stripRef.current ? (initialRows ?? [])[initialRows!.length - 1] : null,
const anchorRef = useRef<{ id: string; stepsPersisted: number } | null>(
(() => {
if (chatId === null || !isStreamingTail(initialRows ?? [])) return null;
const rows = initialRows ?? [];
const tail = rows[rows.length - 1];
return { id: tail.id, stepsPersisted: stepsPersistedOf(tail) };
})(),
);
// Effect-owned backoff timers (not lifecycle flags): the reconnect ladder and the
// 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);
// #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
// to it by the SDK continuation (readUIMessageStream({ message })), so it must be
// present in the store for the attach to continue the RIGHT message.
const initialMessages = useMemo<UIMessage[]>(
() =>
seedRows(
initialRows ?? [],
stripRef.current && autonomousRunsEnabled === true,
).map(rowToUiMessage),
() => (initialRows ?? []).map(rowToUiMessage),
[initialRows],
);
@@ -335,21 +343,16 @@ export default function ChatThread({
(eff: RunEffect, epoch: number) => {
switch (eff.type) {
case "resumeStream": {
// The attach GET. Stamp the outcome's generation (I1). A reconnect
// attempt filters the pinned live row from the store first (the mount
// seed already stripped it), so the live replay's text-start rebuilds it
// without duplicating parts (#430).
// The attach GET. Stamp the outcome's generation (I1). #491 tail-only: the
// store already holds EXACTLY the persisted steps 0..N-1 (the mount seed IS
// persist; a reconnect was re-seeded from persist BEFORE FINISH_DISCONNECT
// scheduled it — see the onFinish disconnect handler), so there is nothing
// to filter here: the SDK continues that seeded message, appending the tail
// (steps >= N) without duplicating the pre-drop partial step.
pendingAttachEpochRef.current = epoch;
// The resumed stream's onFinish is stamped with THIS attach generation
// (F1), so a superseded attempt's late finish is dropped.
turnEpochRef.current = epoch;
if (machineRef.current.phase.name === "reconnecting") {
const anchor = strippedRowRef.current;
if (anchor)
setMessagesRef.current?.((prev) =>
prev.filter((m) => m.id !== anchor.id),
);
}
void resumeStreamRef.current?.();
break;
}
@@ -464,18 +467,23 @@ export default function ChatThread({
new DefaultChatTransport<UIMessage>({
api: "/api/ai-chat/stream",
credentials: "include",
prepareReconnectToStreamRequest: () => ({
// Build the attach URL from the REAL chat id. ?expect=live&anchor=<row id>
// only when a streaming tail was stripped: expect=live opts into a
// finished-retained replay (safe only because the row is stripped and the
// replay rebuilds it), and the anchor pins the replay to OUR run — a
// mismatching (newer) run 204s into the restore+poll path instead.
api: `/api/ai-chat/runs/${chatIdRef.current}/stream${
stripRef.current
? `?expect=live&anchor=${strippedRowRef.current!.id}`
: ""
}`,
}),
prepareReconnectToStreamRequest: () => {
// #491 tail-only attach URL. When there is an anchor (a streaming/active
// tail to resume) build `?anchor=<assistantRowId>&n=<stepsPersisted>`: the
// server returns the TAIL — a synthetic `start` frame + frames for steps
// >= n, then live — which the SDK continuation appends to the seeded row.
// The server 204s (-> restore-noop + poll) when it cannot cover the
// frontier (overflow/rotation gap) or the anchor mismatches (a newer run).
// No anchor (a user tail / pre-first-frame break) => no params.
const anchor = anchorRef.current;
return {
api: `/api/ai-chat/runs/${chatIdRef.current}/stream${
anchor
? `?anchor=${anchor.id}&n=${anchor.stepsPersisted}`
: ""
}`,
};
},
fetch: async (input: RequestInfo | URL, init: RequestInit = {}) => {
if ((init.method ?? "GET") !== "GET") {
// Send path (POST). #488 commit 5: NO client 409 retry ladder anymore
@@ -562,8 +570,9 @@ export default function ChatThread({
// Attach GET outcome -> FSM event. The epoch guard replaces BOTH the one-shot
// 204 guard (noStreamHandledRef) and the unmount gate: a stale/superseded or
// post-DISPOSE outcome is dropped (I1). For a NONE outcome the attachStrategy
// recovery (restore the stripped row + invalidate for a fresh poll) runs first.
// post-DISPOSE outcome is dropped (I1). #491 tail-only: on a NONE outcome there is
// NOTHING to restore the anchor row was never stripped from the view (the seed
// keeps it) — so we only invalidate for a fresh poll + dispatch the FSM event.
const handleAttachOutcome = useCallback(
(ep: number, wasReconnecting: boolean, live: boolean) => {
if (ep !== epochRef.current) return; // stale generation — drop
@@ -575,10 +584,6 @@ export default function ChatThread({
);
return;
}
if (strippedRowRef.current)
setMessagesRef.current?.((prev) =>
mergeById(prev, rowToUiMessage(strippedRowRef.current!)),
);
queryClient.invalidateQueries({
queryKey: AI_CHAT_MESSAGES_RQ_KEY(chatIdRef.current),
});
@@ -661,56 +666,31 @@ export default function ChatThread({
// keeps executing server-side — must win; only a NON-disconnect error (a
// provider 500, `{ isError:true, isDisconnect:false }`) is terminal.
if (isDisconnect) {
if (wasObserver) {
// A resumed/attached OBSERVER stream dropped. Recover via the degraded
// poll (restore the stripped row only when there is no visible content;
// never clobber a fuller on-screen tail, invariant 9). The FSM decides
// reconnect-vs-poll from liveFollow (a live-follow drop reconnects again,
// #488 commit 3; a mount-resume drop polls).
if (mountedRef.current) {
const hasVisible = msgHasVisible;
if (!hasVisible && strippedRowRef.current)
setMessages((prev) =>
mergeById(prev, rowToUiMessage(strippedRowRef.current!)),
);
queryClient.invalidateQueries({
queryKey: AI_CHAT_MESSAGES_RQ_KEY(chatIdRef.current),
});
dispatch({
type: "FINISH_DISCONNECT",
hasVisibleContent: hasVisible,
epoch: stampEpoch,
});
}
if (!mountedRef.current) {
setStopNotice(null);
return;
}
// A LOCAL live turn dropped. #488 commit 2: recover by the RUN-FACT, not by
// the presence of an assistant message — a setup-phase break (before the
// first frame) still leaves a detached run writing to pages. In autonomous
// mode a run is active for the whole turn, so seed the run-fact from the
// start-metadata runId when known, else a sentinel (the attach GET goes by
// chatId, not runId). Pin the assistant row as the strip/anchor when present.
if (autonomousRunsEnabled === true && mountedRef.current) {
const hasAnchor =
message?.role === "assistant" && typeof message.id === "string";
if (hasAnchor) {
strippedRowRef.current = {
id: message.id,
role: "assistant",
content: "",
status: "streaming",
createdAt: new Date().toISOString(),
metadata: { parts: message.parts },
};
stripRef.current = true;
} else {
strippedRowRef.current = null;
stripRef.current = false;
}
// No detached run to recover (legacy, non-autonomous): a plain disconnect —
// terminal notice, no reconnect. (An observer only exists in autonomous mode,
// so this is always a local turn.)
if (autonomousRunsEnabled !== true) {
dispatch({
type: "RUN_FACT",
runFact: { runId: extractRunId(message) ?? "pending" },
type: "FINISH_DISCONNECT",
hasVisibleContent: false,
epoch: stampEpoch,
});
setStopNotice("disconnect");
return;
}
// A mount-resume OBSERVER (one-shot resume, NOT live-follow) drop falls to
// the degraded POLL, which merges by id — it does NOT attach, so there is
// nothing to re-seed. #491 tail-only: the anchor row was never removed from
// the view (the seed keeps it; the continuation only APPENDED), so nothing to
// restore either. The FSM routes this to `polling` (ownership observer,
// !liveFollow).
if (wasObserver && !machineRef.current.ctx.liveFollow) {
queryClient.invalidateQueries({
queryKey: AI_CHAT_MESSAGES_RQ_KEY(chatIdRef.current),
});
dispatch({
type: "FINISH_DISCONNECT",
@@ -718,14 +698,92 @@ export default function ChatThread({
epoch: stampEpoch,
});
setStopNotice(null);
} else {
return;
}
// We will (re-)ENTER THE RECONNECT LADDER (an attach): a LOCAL live turn's
// first drop, OR a live-follow observer's SUBSEQUENT drop (#488 commit 3).
// #488 commit 2: recover by the RUN-FACT, not by the presence of an assistant
// message — a setup-phase break still leaves a detached run writing to pages.
//
// #491 tail-only (THE crux): the live store holds a PARTIAL step that is AHEAD
// of the persisted boundary; tail-applying the reconnect's step frames over it
// would DUPLICATE that partial step. So entering reconnecting is ALWAYS via a
// RE-SEED FROM PERSIST — never the live store. Fetch the authoritative
// persisted assistant row (`getRun` returns the projected `message`), replace
// the live partial by id (mergeById -> the store now holds EXACTLY steps
// 0..N-1), and set the anchor to `{ id, n = stepsPersisted }`. Only AFTER the
// re-seed is applied do we enter the ladder (FINISH_DISCONNECT schedules the
// backoff) — so the attach can never tail-apply over the live partial.
const cid = chatIdRef.current;
// The live-message runId is the run-fact source (the attach GET keys on
// chatId, so a sentinel still recovers a setup-phase break).
const runId = extractRunId(message ?? undefined) ?? "pending";
const enterReconnect = (fact: string): void => {
if (!mountedRef.current) return;
// Epoch-stamp the run-fact too (I1): the getRun rtt widens the
// onFinish->dispatch window, so a concurrent SEND_LOCAL during it must be
// able to drop this stale RUN_FACT (else it clobbers the new turn's
// runFact.runId). Consistent with the postRun RUN_FACT stamp.
dispatch({ type: "RUN_FACT", runFact: { runId: fact }, epoch: stampEpoch });
dispatch({
type: "FINISH_DISCONNECT",
hasVisibleContent: false,
hasVisibleContent: msgHasVisible,
epoch: stampEpoch,
});
setStopNotice("disconnect");
};
// Restore the STRUCTURAL guarantee that the live partial is never the
// tail-apply base: drop the live partial from the store by id and null the
// anchor, so the reconnect replays from step 0 into a CLEAN store (a full
// rebuild) or, past any rotation, 204s -> degraded poll. Used on BOTH the
// no-persisted-row and getRun-FAILURE paths — after this there is no path
// where the attach tail-applies frames onto a row that already has them
// (the #137/#161 duplication class).
const dropLivePartialAndReplayFromStart = (): void => {
if (message?.role === "assistant" && typeof message.id === "string") {
const liveId = message.id;
setMessagesRef.current?.((prev) =>
prev.filter((m) => m.id !== liveId),
);
}
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.
dropLivePartialAndReplayFromStart();
enterReconnect(runId);
});
} else {
dropLivePartialAndReplayFromStart();
enterReconnect(runId);
}
setStopNotice(null);
return;
}
// A NON-disconnect stream error (a provider 500 etc.) -> terminal error banner.
@@ -746,11 +804,10 @@ export default function ChatThread({
if (mountedRef.current) {
const hasVisible = msgHasVisible;
if (!hasVisible) {
// Starved replay: restore the stripped row + poll to the real terminal.
if (strippedRowRef.current)
setMessages((prev) =>
mergeById(prev, rowToUiMessage(strippedRowRef.current!)),
);
// Starved replay (the tail carried no new steps). #491 tail-only: the
// seeded steps 0..N-1 are still on screen (the SDK continuation never
// wiped them — `start` does not reset parts), so there is nothing to
// restore; just poll to the real terminal.
queryClient.invalidateQueries({
queryKey: AI_CHAT_MESSAGES_RQ_KEY(chatIdRef.current),
});
@@ -863,12 +920,12 @@ export default function ChatThread({
const tail = rows[rows.length - 1];
if (!tail || tail.role !== "assistant") return;
setMessages((prev) => mergeById(prev, rowToUiMessage(tail)));
// Anchor-mismatch coherence: a restored stripped row A that a DIFFERENT run's
// row B has replaced as the tail would linger as an orphan — settle A from
// fresh history so no phantom row survives.
const stripped = strippedRowRef.current;
if (stripped && stripped.id !== tail.id) {
const historical = rows.find((r) => r.id === stripped.id);
// Anchor-mismatch coherence: if a DIFFERENT run's row B has replaced our anchor
// row A as the tail, A would linger as an orphan — reconcile A by id from FRESH
// PERSISTED history (not the pinned live row) so no phantom row survives.
const anchor = anchorRef.current;
if (anchor && anchor.id !== tail.id) {
const historical = rows.find((r) => r.id === anchor.id);
if (historical)
setMessages((prev) => mergeById(prev, rowToUiMessage(historical)));
}
@@ -27,6 +27,7 @@ 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.
@@ -114,3 +115,89 @@ 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.
});
});
@@ -0,0 +1,112 @@
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);
});
});
@@ -4,6 +4,7 @@ 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";
@@ -86,17 +87,39 @@ 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 marked + DOMPurify pass. Props are
* primitives, so React.memo's default shallow compare is exactly right (the
* `text` string is compared by value).
* 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 streamingdone flip because
* `messageSignature` tracks each part's `state` (and `turnStreaming` flips at
* turn end), so the incremental view always converges to this single render.
*/
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 (
@@ -179,47 +202,10 @@ function MessageItem({
{resolveAssistantName(assistantName) ?? t("AI agent")}
</Text>
{message.parts.map((part, index) => {
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}
/>
);
}
// 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 (isToolPart(part.type)) {
return (
<ToolCallCard
@@ -232,7 +218,76 @@ function MessageItem({
);
}
return null;
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;
}
}
})}
{/* A persisted turn error (server stored it in metadata.error). Rendered
here so it survives a thread remount and shows in reopened history. */}
@@ -0,0 +1,96 @@
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>
),
)}
</>
);
}
@@ -57,6 +57,31 @@ export async function stopRun(
return req.data;
}
/**
* Delta poll (#491): the chat's message rows changed since `cursor` (a DB-clock
* timestamp echoed from the previous poll) plus the current run fact, in ONE
* round-trip the degraded-poll fallback's payload, replacing the old "refetch
* ALL infinite-query pages every 2.5s with full parts" poll. Omit `cursor` on the
* first poll (returns just a fresh cursor, no rows, to start the chain). The
* overlap window guarantees occasional REPEATS, so the caller MUST merge rows
* idempotently by id (mergeById). Owner-gated server-side.
*/
export async function getAiChatMessagesDelta(
chatId: string,
cursor?: string,
): Promise<{
rows: IAiChatMessageRow[];
cursor: string;
run: { id: string; status: string } | null;
}> {
const req = await api.post<{
rows: IAiChatMessageRow[];
cursor: string;
run: { id: string; status: string } | null;
}>("/ai-chat/messages/delta", { chatId, cursor });
return req.data;
}
/**
* #488: the run-fact "is a run active on this chat?" first-class from the
* server (POST /ai-chat/run). Called on mount to seed the client FSM's run-fact
@@ -48,6 +48,7 @@ Legend: **†** = command-transition (bumps `epoch`, I1). Effects in `[…]`.
| `RETRY` (manual, stalled banner) | stalled | polling(attach-none) **†** | `[armPoll]` |
| `POLL_TERMINAL` (settled tail merged) | polling, reconnecting, stopping | idle | `[disarmPoll, cancelReconnect]`, runFact←null (I4) |
| `POLL_IDLE_CAP` (inactivity cap) | polling, reconnecting | stalled | `[disarmPoll, cancelReconnect]` (commit 4a — no more silent) |
| `POLL_IDLE_CAP` (inactivity cap) | stopping | idle | `[disarmPoll, cancelReconnect]`, runFact←null (Review #4: a Stop-armed poll with no SDK/terminal backstop gets a bounded exit — NOT `stalled`, Stop was already pressed so nothing to retry) |
| `RUN_FACT{null}` (POST /run → null/terminal, 204) | reconnecting/attaching/polling/stopping | idle | `[cancelReconnect, disarmPoll]`, runFact←null (I3 fresh-negative gate) |
| `RUN_FACT{runId}` | any | (same) | runFact←runId (pessimism toward an attempt) |
| `STOP_REQUESTED` (user Stop) | streaming, reconnecting, polling | stopping **†** | `[stopRun, abortAttach, cancelReconnect, armPoll]` (poll drives the terminal — I4 exit by data) |
@@ -121,8 +122,7 @@ holds. **Pending column: empty.**
| 11 | `stopPendingRef` | **FSM phase `stopping`** | the deferred stop fires from the chat-id adoption effect while `stopping` |
| 12 | `mountedRef` | **retained (React liveness)** | orthogonal to run-lifecycle; gates imperative onFinish side-effects post-unmount. Epoch (I1) handles stale COMMAND-outcomes; DISPOSE bumps it |
| 13 | `attemptResumeRef` | **FSM `ATTACH_START` + run-fact** | mount arms attach ONLY on a confirmed active run (commit 4b: streaming-tail status, or POST /run for a user tail) |
| 14 | `stripRef` | **data** (attachStrategy) | strip+replay detail; the `resumeStream` effect reads it |
| 15 | `strippedRowRef` | **data** (attachStrategy) | the anchor row |
| 14–15 | `anchorRef {id, stepsPersisted}` | **data** (attachStrategy) | #491 tail-only: replaced `stripRef`/`strippedRowRef`. The PERSISTED assistant row that pins the run (server invariant 6) + its step frontier N; feeds `?anchor=<id>&n=<stepsPersisted>`. No strip — the seed keeps every row; entering reconnecting re-seeds from persist |
| 16 | `attachAbortRef` | **effect-owned controller** | aborted by the `abortAttach` effect in cleanup (I5) |
| 17–25 | `chatIdRef`, `openPageRef`, `getEditorSelectionRef`, `roleIdRef`, `stableIdRef`, `queuedRef`, `sendMessageRef`, `statusRef`, `lastForwardedChatIdRef` | **data** (identity/send mirrors) | unchanged — not lifecycle flags |
| NEW | `pendingSupersedeRef` | **data** (send-plumbing) | the runId injected into the next `POST /stream {supersede}`; the single replacement for the 3 DELETED one-shots (#8/#9/#10) — net −2 refs |
@@ -151,8 +151,12 @@ message. Sources, in the order they update `ctx.runFact`:
3. **Attach outcomes:** `ATTACH_LIVE` (2xx) confirms active; a 204 on a non-stripped
path is an authoritative NEGATIVE fact → the runtime dispatches `RUN_FACT{null}`,
which cancels recovery (I3 fresh-negative gate).
4. **Poll (future resume-stack iteration #491):** the delta will carry the run field;
until then the poll drives to a terminal ROW, dispatched as `POLL_TERMINAL`.
4. **Poll (#491, implemented):** the degraded poll now hits the delta endpoint
(`POST /ai-chat/messages/delta`), which ALREADY carries the run fact
(`run: {id, status} | null`) alongside the changed rows. The client does NOT yet
consume that run field — it still drives to a terminal ROW (merged by id),
dispatched as `POLL_TERMINAL` — so the run field rides the wire for a future
client that settles straight off it.
Pessimism rule: a stale-but-positive fact PERMITS entering recovery (attach); the
204 then cuts it. A fresh negative fact gates recovery OUT immediately.
@@ -178,6 +182,9 @@ Pessimism rule: a stale-but-positive fact PERMITS entering recovery (attach); th
/run) are effect-owned and aborted in cleanup (`abortAttach` on `DISPOSE`), not
render-phase refs. A client abort of an already-sent POST does not cancel the
server action, so disarming on unmount is safe.
- **attachStrategy** (strip+replay today) is behind the `resumeStream` effect; the
resume-stack iteration (#491) swaps it to tail-only WITHOUT touching the FSM.
- **attachStrategy** is behind the `resumeStream` effect; #491 swapped it to
tail-only (`?anchor=&n=`, `anchorRef` data) WITHOUT touching the FSM. Entering
reconnecting always re-seeds from persist; on a getRun failure the live partial
is dropped + replay-from-start so it is never the tail-apply base (no #137/#161
duplication).
- **Queue** stays a data structure; flush/interrupt decisions are transitions.
@@ -181,6 +181,12 @@ export interface IAiChatMessageRow {
toolCalls?: unknown;
metadata?: {
parts?: UIMessage["parts"];
// #491 step-alignment anchor: the count of FINISHED steps whose parts are in
// THIS row, written atomically with `parts` server-side (flushAssistant). The
// resume client reads it as its persisted step frontier N — the tail-only
// attach asks the run-stream registry for the frames of step N onward (the
// seed already carries steps 0..N-1). Absent on pre-#491 rows -> read as 0.
stepsPersisted?: number;
// AI SDK v6 `totalUsage` persisted on assistant rows. Legacy cumulative
// figure (sum of every step's usage for the turn); kept for back-compat and
// as the fallback for older rows that have no `contextTokens`.
@@ -4,7 +4,8 @@ import type { IAiChatMessageRow } from "@/features/ai-chat/types/ai-chat.types.t
import {
isStreamingTail,
isSettledAssistantTail,
seedRows,
stepsPersistedOf,
mergeDeltaRowsIntoPages,
mergeById,
} from "./resume-helpers.ts";
@@ -12,8 +13,18 @@ function row(
id: string,
role: string,
status?: string,
stepsPersisted?: number,
): IAiChatMessageRow {
return { id, role, content: "", status, createdAt: "2026-01-01T00:00:00Z" };
return {
id,
role,
content: "",
status,
createdAt: "2026-01-01T00:00:00Z",
...(stepsPersisted !== undefined
? { metadata: { stepsPersisted } }
: {}),
};
}
function makeMsg(id: string, text: string): UIMessage {
@@ -65,23 +76,92 @@ describe("isSettledAssistantTail", () => {
});
});
describe("seedRows", () => {
const rows = [row("u1", "user"), row("a1", "assistant", "streaming")];
it("returns the rows unchanged when not stripping", () => {
expect(seedRows(rows, false)).toBe(rows);
describe("stepsPersistedOf", () => {
it("reads metadata.stepsPersisted", () => {
expect(stepsPersistedOf(row("a1", "assistant", "streaming", 3))).toBe(3);
expect(stepsPersistedOf(row("a1", "assistant", "streaming", 0))).toBe(0);
});
it("drops the last row when stripping", () => {
const seeded = seedRows(rows, true);
expect(seeded).toHaveLength(1);
expect(seeded[0].id).toBe("u1");
it("defaults to 0 for a pre-#491 row (absent), null/undefined, or a bad value", () => {
expect(stepsPersistedOf(row("a1", "assistant", "streaming"))).toBe(0);
expect(stepsPersistedOf(null)).toBe(0);
expect(stepsPersistedOf(undefined)).toBe(0);
expect(
stepsPersistedOf({
id: "a1",
role: "assistant",
content: "",
createdAt: "x",
metadata: { stepsPersisted: -2 },
}),
).toBe(0);
});
it("returns an empty list when stripping a single-row list", () => {
expect(seedRows([row("a1", "assistant", "streaming")], true)).toHaveLength(
0,
);
it("floors a non-integer count", () => {
expect(
stepsPersistedOf({
id: "a1",
role: "assistant",
content: "",
createdAt: "x",
metadata: { stepsPersisted: 2.9 },
}),
).toBe(2);
});
});
describe("mergeDeltaRowsIntoPages", () => {
const pages = () => [
{ items: [row("u1", "user"), row("a1", "assistant", "streaming", 1)], meta: {} },
];
it("returns the pages unchanged for an empty delta", () => {
const p = pages();
expect(mergeDeltaRowsIntoPages(p, [])).toBe(p);
});
it("appends a genuinely new row to the last page in chronological order", () => {
const merged = mergeDeltaRowsIntoPages(pages(), [row("a2", "assistant", "streaming", 0)]);
expect(merged[0].items.map((i) => i.id)).toEqual(["u1", "a1", "a2"]);
});
it("replaces a grown row in place (per-step growth), never appends a duplicate", () => {
const merged = mergeDeltaRowsIntoPages(pages(), [
row("a1", "assistant", "streaming", 2),
]);
expect(merged[0].items.map((i) => i.id)).toEqual(["u1", "a1"]);
// the in-place replacement carries the grown step frontier.
expect(stepsPersistedOf(merged[0].items[1])).toBe(2);
});
it("does not mutate the input pages", () => {
const input = pages();
const before = input[0].items.slice();
mergeDeltaRowsIntoPages(input, [row("a2", "assistant", "streaming", 0)]);
expect(input[0].items).toEqual(before); // untouched
});
// #491 CONTRACT: the delta overlap window re-delivers the same rows, so merging
// MUST be idempotent — applying a delta twice equals applying it once (no growth,
// no reorder). A regression re-introduces duplicate assistant bubbles per poll.
it("is idempotent: applying the same delta twice equals once", () => {
const delta = [
row("a1", "assistant", "streaming", 2), // grown existing row
row("a2", "assistant", "streaming", 0), // new row
];
const once = mergeDeltaRowsIntoPages(pages(), delta);
const twice = mergeDeltaRowsIntoPages(once, delta);
const thrice = mergeDeltaRowsIntoPages(twice, delta);
expect(once[0].items.map((i) => i.id)).toEqual(["u1", "a1", "a2"]);
expect(twice[0].items.map((i) => i.id)).toEqual(["u1", "a1", "a2"]);
expect(twice).toEqual(once);
expect(thrice).toEqual(once);
});
it("seeds a first page when the cache is empty", () => {
const merged = mergeDeltaRowsIntoPages([], [row("u1", "user")]);
expect(merged).toHaveLength(1);
expect(merged[0].items.map((i) => i.id)).toEqual(["u1"]);
});
});
@@ -109,4 +189,37 @@ describe("mergeById", () => {
expect(mergeById(prev, null)).toBe(prev);
expect(mergeById(prev, undefined)).toBe(prev);
});
// #491 CONTRACT: the delta poll's overlap window GUARANTEES the same row is
// re-delivered across close polls, so merging must be IDEMPOTENT by id — merging
// the same row (or an equal-length list of rows) twice must not duplicate or
// reorder. This is the property the whole delta-poll design leans on; a
// regression here would re-introduce duplicate assistant bubbles on every poll.
it("is idempotent by id: re-merging the same row does not duplicate or reorder", () => {
const seed = [makeMsg("u1", "hi"), makeMsg("a1", "step 1")];
const repeat = makeMsg("a1", "step 1"); // the SAME row the overlap re-delivers
const once = mergeById(seed, repeat);
const twice = mergeById(once, repeat);
const thrice = mergeById(twice, repeat);
// Length is stable (no growth), order is stable (user then assistant).
expect(once.map((m) => m.id)).toEqual(["u1", "a1"]);
expect(twice.map((m) => m.id)).toEqual(["u1", "a1"]);
expect(thrice.map((m) => m.id)).toEqual(["u1", "a1"]);
// The repeated merge converges: the row is replaced in place, never appended.
expect(twice[1]).toBe(repeat);
});
it("is idempotent across a batch of repeated + grown rows (delta re-delivery)", () => {
// A delta poll re-delivers a1 (unchanged) and a2 (grown one step). Applying the
// batch twice must equal applying it once — the poll can re-send either.
const start = [makeMsg("u1", "hi"), makeMsg("a1", "done")];
const batch = [makeMsg("a1", "done"), makeMsg("a2", "grown step 2")];
const apply = (list: typeof start) =>
batch.reduce((acc, row) => mergeById(acc, row), list);
const once = apply(start);
const twice = apply(once);
expect(once.map((m) => m.id)).toEqual(["u1", "a1", "a2"]);
expect(twice.map((m) => m.id)).toEqual(["u1", "a1", "a2"]);
expect(twice).toEqual(once);
});
});
@@ -11,9 +11,10 @@ import type { IAiChatMessageRow } from "@/features/ai-chat/types/ai-chat.types.t
/**
* A STREAMING tail: the last persisted row is an assistant row still marked
* `status === 'streaming'`. Such a tail is stripped from the seed and rebuilt by
* the replay (`expect=live`), since the SDK's `text-start` always pushes a new
* part and replaying over a seeded in-progress row would duplicate its text.
* `status === 'streaming'`. #491 (tail-only): such a tail is seeded UNCHANGED
* it carries the persisted steps 0..N-1 and the run-stream registry's tail
* (frames for steps >= N) is APPENDED to it by the SDK's `readUIMessageStream`
* continuation. Only the presence of this tail decides WHETHER to attach.
*/
export function isStreamingTail(rows: IAiChatMessageRow[]): boolean {
const tail = rows[rows.length - 1];
@@ -32,15 +33,61 @@ export function isSettledAssistantTail(rows: IAiChatMessageRow[]): boolean {
}
/**
* Seed rows for `useChat`: return the rows unchanged, or without the last row when
* `strip` is set (the streaming tail is stripped so the live replay rebuilds it
* without duplicating parts).
* #491 tail-only anchor: the count of FINISHED steps whose parts are persisted in
* THIS assistant row (`metadata.stepsPersisted`), written atomically with `parts`
* server-side. The resume client reads it as its persisted step frontier N the
* tail-only attach asks the run-stream registry for the frames of step N onward
* (the seed already carries steps 0..N-1). Absent on pre-#491 rows => 0.
*/
export function seedRows(
export function stepsPersistedOf(
row: IAiChatMessageRow | null | undefined,
): number {
const n = row?.metadata?.stepsPersisted;
return typeof n === "number" && n >= 0 ? Math.floor(n) : 0;
}
/** One page of the messages infinite-query cache (`{ items, meta }`). */
export interface IMessagePage {
items: IAiChatMessageRow[];
meta: unknown;
}
/**
* #491 delta-poll merge: upsert the delta poll's `rows` into the messages
* infinite-query page structure IDEMPOTENTLY by id. The delta endpoint's overlap
* window GUARANTEES occasional REPEATS, so this MUST converge: a row already
* present is REPLACED IN PLACE (per-step growth of an in-progress row), a new row
* is APPENDED to the last page in chronological order (the server returns delta
* rows oldest-first). Applying the same delta twice equals applying it once. Never
* mutates the input pages (returns fresh page objects with cloned item arrays).
*/
export function mergeDeltaRowsIntoPages(
pages: IMessagePage[],
rows: IAiChatMessageRow[],
strip: boolean,
): IAiChatMessageRow[] {
return strip ? rows.slice(0, -1) : rows;
): IMessagePage[] {
if (rows.length === 0) return pages;
const next: IMessagePage[] = pages.map((p) => ({
...p,
items: p.items.slice(),
}));
const locate = (id: string): [number, number] | null => {
for (let pi = 0; pi < next.length; pi++) {
const ii = next[pi].items.findIndex((it) => it.id === id);
if (ii !== -1) return [pi, ii];
}
return null;
};
for (const row of rows) {
const at = locate(row.id);
if (at) {
next[at[0]].items[at[1]] = row; // replace in place — idempotent by id
} else if (next.length > 0) {
next[next.length - 1].items.push(row); // append chronologically
} else {
next.push({ items: [row], meta: undefined });
}
}
return next;
}
/**
@@ -0,0 +1,83 @@
import { describe, it, expect } from "vitest";
import { readUIMessageStream, type UIMessage } from "ai";
import pkg from "../../../../package.json";
/**
* PIN-SPEC TRIP-WIRE (#491). The tail-only attach continuation relies on THREE
* behaviors of `ai@6.0.207`, verified line-by-line in the issue. Without this
* test, an `ai` bump could silently break attach (the client would append the
* live tail to the wrong message, or duplicate a step):
*
* 1. `readUIMessageStream({ message })` CONTINUES the passed message it does
* not start a fresh one so the tail streamed after a re-seed is appended to
* the seeded assistant row (the same DB id).
* 2. A `start` frame does NOT reset the existing message's parts (so the seeded
* steps 0..N-1 survive; the synthetic `start` the registry prepends only
* carries the run-fact metadata).
* 3. Text parts do NOT cross a `finish-step` boundary a new `text-start` after
* `finish-step` is a NEW part so the reconstructed steps stay separated and
* the step frontier stays meaningful.
*
* If an `ai` upgrade changes any of these, this test fails LOUD instead of the
* resume path silently corrupting.
*/
describe("ai SDK continuation trip-wire (#491, tail-only attach)", () => {
it("is pinned to the exact ai version the continuation was verified against", () => {
// A caret/range bump is exactly what would silently break attach — require an
// exact pin. Bumping ai MUST re-verify the behavior asserted below, then this.
expect((pkg as { dependencies: Record<string, string> }).dependencies.ai).toBe(
"6.0.207",
);
});
it("continues the seeded message: start does not reset parts, the tail appends as new parts", async () => {
// A seeded assistant row with ONE finished step already reconstructed.
const seeded: UIMessage = {
id: "assistant-1",
role: "assistant",
parts: [
{ type: "step-start" },
{ type: "text", text: "STEP0", state: "done" },
],
} as UIMessage;
// The tail the registry delivers on re-attach: a synthetic start (run-fact),
// then step 1's frames, then finish. As UI-message chunks (what the SSE frames
// decode to).
const chunks = [
{ type: "start", messageMetadata: { runId: "r1", chatId: "c1" } },
{ type: "start-step" },
{ type: "text-start", id: "t1" },
{ type: "text-delta", id: "t1", delta: "STEP1" },
{ type: "text-end", id: "t1" },
{ type: "finish-step" },
{ type: "finish" },
];
const stream = new ReadableStream({
start(c) {
for (const ch of chunks) c.enqueue(ch);
c.close();
},
});
let last: UIMessage | undefined;
for await (const msg of readUIMessageStream({ message: seeded, stream })) {
last = msg;
}
expect(last).toBeDefined();
// Same message id (continuation, not a fresh message).
expect(last!.id).toBe("assistant-1");
// The seeded step-0 parts SURVIVED the `start` frame, and step 1 was appended
// as SEPARATE parts (text did not cross the finish-step boundary).
const shape = last!.parts.map((p) => `${p.type}:${(p as { text?: string }).text ?? ""}`);
expect(shape).toEqual([
"step-start:",
"text:STEP0",
"step-start:",
"text:STEP1",
]);
// The run-fact metadata from the synthetic start frame is applied.
expect(last!.metadata).toMatchObject({ runId: "r1", chatId: "c1" });
});
});
@@ -3,11 +3,20 @@ import { atom } from "jotai";
// import would drag the whole @tiptap/core engine into the eager graph of every
// shell component that reads one of these atoms.
import type { Editor } from "@tiptap/core";
import type { HocuspocusProvider } from "@hocuspocus/provider";
import { PageEditMode } from "@/features/user/types/user.types.ts";
import type { DictationUnavailableReason } from "@/features/dictation/dictation-status";
export const pageEditorAtom = atom<Editor | null>(null);
// #370 — the active page's collab provider, published by the page editor so the
// header menu can emit the "save-version" stateless signal (Cmd+S / button).
// Null when the page is read-only / collab isn't connected. A typed initial
// value (rather than an explicit generic) keeps jotai's overload resolution on
// the writable PrimitiveAtom branch.
const initialCollabProvider: HocuspocusProvider | null = null;
export const collabProviderAtom = atom(initialCollabProvider);
export const titleEditorAtom = atom<Editor | null>(null);
export const readOnlyEditorAtom = atom<Editor | null>(null);
@@ -31,11 +31,18 @@ import { useAtom, useAtomValue, useSetAtom } from "jotai";
import useCollaborationUrl from "@/features/editor/hooks/use-collaboration-url";
import { currentUserAtom } from "@/features/user/atoms/current-user-atom";
import {
collabProviderAtom,
currentPageEditModeAtom,
dictationAvailabilityAtom,
pageEditorAtom,
yjsConnectionStatusAtom,
} from "@/features/editor/atoms/editor-atoms";
import { notifications } from "@mantine/notifications";
import {
VERSION_SAVED_MESSAGE_TYPE,
type VersionSavedMessage,
saveVersionPending,
} from "@/features/page-history/version-messages";
import { asideStateAtom } from "@/components/layouts/global/hooks/atoms/sidebar-atom";
import {
activeCommentIdAtom,
@@ -124,6 +131,7 @@ export default function PageEditor({
const [currentUser] = useAtom(currentUserAtom);
const [, setEditor] = useAtom(pageEditorAtom);
const setCollabProvider = useSetAtom(collabProviderAtom);
const [, setAsideState] = useAtom(asideStateAtom);
const [, setActiveCommentId] = useAtom(activeCommentIdAtom);
const [showCommentPopup, setShowCommentPopup] = useAtom(showCommentPopupAtom);
@@ -181,6 +189,24 @@ export default function PageEditor({
const onStatelessHandler = ({ payload }: onStatelessParameters) => {
try {
const message = JSON.parse(payload);
// #370 — a version was saved somewhere; live-refresh the history panel
// on every client. Only the client that pressed Save (tracked by the
// module-level flag) shows the confirmation toast.
if (message?.type === VERSION_SAVED_MESSAGE_TYPE) {
const versionMsg = message as VersionSavedMessage;
queryClient.invalidateQueries({
queryKey: ["page-history-list"],
});
if (saveVersionPending.current) {
saveVersionPending.current = false;
notifications.show({
message: versionMsg.alreadySaved
? t("Already saved as the latest version")
: t("Version saved"),
});
}
return;
}
if (message?.type !== "page.updated" || !message.updatedAt) return;
const pageData = queryClient.getQueryData<IPage>(["pages", slugId]);
if (pageData) {
@@ -238,12 +264,16 @@ export default function PageEditor({
local.on("synced", onLocalSyncedHandler);
providersRef.current = { socket, local, remote };
// #370 — publish the provider so the header menu can emit save-version.
setCollabProvider(remote);
setProvidersReady(true);
} else {
setCollabProvider(providersRef.current.remote);
setProvidersReady(true);
}
// Only destroy on final unmount
return () => {
setCollabProvider(null);
providersRef.current?.socket.destroy();
providersRef.current?.remote.destroy();
providersRef.current?.local.destroy();
@@ -1,4 +1,11 @@
import { Text, Group, UnstyledButton, Avatar, Tooltip } from "@mantine/core";
import {
Text,
Group,
UnstyledButton,
Avatar,
Tooltip,
Badge,
} from "@mantine/core";
import { CustomAvatar } from "@/components/ui/custom-avatar.tsx";
import { AgentAvatarStack } from "@/components/ui/agent-avatar-stack.tsx";
import { formattedDate } from "@/lib/time";
@@ -7,36 +14,59 @@ import clsx from "clsx";
import { IPageHistory } from "@/features/page-history/types/page.types";
import { memo, useCallback } from "react";
import { useSetAtom } from "jotai";
import { useTranslation } from "react-i18next";
import { historyAtoms } from "@/features/page-history/atoms/history-atoms.ts";
const MAX_VISIBLE_AVATARS = 5;
/**
* #370 map a snapshot's intentionality tier to its badge. `version: true`
* marks the intentional points (manual / agent); autosaves (boundary / idle /
* legacy null) are non-versions and get dimmed in the list.
*/
type HistoryKindMeta = { labelKey: string; color: string; version: boolean };
export function historyKindMeta(kind?: string | null): HistoryKindMeta {
switch (kind) {
case "manual":
return { labelKey: "Saved", color: "blue", version: true };
case "agent":
return { labelKey: "Agent version", color: "violet", version: true };
case "boundary":
return { labelKey: "Boundary", color: "gray", version: false };
default: // "idle" | null | undefined (legacy autosave)
return { labelKey: "Autosave", color: "gray", version: false };
}
}
interface HistoryItemProps {
historyItem: IPageHistory;
index: number;
onSelect: (id: string, index: number) => void;
onHover?: (id: string, index: number) => void;
// The previous snapshot for diff/restore is resolved by id from the FULL list
// in the parent (resolvePrevSnapshotId), so the item only needs to report its
// own id — never a list index (which would be the filtered-view index).
onSelect: (id: string) => void;
onHover?: (id: string) => void;
onHoverEnd?: () => void;
isActive: boolean;
}
const HistoryItem = memo(function HistoryItem({
historyItem,
index,
onSelect,
onHover,
onHoverEnd,
isActive,
}: HistoryItemProps) {
const setHistoryModalOpen = useSetAtom(historyAtoms);
const { t } = useTranslation();
const kindMeta = historyKindMeta(historyItem.kind);
const handleClick = useCallback(() => {
onSelect(historyItem.id, index);
}, [onSelect, historyItem.id, index]);
onSelect(historyItem.id);
}, [onSelect, historyItem.id]);
const handleMouseEnter = useCallback(() => {
onHover?.(historyItem.id, index);
}, [onHover, historyItem.id, index]);
onHover?.(historyItem.id);
}, [onHover, historyItem.id]);
const contributors = historyItem.contributors;
const hasContributors = contributors && contributors.length > 0;
@@ -49,8 +79,20 @@ const HistoryItem = memo(function HistoryItem({
onMouseEnter={handleMouseEnter}
onMouseLeave={onHoverEnd}
className={clsx(classes.history, { [classes.active]: isActive })}
// #370 — dim autosnapshots so intentional versions stand out.
style={{ opacity: kindMeta.version ? 1 : 0.55 }}
>
<Text size="sm">{formattedDate(new Date(historyItem.createdAt))}</Text>
<Group gap={6} wrap="nowrap" justify="space-between">
<Text size="sm">{formattedDate(new Date(historyItem.createdAt))}</Text>
<Badge
size="xs"
radius="sm"
variant={kindMeta.version ? "filled" : "light"}
color={kindMeta.color}
>
{t(kindMeta.labelKey)}
</Badge>
</Group>
<Group gap={6} wrap="nowrap" mt={4}>
{hasContributors ? (
@@ -2,14 +2,16 @@ import {
usePageHistoryListQuery,
prefetchPageHistory,
} from "@/features/page-history/queries/page-history-query";
import HistoryItem from "@/features/page-history/components/history-item";
import HistoryItem, {
historyKindMeta,
} from "@/features/page-history/components/history-item";
import {
activeHistoryIdAtom,
activeHistoryPrevIdAtom,
historyAtoms,
} from "@/features/page-history/atoms/history-atoms";
import { useAtom, useSetAtom } from "jotai";
import { useCallback, useEffect, useMemo, useRef } from "react";
import { useCallback, useEffect, useMemo, useRef, useState } from "react";
import {
Button,
ScrollArea,
@@ -17,9 +19,12 @@ import {
Divider,
Loader,
Center,
Switch,
Text,
} from "@mantine/core";
import { useTranslation } from "react-i18next";
import { useHistoryRestore } from "@/features/page-history/hooks";
import { resolvePrevSnapshotId } from "@/features/page-history/utils/resolve-prev-snapshot";
const PREFETCH_DELAY_MS = 150;
@@ -47,6 +52,22 @@ function HistoryList({ pageId }: Props) {
[pageHistoryData],
);
// #370 — "only versions" filter: hide autosnapshots (idle/boundary/legacy
// null), keep only intentional points (manual/agent). Filtering is over the
// already-loaded pages; the diff/restore still targets the true previous
// snapshot, so items carry their index within the FULL list.
const [onlyVersions, setOnlyVersions] = useState(false);
// Reuse historyKindMeta().version — the SAME predicate the badge (HistoryItem)
// uses to mark intentional points — so the "Only versions" filter and the badge
// can never drift apart when a future intentional kind is added.
const visibleItems = useMemo(
() =>
onlyVersions
? historyItems.filter((item) => historyKindMeta(item.kind).version)
: historyItems,
[historyItems, onlyVersions],
);
const loadMoreRef = useRef<HTMLDivElement>(null);
const prefetchTimeoutRef = useRef<ReturnType<typeof setTimeout> | null>(null);
@@ -60,11 +81,13 @@ function HistoryList({ pageId }: Props) {
}, []);
const handleHover = useCallback(
(historyId: string, index: number) => {
(historyId: string) => {
clearPrefetchTimeout();
prefetchTimeoutRef.current = setTimeout(() => {
prefetchPageHistory(historyId);
const prevId = historyItems[index + 1]?.id;
// The true previous snapshot in the FULL list (not the previous visible
// one under the "only versions" filter).
const prevId = resolvePrevSnapshotId(historyItems, historyId);
if (prevId) {
prefetchPageHistory(prevId);
}
@@ -78,9 +101,11 @@ function HistoryList({ pageId }: Props) {
}, [clearPrefetchTimeout]);
const handleSelect = useCallback(
(id: string, index: number) => {
(id: string) => {
setActiveHistoryId(id);
setActiveHistoryPrevId(historyItems[index + 1]?.id ?? "");
// Baseline = true previous snapshot in the FULL list, so the "only
// versions" filter never diffs/restores against the wrong item.
setActiveHistoryPrevId(resolvePrevSnapshotId(historyItems, id));
},
[historyItems, setActiveHistoryId, setActiveHistoryPrevId],
);
@@ -128,12 +153,27 @@ function HistoryList({ pageId }: Props) {
return (
<div>
<Group px="xs" py={6} justify="flex-end">
<Switch
size="xs"
checked={onlyVersions}
onChange={(e) => setOnlyVersions(e.currentTarget.checked)}
label={t("Only versions")}
/>
</Group>
<ScrollArea h={620} w="100%" type="scroll" scrollbarSize={5}>
{historyItems.map((historyItem, index) => (
{onlyVersions && visibleItems.length === 0 && (
<Center py="md">
<Text size="sm" c="dimmed">
{t("No saved versions yet.")}
</Text>
</Center>
)}
{visibleItems.map((historyItem) => (
<HistoryItem
key={historyItem.id}
historyItem={historyItem}
index={index}
onSelect={handleSelect}
onHover={handleHover}
onHoverEnd={clearPrefetchTimeout}
@@ -24,6 +24,10 @@ export interface IPageHistory {
updatedAt: string;
lastUpdatedBy: IPageHistoryUser;
contributors?: IPageHistoryUser[];
// #370 — intentionality tier: 'manual'/'agent' are versions (intentional
// points), 'idle'/'boundary' are autosnapshots; null/undefined = legacy
// autosave. Derived server-side, drives the history badge + "versions" filter.
kind?: "manual" | "agent" | "idle" | "boundary" | null;
// Provenance markers copied off the page row when the snapshot was saved.
// `'agent'` marks a version written by the AI agent; `lastUpdatedAiChatId`
// (when present) deep-links to the chat that produced the edit.
@@ -0,0 +1,42 @@
import { describe, it, expect } from "vitest";
import { resolvePrevSnapshotId } from "./resolve-prev-snapshot";
// #370 F4 — the risky client path: with the "only versions" filter active, diff
// and restore must still baseline against the TRUE previous snapshot in the FULL
// list, never the previous VISIBLE version (which would skip the autosnapshots
// between two versions). These pin that the resolution is by FULL-list order.
describe("resolvePrevSnapshotId", () => {
// Newest-first, as the history list stores it: a version, then two autosaves,
// then an older version.
const full = [
{ id: "v2", kind: "manual" },
{ id: "a2", kind: "idle" },
{ id: "a1", kind: "boundary" },
{ id: "v1", kind: "manual" },
{ id: "a0", kind: null },
];
it("returns the immediate FULL-list successor, not the previous visible version", () => {
// Selecting v2 while filtered to versions-only must baseline against a2 (the
// real chronological predecessor), NOT v1 (the previous visible version).
expect(resolvePrevSnapshotId(full, "v2")).toBe("a2");
});
it("resolves an autosnapshot's predecessor by full-list order", () => {
expect(resolvePrevSnapshotId(full, "a1")).toBe("v1");
});
it("returns '' for the oldest item (no predecessor)", () => {
expect(resolvePrevSnapshotId(full, "a0")).toBe("");
});
it("returns '' for an id not in the list", () => {
expect(resolvePrevSnapshotId(full, "missing")).toBe("");
});
it("does not depend on a filtered subset — same result whatever is visible", () => {
// The helper only ever sees the full list; a filtered view cannot change the
// baseline it computes.
expect(resolvePrevSnapshotId(full, "v1")).toBe("a0");
});
});
@@ -0,0 +1,22 @@
/**
* #370 resolve the TRUE previous snapshot for a history item.
*
* The history panel can be filtered to "only versions" (manual/agent), but diff
* and restore must always compare against the immediately-preceding snapshot in
* the FULL, unfiltered list NOT the previous VISIBLE item. Comparing against
* the previous visible version would silently skip the autosnapshots between two
* versions and diff/restore the wrong baseline.
*
* Given the full (newest-first) list and an item id, this returns the id of the
* item right after it in the full list (its chronological predecessor), or "" if
* it is the oldest / not found. Pure and list-order-preserving so it can be unit
* tested without mounting the component.
*/
export function resolvePrevSnapshotId(
fullItems: ReadonlyArray<{ id: string }>,
id: string,
): string {
const index = fullItems.findIndex((item) => item.id === id);
if (index === -1) return "";
return fullItems[index + 1]?.id ?? "";
}
@@ -0,0 +1,28 @@
/**
* #370 page-version stateless wire formats. Kept in one place so the client
* emitter (Save hotkey / button) and the client listener (page-editor) agree
* with the server (PersistenceExtension) on the message shapes.
*/
/** Client server: "save a version now". The server derives the tier
* (manual/agent) from the signed connection actor, never from this payload. */
export const SAVE_VERSION_MESSAGE_TYPE = "save-version";
/** Server → all clients: a version was saved (or promoted / already existed). */
export const VERSION_SAVED_MESSAGE_TYPE = "version.saved";
export interface VersionSavedMessage {
type: typeof VERSION_SAVED_MESSAGE_TYPE;
historyId: string;
kind: "manual" | "agent";
/** True when the latest snapshot was already a manual version (a no-op save). */
alreadySaved: boolean;
}
/**
* Cross-component coordination flag so only the client that pressed Save shows
* the confirmation toast, while every other client silently refreshes its
* history panel on the broadcast. A module-level ref avoids stale-closure
* pitfalls in the editor's long-lived stateless handler.
*/
export const saveVersionPending = { current: false };
@@ -3,6 +3,7 @@ import {
IconArrowRight,
IconArrowsHorizontal,
IconClockHour4,
IconDeviceFloppy,
IconDots,
IconEye,
IconEyeOff,
@@ -17,7 +18,7 @@ import {
IconTrash,
IconWifiOff,
} from "@tabler/icons-react";
import React, { useEffect, useRef, useState } from "react";
import React, { useCallback, useEffect, useRef, useState } from "react";
import { useAsideTriggerProps } from "@/hooks/use-toggle-aside.tsx";
import { useAtom, useAtomValue } from "jotai";
import { historyAtoms } from "@/features/page-history/atoms/history-atoms.ts";
@@ -39,9 +40,14 @@ import { Trans, useTranslation } from "react-i18next";
import ExportModal from "@/components/common/export-modal";
import { convertProseMirrorToMarkdown } from "@docmost/prosemirror-markdown/browser";
import {
collabProviderAtom,
pageEditorAtom,
yjsConnectionStatusAtom,
} from "@/features/editor/atoms/editor-atoms.ts";
import {
SAVE_VERSION_MESSAGE_TYPE,
saveVersionPending,
} from "@/features/page-history/version-messages.ts";
import { formattedDate } from "@/lib/time.ts";
import { PageEditModeToggle } from "@/features/user/components/page-state-pref.tsx";
import MovePageModal from "@/features/page/components/move-page-modal.tsx";
@@ -72,9 +78,34 @@ export default function PageHeaderMenu({ readOnly }: PageHeaderMenuProps) {
});
const isDeleted = !!page?.deletedAt;
const [workspace] = useAtom(workspaceAtom);
const collabProvider = useAtomValue(collabProviderAtom);
// Community public-sharing entry point (replaces the removed EE PageShareModal)
const workspaceSharingDisabled = workspace?.settings?.sharing?.disabled === true;
// #370 — explicit "save a version" (Cmd+S / Save button). One path for the
// human; the server derives the tier from the signed actor. Readers can't save
// (the button is hidden and the collab connection is read-only server-side).
const handleSaveVersion = useCallback(() => {
if (readOnly || !collabProvider) return;
// Flag this client as the initiator so only it shows the confirmation toast;
// a safety timeout clears it if no broadcast comes back (e.g. offline).
saveVersionPending.current = true;
window.setTimeout(() => {
saveVersionPending.current = false;
}, 5000);
collabProvider.sendStateless(
JSON.stringify({ type: SAVE_VERSION_MESSAGE_TYPE }),
);
}, [readOnly, collabProvider]);
// mod+S must also block the browser's "Save page" dialog. `triggerOnContent-
// Editable` + empty ignore-list so it fires while typing in the editor/title.
useHotkeys(
[["mod+S", handleSaveVersion, { preventDefault: true }]],
[],
true,
);
useHotkeys(
[
[
@@ -133,15 +164,16 @@ export default function PageHeaderMenu({ readOnly }: PageHeaderMenuProps) {
</ActionIcon>
</Tooltip>
<PageActionMenu readOnly={readOnly} />
<PageActionMenu readOnly={readOnly} onSaveVersion={handleSaveVersion} />
</>
);
}
interface PageActionMenuProps {
readOnly?: boolean;
onSaveVersion?: () => void;
}
function PageActionMenu({ readOnly }: PageActionMenuProps) {
function PageActionMenu({ readOnly, onSaveVersion }: PageActionMenuProps) {
const { t } = useTranslation();
const [, setHistoryModalOpen] = useAtom(historyAtoms);
const clipboard = useClipboard({ timeout: 500 });
@@ -303,6 +335,20 @@ function PageActionMenu({ readOnly }: PageActionMenuProps) {
</Group>
</Menu.Item>
{!readOnly && (
<Menu.Item
leftSection={<IconDeviceFloppy size={16} />}
onClick={onSaveVersion}
rightSection={
<Text size="xs" c="dimmed">
{t("Ctrl+S")}
</Text>
}
>
{t("Save version")}
</Menu.Item>
)}
<Menu.Item
leftSection={<IconHistory size={16} />}
onClick={openHistoryModal}
@@ -16,6 +16,7 @@ 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: [
@@ -31,6 +32,7 @@ import { EnvironmentModule } from '../integrations/environment/environment.modul
exports: [CollaborationGateway],
imports: [
TokenModule,
ApiKeyModule,
WatcherModule,
StorageModule.forRootAsync({
imports: [EnvironmentModule],
+30 -4
View File
@@ -1,10 +1,36 @@
export const HISTORY_INTERVAL = 5 * 60 * 1000;
export const HISTORY_FAST_INTERVAL = 60 * 1000;
export const HISTORY_FAST_THRESHOLD = 5 * 60 * 1000;
// #348 — debounce window for the per-page RAG re-embed job. Repeated saves
// within this window collapse to a single delayed job (coalesced by a stable
// jobId), so active editing does not pile up expensive re-embeds (external API
// + page_embeddings rewrite, concurrency 1). The worker reads the CURRENT page
// state at run time, so the last content within the window wins.
export const EMBED_DEBOUNCE_MS = 30 * 1000;
/**
* #370 page-history intentionality tiers. Domain of `page_history.kind`.
* - 'manual' / 'agent' Tier 1 versions (intentional points)
* - 'idle' / 'boundary' Tier 0 autosnapshots (safety net)
* A legacy `null` kind is treated as an autosave.
*/
export type PageHistoryKind = 'manual' | 'agent' | 'idle' | 'boundary';
/**
* #370 trailing idle-flush windows. A page's pending idle snapshot is
* re-armed on every store and fires this long after edits go quiet, so a burst
* of edits collapses into a single autosnapshot instead of one-per-store. Human
* sessions are noisier and less risky, so they flush less often than the agent.
*/
export const IDLE_INTERVAL_USER = 60 * 60 * 1000; // 60m
export const IDLE_INTERVAL_AGENT = 15 * 60 * 1000; // 15m
/**
* #370 max-wait ceiling for the idle flush. Pure trailing debounce starves the
* safety net: hocuspocus stores at least every ~45s, so a CONTINUOUS editing
* session would re-arm the trailing timer forever and never take an idle
* snapshot until edits finally go quiet (up to IDLE_INTERVAL_USER = 60m). This
* ceiling bounds the actual wait from the FIRST edit of a burst, so an idle
* snapshot fires at least this often during a long unbroken session restoring
* a recovery point cadence closer to the old heuristic without one-per-store
* noise. Mirrors hocuspocus's own maxDebounce idea.
*/
export const IDLE_MAX_WAIT_USER = 10 * 60 * 1000; // 10m
export const IDLE_MAX_WAIT_AGENT = 5 * 60 * 1000; // 5m
@@ -52,6 +52,7 @@ 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.
@@ -79,12 +80,15 @@ 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);
@@ -231,4 +235,73 @@ 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);
});
});
});
@@ -14,20 +14,37 @@ 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 { JwtCollabPayload, JwtType } from '../../core/auth/dto/jwt-payload';
import {
JwtApiKeyPayload,
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) {
@@ -54,6 +71,36 @@ 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;
@@ -1,84 +1,93 @@
import { computeHistoryJob, resolveSource } from './persistence.extension';
import {
computeHistoryJob,
resolveSource,
} from './persistence.extension';
import {
HISTORY_FAST_INTERVAL,
HISTORY_FAST_THRESHOLD,
HISTORY_INTERVAL,
IDLE_INTERVAL_AGENT,
IDLE_INTERVAL_USER,
IDLE_MAX_WAIT_AGENT,
IDLE_MAX_WAIT_USER,
} from '../constants';
// A fixed clock + fixed createdAt make pageAge deterministic.
const NOW = 1_700_000_000_000;
const PAGE_ID = '550e8400-e29b-41d4-a716-446655440000';
// Build a minimal page whose age (NOW - createdAt) is exactly `ageMs`.
const pageAged = (ageMs: number) => ({
id: PAGE_ID,
createdAt: new Date(NOW - ageMs),
});
const page = { id: PAGE_ID };
describe('computeHistoryJob', () => {
it('agent edit → delay MUST be 0 and job id is source-keyed', () => {
// INVARIANT (§15 H2 / persistence.extension): the agent delay MUST stay 0.
// The worker re-reads the page row at run time, so any non-zero delay risks
// snapshotting content a later human edit has already overwritten. This is
// the load-bearing assertion of this spec — do not relax it.
const { jobId, delay } = computeHistoryJob(pageAged(0), 'agent', NOW);
expect(delay).toBe(0);
expect(jobId).toBe(`${PAGE_ID}-agent`);
});
it('agent edit on an OLD page is still delay 0 (age never applies to agents)', () => {
// Even when the page is far older than the fast threshold, the agent path
// must short-circuit to 0 — age-based debounce is a human-only concern.
const { jobId, delay } = computeHistoryJob(
pageAged(HISTORY_FAST_THRESHOLD + 60_000),
'agent',
NOW,
);
expect(delay).toBe(0);
expect(jobId).toBe(`${PAGE_ID}-agent`);
});
it('human edit on a YOUNG page (age < threshold) → fast interval, bare job id', () => {
const { jobId, delay } = computeHistoryJob(
pageAged(HISTORY_FAST_THRESHOLD - 1),
'user',
NOW,
);
expect(delay).toBe(HISTORY_FAST_INTERVAL);
describe('computeHistoryJob (#370 — shared trailing idle pipeline)', () => {
it('human edit → user idle window, bare page.id job', () => {
// Humans and the agent now share ONE idle job per page (jobId = page.id).
// The agent's old delay=0 fast path is GONE — intentional agent points now
// arrive via the explicit save-version signal, not a zero-delay snapshot.
const { jobId, delay } = computeHistoryJob(page, 'user');
expect(delay).toBe(IDLE_INTERVAL_USER);
expect(jobId).toBe(PAGE_ID);
});
it('human edit on an OLD page (age > threshold) → standard interval', () => {
const { jobId, delay } = computeHistoryJob(
pageAged(HISTORY_FAST_THRESHOLD + 1),
'user',
NOW,
);
expect(delay).toBe(HISTORY_INTERVAL);
it('agent edit → agent idle window (shorter), still the bare page.id job', () => {
const { jobId, delay } = computeHistoryJob(page, 'agent');
expect(delay).toBe(IDLE_INTERVAL_AGENT);
// No `-agent` suffix anymore: the agent joins the common idle pipeline.
expect(jobId).toBe(PAGE_ID);
});
it('boundary: pageAge EXACTLY === threshold takes the slow branch (the `<` is strict)', () => {
// Off-by-one guard: the condition is `pageAge < HISTORY_FAST_THRESHOLD`, so
// an age of exactly the threshold is NOT "fast" — it must use HISTORY_INTERVAL.
const { delay } = computeHistoryJob(
pageAged(HISTORY_FAST_THRESHOLD),
'user',
NOW,
);
expect(delay).toBe(HISTORY_INTERVAL);
it('agent flushes sooner than a human', () => {
expect(IDLE_INTERVAL_AGENT).toBeLessThan(IDLE_INTERVAL_USER);
});
it('treats any non-"agent" source string as human', () => {
// resolveSource only ever yields 'agent' | 'user', but guard the contract:
// the agent branch keys strictly on === 'agent'.
const { jobId, delay } = computeHistoryJob(pageAged(0), 'user', NOW);
expect(delay).toBe(HISTORY_FAST_INTERVAL);
it('treats any non-"agent" source string as human (keys strictly on === agent)', () => {
const { jobId, delay } = computeHistoryJob(page, 'user');
expect(delay).toBe(IDLE_INTERVAL_USER);
expect(jobId).toBe(PAGE_ID);
});
// #370 review round-1 WARNING: the max-wait ceiling prevents autosnapshot
// starvation during a continuous editing session (the trailing timer would
// otherwise re-arm forever and never fire).
describe('max-wait ceiling', () => {
const T0 = 1_000_000; // arbitrary fixed epoch for deterministic tests
it('once a burst is armed, delay clamps to the remaining max-wait budget', () => {
// 1 minute into the burst the USER interval (60m) far exceeds the remaining
// max-wait budget (10m - 1m = 9m), so the delay is clamped DOWN to that
// remaining budget — the full interval is NOT used once a ceiling applies.
const { delay } = computeHistoryJob(page, 'user', T0, T0 + 60_000);
expect(delay).toBe(IDLE_MAX_WAIT_USER - 60_000);
});
it('never waits longer than the max-wait budget from the burst start', () => {
// A store arriving right at the ceiling → delay 0 (fire promptly).
const { delay } = computeHistoryJob(
page,
'user',
T0,
T0 + IDLE_MAX_WAIT_USER,
);
expect(delay).toBe(0);
});
it('past the ceiling never returns a negative delay', () => {
const { delay } = computeHistoryJob(
page,
'user',
T0,
T0 + IDLE_MAX_WAIT_USER + 5 * 60_000,
);
expect(delay).toBe(0);
});
it('the agent ceiling is shorter than the user ceiling', () => {
expect(IDLE_MAX_WAIT_AGENT).toBeLessThan(IDLE_MAX_WAIT_USER);
const { delay } = computeHistoryJob(
page,
'agent',
T0,
T0 + IDLE_MAX_WAIT_AGENT,
);
expect(delay).toBe(0);
});
it('without a burstStart there is no ceiling (backward-compatible)', () => {
expect(computeHistoryJob(page, 'user').delay).toBe(IDLE_INTERVAL_USER);
expect(computeHistoryJob(page, 'agent').delay).toBe(IDLE_INTERVAL_AGENT);
});
});
});
describe('resolveSource (truth table)', () => {
@@ -40,11 +40,12 @@ describe('PersistenceExtension.onStoreDocument — Approach-A boundary snapshot'
let pageHistoryRepo: {
saveHistory: jest.Mock;
findPageLastHistory: jest.Mock;
updateHistoryKind: jest.Mock;
};
let aiQueue: { add: jest.Mock };
let historyQueue: { add: jest.Mock };
let historyQueue: { add: jest.Mock; remove: jest.Mock };
let notificationQueue: { add: jest.Mock };
let collabHistory: { addContributors: jest.Mock };
let collabHistory: { addContributors: jest.Mock; popContributors: jest.Mock };
let transclusionService: {
syncPageTransclusions: jest.Mock;
syncPageReferences: jest.Mock;
@@ -93,13 +94,22 @@ describe('PersistenceExtension.onStoreDocument — Approach-A boundary snapshot'
pageHistoryRepo = {
saveHistory: jest.fn().mockImplementation(async () => {
callOrder.push('saveHistory');
return { id: 'history-1' };
}),
findPageLastHistory: jest.fn().mockResolvedValue(null),
updateHistoryKind: jest.fn().mockResolvedValue(undefined),
};
aiQueue = { add: jest.fn().mockResolvedValue(undefined) };
historyQueue = { add: jest.fn().mockResolvedValue(undefined) };
historyQueue = {
add: jest.fn().mockResolvedValue(undefined),
// #370 — enqueuePageHistory now removes any pending idle job before re-adding.
remove: jest.fn().mockResolvedValue(undefined),
};
notificationQueue = { add: jest.fn().mockResolvedValue(undefined) };
collabHistory = { addContributors: jest.fn().mockResolvedValue(undefined) };
collabHistory = {
addContributors: jest.fn().mockResolvedValue(undefined),
popContributors: jest.fn().mockResolvedValue([]),
};
transclusionService = {
syncPageTransclusions: jest.fn().mockResolvedValue(undefined),
syncPageReferences: jest.fn().mockResolvedValue(undefined),
@@ -165,6 +175,50 @@ describe('PersistenceExtension.onStoreDocument — Approach-A boundary snapshot'
expect(pageRepo.updatePage.mock.calls[0][0].lastUpdatedSource).toBe('user');
});
// #370 review round-1 SUGGESTION: the boundary was GENERALIZED from a
// user→agent special-case to ANY lastUpdatedSource transition. These pin the
// generalized behaviour it was rebuilt for.
describe('generalized boundary — any source transition', () => {
// Same persisted page but with an explicit prior source.
const pageWithPriorSource = (prior: string | null) => ({
...persistedHumanPage('NEW CONTENT'),
lastUpdatedSource: prior,
});
it('agent→user transition fires the boundary (pins the prior agent revision)', async () => {
const document = ydocFor(doc('NEW CONTENT'));
pageRepo.findById.mockResolvedValue(pageWithPriorSource('agent'));
pageHistoryRepo.findPageLastHistory.mockResolvedValue(null);
await ext.onStoreDocument(buildData(document, 'user') as any);
expect(pageHistoryRepo.saveHistory).toHaveBeenCalledTimes(1);
expect(callOrder).toEqual(['saveHistory', 'updatePage']);
expect(pageRepo.updatePage.mock.calls[0][0].lastUpdatedSource).toBe('user');
});
it('git→user transition fires the boundary (git-sync overwrite is a source change)', async () => {
const document = ydocFor(doc('NEW CONTENT'));
pageRepo.findById.mockResolvedValue(pageWithPriorSource('git'));
pageHistoryRepo.findPageLastHistory.mockResolvedValue(null);
await ext.onStoreDocument(buildData(document, 'user') as any);
expect(pageHistoryRepo.saveHistory).toHaveBeenCalledTimes(1);
expect(callOrder).toEqual(['saveHistory', 'updatePage']);
});
it('a null prior source (first-ever edit) does NOT fire the boundary', async () => {
const document = ydocFor(doc('NEW CONTENT'));
pageRepo.findById.mockResolvedValue(pageWithPriorSource(null));
await ext.onStoreDocument(buildData(document, 'agent') as any);
expect(pageHistoryRepo.saveHistory).not.toHaveBeenCalled();
expect(pageRepo.updatePage).toHaveBeenCalledTimes(1);
});
});
it('idempotency: unchanged content → no updatePage, no history, no queues', async () => {
// The Y.Doc content equals the persisted content deeply → early skip.
// A Y.Doc round-trip normalizes attrs (e.g. paragraph indent), so derive
@@ -479,4 +533,231 @@ describe('PersistenceExtension.onStoreDocument — Approach-A boundary snapshot'
// Contributors keyed by the UUID so they match the PAGE_HISTORY job (page.id).
expect(collabHistory.addContributors.mock.calls[0][0]).toBe(PAGE_ID);
});
// #370 — explicit save-version (Cmd+S / agent save tool) over the stateless
// seam. The tier is derived from the SIGNED connection actor, the store path
// is reused, and promote-not-dup avoids duplicating heavy content rows.
describe('save-version (#370)', () => {
const emitSave = (document: any, actor: 'user' | 'agent') =>
ext.onStateless({
connection: {
readOnly: false,
context: { user: { id: USER_ID, name: 'Alice' }, actor },
} as any,
documentName: `page.${PAGE_ID}`,
document: document as any,
payload: JSON.stringify({ type: 'save-version' }),
} as any);
// findById returns a page whose content already equals the live doc, so the
// store path is a no-op and we isolate the versioning decision.
const pageMatchingDoc = (document: any) => ({
...persistedHumanPage('IGNORED'),
content: TiptapTransformer.fromYdoc(document, 'default'),
});
it('human save with no prior snapshot → writes a manual version + broadcasts', async () => {
const document = ydocFor(doc('VERSION ME'));
pageRepo.findById.mockResolvedValue(pageMatchingDoc(document));
pageHistoryRepo.findPageLastHistory.mockResolvedValue(null);
await emitSave(document, 'user');
expect(pageHistoryRepo.saveHistory).toHaveBeenCalledTimes(1);
expect(pageHistoryRepo.saveHistory.mock.calls[0][1]).toEqual(
expect.objectContaining({ kind: 'manual' }),
);
// The pending idle autosnapshot is cancelled by the explicit version.
expect(historyQueue.remove).toHaveBeenCalledWith(PAGE_ID);
const msg = JSON.parse(
(document as any).broadcastStateless.mock.calls[(document as any).broadcastStateless.mock.calls.length - 1][0],
);
expect(msg).toMatchObject({
type: 'version.saved',
kind: 'manual',
alreadySaved: false,
});
});
it('agent save derives kind=agent from the signed actor', async () => {
const document = ydocFor(doc('AGENT VERSION'));
pageRepo.findById.mockResolvedValue(pageMatchingDoc(document));
pageHistoryRepo.findPageLastHistory.mockResolvedValue(null);
await emitSave(document, 'agent');
expect(pageHistoryRepo.saveHistory.mock.calls[pageHistoryRepo.saveHistory.mock.calls.length - 1][1]).toEqual(
expect.objectContaining({ kind: 'agent' }),
);
});
it('promote-not-dup: latest snapshot is an autosave with identical content → upgrades in place', async () => {
const document = ydocFor(doc('SAME'));
const page = pageMatchingDoc(document);
pageRepo.findById.mockResolvedValue(page);
pageHistoryRepo.findPageLastHistory.mockResolvedValue({
id: 'auto-1',
content: page.content,
kind: 'idle',
});
await emitSave(document, 'user');
// No heavy new content row — the existing autosave is promoted to manual.
expect(pageHistoryRepo.updateHistoryKind).toHaveBeenCalledWith(
'auto-1',
'manual',
expect.anything(),
);
expect(pageHistoryRepo.saveHistory).not.toHaveBeenCalled();
const msg = JSON.parse(
(document as any).broadcastStateless.mock.calls[(document as any).broadcastStateless.mock.calls.length - 1][0],
);
expect(msg).toMatchObject({ historyId: 'auto-1', alreadySaved: false });
});
it('no-op when the latest snapshot is already a manual version of this content', async () => {
const document = ydocFor(doc('ALREADY SAVED'));
const page = pageMatchingDoc(document);
pageRepo.findById.mockResolvedValue(page);
pageHistoryRepo.findPageLastHistory.mockResolvedValue({
id: 'ver-1',
content: page.content,
kind: 'manual',
});
await emitSave(document, 'user');
expect(pageHistoryRepo.updateHistoryKind).not.toHaveBeenCalled();
expect(pageHistoryRepo.saveHistory).not.toHaveBeenCalled();
const msg = JSON.parse(
(document as any).broadcastStateless.mock.calls[(document as any).broadcastStateless.mock.calls.length - 1][0],
);
expect(msg).toMatchObject({ alreadySaved: true, kind: 'manual' });
});
it('a read-only connection cannot save a version', async () => {
const document = ydocFor(doc('READER'));
pageRepo.findById.mockResolvedValue(pageMatchingDoc(document));
await ext.onStateless({
connection: {
readOnly: true,
context: { user: { id: USER_ID }, actor: 'user' },
} as any,
documentName: `page.${PAGE_ID}`,
document: document as any,
payload: JSON.stringify({ type: 'save-version' }),
} as any);
expect(pageHistoryRepo.saveHistory).not.toHaveBeenCalled();
expect(pageHistoryRepo.updateHistoryKind).not.toHaveBeenCalled();
});
// #370 F8-twin — a COMMIT abort (serialization/deadlock/conn-drop) rejects
// OUTSIDE the tx callback, AFTER the destructive popContributors (SPOP) and
// saveHistory ran but the INSERT rolled back. onStateless has no retry, so
// the outer catch MUST re-add (SADD) the popped set or attribution is lost
// irrecoverably. MUTATION: drop the outer catch → addContributors is never
// called → this reddens.
it('restores popped contributors when the commit aborts after the callback', async () => {
const document = ydocFor(doc('VERSION ME'));
pageRepo.findById.mockResolvedValue(pageMatchingDoc(document));
// No matching snapshot → fresh version branch → pops contributors.
pageHistoryRepo.findPageLastHistory.mockResolvedValue(null);
collabHistory.popContributors.mockResolvedValue(['u1', 'u2']);
// A db whose commit REJECTS after the callback body resolved: the SPOP and
// saveHistory already ran, then the tx aborts. onStoreDocument's flush uses
// the same db but its content matches (no-op branch) and its own retry loop
// swallows the throw, so only the versioning tx exercises the restore.
const commitFailingDb = {
transaction: () => ({
execute: async (fn: (trx: any) => Promise<any>) => {
await fn(trxStub);
throw new Error('commit aborted (serialization_failure)');
},
}),
};
const ext2 = new PersistenceExtension(
pageRepo as any,
pageHistoryRepo as any,
commitFailingDb as any,
aiQueue as any,
historyQueue as any,
notificationQueue as any,
collabHistory as any,
transclusionService as any,
);
jest.spyOn(ext2['logger'], 'debug').mockImplementation(() => undefined);
jest.spyOn(ext2['logger'], 'warn').mockImplementation(() => undefined);
jest.spyOn(ext2['logger'], 'error').mockImplementation(() => undefined);
await expect(
ext2.onStateless({
connection: {
readOnly: false,
context: { user: { id: USER_ID, name: 'Alice' }, actor: 'user' },
} as any,
documentName: `page.${PAGE_ID}`,
document: document as any,
payload: JSON.stringify({ type: 'save-version' }),
} as any),
).rejects.toThrow();
// Attribution preserved: the popped set is SADD-restored, keyed by the page
// UUID it was popped under.
expect(collabHistory.addContributors).toHaveBeenCalledWith(PAGE_ID, [
'u1',
'u2',
]);
});
// #370 #260 — for a `page.<slugId>` document the idle job is armed under the
// page UUID (computeHistoryJob's jobId = page.id), so the supersede-remove
// must target page.id, not the raw slugId doc-name id, or it silently misses.
it('cancels the superseded idle job by the page UUID for a slugId doc', async () => {
const SLUG = 'slug-1'; // persistedHumanPage.slugId
const document = ydocFor(doc('VERSION ME'));
pageRepo.findById.mockResolvedValue(pageMatchingDoc(document));
pageHistoryRepo.findPageLastHistory.mockResolvedValue(null);
await ext.onStateless({
connection: {
readOnly: false,
context: { user: { id: USER_ID, name: 'Alice' }, actor: 'user' },
} as any,
documentName: `page.${SLUG}`,
document: document as any,
payload: JSON.stringify({ type: 'save-version' }),
} as any);
// remove() keyed by the UUID (the real jobId), never the slugId.
expect(historyQueue.remove).toHaveBeenCalledWith(PAGE_ID);
expect(historyQueue.remove).not.toHaveBeenCalledWith(SLUG);
});
});
// #370 — the in-memory idle-burst marker must be dropped on doc unload (like
// its sibling per-document maps) or it grows unbounded for every page that was
// edited but never manually saved. MUTATION: drop the afterUnloadDocument
// delete → the entry survives → this reddens.
describe('idleBurstStart housekeeping', () => {
it('afterUnloadDocument clears the idle-burst marker armed by a store', async () => {
const document = ydocFor(doc('EDIT'));
pageRepo.findById.mockResolvedValue(persistedHumanPage('EDIT'));
await ext.onStoreDocument(buildData(document, 'user') as any);
const map = ext['idleBurstStart'] as Map<string, number>;
// Keyed by documentName (buildData uses `page.${PAGE_ID}`).
expect(map.has(`page.${PAGE_ID}`)).toBe(true);
await ext.afterUnloadDocument({
documentName: `page.${PAGE_ID}`,
} as any);
expect(map.has(`page.${PAGE_ID}`)).toBe(false);
});
});
});
@@ -37,9 +37,11 @@ import { Page } from '@docmost/db/types/entity.types';
import { CollabHistoryService } from '../services/collab-history.service';
import {
EMBED_DEBOUNCE_MS,
HISTORY_FAST_INTERVAL,
HISTORY_FAST_THRESHOLD,
HISTORY_INTERVAL,
IDLE_INTERVAL_AGENT,
IDLE_INTERVAL_USER,
IDLE_MAX_WAIT_AGENT,
IDLE_MAX_WAIT_USER,
PageHistoryKind,
} from '../constants';
import { TransclusionService } from '../../core/page/transclusion/transclusion.service';
import {
@@ -56,6 +58,16 @@ import { hasTransclusionFamilyNodes } from '../../core/page/transclusion/utils/t
*/
export const INTENTIONAL_CLEAR_MESSAGE_TYPE = 'intentional-clear';
/**
* #370 wire format of the clientserver "save a version" signal. Sent by the
* human (Cmd+S / Save button) and by the agent's explicit save tool over the
* SAME stateless channel. The intentionality tier ('manual' vs 'agent') is
* derived SERVER-SIDE from the signed connection actor, never from this
* payload, so a version's type is unforgeable. The document is taken from the
* connection (not the payload), so the signal cannot be aimed at another page.
*/
export const SAVE_VERSION_MESSAGE_TYPE = 'save-version';
/**
* #251 how long an intentional-clear signal stays "pending" before it is
* ignored. The signal is set on the clearing keystroke but consumed by the
@@ -92,35 +104,39 @@ export function resolveSource(
}
/**
* Compute the BullMQ job id + delay for a page-history snapshot job. Pure so
* the data-loss-sensitive timing arithmetic is unit-testable; `now` is injected
* (caller passes `Date.now()`) for determinism.
* #370 compute the BullMQ job id + delay for a page's trailing idle-flush
* autosnapshot. Pure so the timing is unit-testable.
*
* - Agent edits: delay 0 and a source-keyed job id `${page.id}-agent`. The
* delay MUST stay 0 the worker re-reads the page row at run time, so any
* delay risks reading content a later human edit has already overwritten
* (mis-tagged snapshot). 0 minimizes that window. The `-agent` suffix keeps
* the job from coalescing with the bare-page.id human job.
* - Human edits: age-based debounce so rapid human edits coalesce into one
* snapshot; job id is the bare `page.id`.
*
* BullMQ forbids ':' in custom job ids (Redis key separator), so '-' is used;
* page.id is a UUID, so `${page.id}-agent` cannot collide with a human job.
* Both humans and the agent now share ONE idle pipeline (the agent's old
* `delay=0` fast path is gone intentional agent points arrive via the
* explicit save-version signal instead). The job id is the bare `page.id`, so a
* page has at most one pending idle job; the caller removes-and-re-adds it on
* every store to keep it debounced to the trailing edge of an edit burst. The
* window differs by source only: the agent flushes sooner than a human.
*/
export function computeHistoryJob(
page: Pick<Page, 'id' | 'createdAt'>,
page: Pick<Page, 'id'>,
source: string,
now: number,
// Epoch ms of the FIRST edit in the current burst (when the pending idle job
// was first armed). Used to enforce the max-wait ceiling so a continuous
// editing session cannot re-arm the trailing timer forever. `now` is injectable
// for tests; both default to a live clock / no ceiling when omitted.
burstStart?: number,
now: number = Date.now(),
): { jobId: string; delay: number } {
const isAgent = source === 'agent';
const pageAge = now - new Date(page.createdAt).getTime();
const delay = isAgent
? 0
: pageAge < HISTORY_FAST_THRESHOLD
? HISTORY_FAST_INTERVAL
: HISTORY_INTERVAL;
const jobId = isAgent ? `${page.id}-agent` : page.id;
return { jobId, delay };
const interval = isAgent ? IDLE_INTERVAL_AGENT : IDLE_INTERVAL_USER;
const maxWait = isAgent ? IDLE_MAX_WAIT_AGENT : IDLE_MAX_WAIT_USER;
let delay = interval;
if (burstStart !== undefined) {
// Time already elapsed since the burst's first edit; the snapshot must fire
// no later than `maxWait` after that, so shrink the trailing delay to the
// remaining budget (never negative, so BullMQ fires it promptly).
const remaining = burstStart + maxWait - now;
delay = Math.max(0, Math.min(interval, remaining));
}
return { jobId: page.id, delay };
}
@Injectable()
@@ -132,6 +148,28 @@ export class PersistenceExtension implements Extension {
// coalescing window" per document and OR it across all edits in the window,
// so the snapshot is marked 'agent' regardless of who wrote last.
private agentTouched: Map<string, boolean> = new Map();
// #370 — epoch ms of the FIRST edit in the current idle-flush burst. Keyed by
// documentName (like its sibling per-document maps above), NOT by page.id, so
// it can be cleaned in afterUnloadDocument alongside `contributors` /
// `agentTouched` / `intentionalClear` when the doc unloads — otherwise any page
// that was edited but never manually saved (the common case) would keep its
// entry forever and the Map would grow unbounded in this long-lived process.
// Set when the pending idle job is first armed (empty entry), read to enforce
// the max-wait ceiling in computeHistoryJob, and cleared on doc unload or when
// a manual save cancels the idle job so the next burst starts a fresh window.
//
// Single-process assumption (like `contributors` / `agentTouched` above): this
// lives only in THIS collab process's memory. A restart, or a page's ownership
// moving to another node, loses the burst-start marker. Consequence: a burst
// that spans the restart looks like a fresh burst to the surviving process, so
// its max-wait ceiling is re-anchored to the first post-restart edit — a single
// continuous session straddling a restart can therefore wait up to ~2× the cap
// for its idle snapshot (once for the lost pre-restart window, once for the new
// one). Bounded and benign (it only DELAYS a safety-net autosnapshot; manual
// saves are unaffected and the next quiet period always flushes), but the
// assumption and its consequence are recorded here so no one mistakes the
// in-memory marker for a durable, cross-process guarantee.
private idleBurstStart: Map<string, number> = new Map();
// #251 — per-document "intentional clear pending" flags. Keyed by
// documentName, value = expiry timestamp (ms). Set by onStateless when the
// client reports a deliberate clear; consumed once by the next
@@ -363,20 +401,19 @@ export class PersistenceExtension implements Extension {
//this.logger.debug('Contributors error:' + err?.['message']);
}
// Approach A — boundary snapshot before the agent's first edit.
// When this store is the agent's and the page's currently persisted
// state was authored by a human, pin that human state as its own
// history version BEFORE the agent overwrites it. `page` still holds
// the OLD content/provenance here, so saveHistory(page) captures the
// pre-agent state tagged 'user'. The agent's new content is
// snapshotted later by the debounced PAGE_HISTORY job ('agent'). Skip
// if the prior state is already agent-authored (boundary already
// pinned on the user->agent transition), if the page is effectively
// empty, or if the latest existing snapshot already equals this human
// state (avoid duplicates).
// #370 — boundary snapshot on ANY source transition. When the store
// flips the page's provenance (user↔agent↔git), pin the OUTGOING
// state as its own history version BEFORE the incoming source
// overwrites it. `page` still holds the OLD content/provenance here,
// so saveHistory(page) captures the pre-transition state tagged with
// its own source, kind='boundary'. The incoming content is snapshotted
// later by the debounced idle job. Skip if the page is effectively
// empty or if the latest existing snapshot already equals this state
// (the shared isDeepStrictEqual gate — avoids duplicates). Generalizing
// beyond the old user→agent special-case also covers git-sync for free.
if (
lastUpdatedSource === 'agent' &&
page.lastUpdatedSource !== 'agent'
page.lastUpdatedSource &&
page.lastUpdatedSource !== lastUpdatedSource
) {
// pageHistory.pageId is uuid-typed; use page.id (never the doc-name
// slugId) so a `page.<slugId>` doc cannot throw 22P02 here (#260).
@@ -384,15 +421,13 @@ export class PersistenceExtension implements Extension {
page.id,
{ includeContent: true, trx },
);
const humanBaselineMissing =
const baselineMissing =
!lastHistory ||
!isDeepStrictEqual(lastHistory.content, page.content);
if (
!isEmptyParagraphDoc(page.content as any) &&
humanBaselineMissing
) {
if (!isEmptyParagraphDoc(page.content as any) && baselineMissing) {
await this.pageHistoryRepo.saveHistory(page, {
contributorIds: page.contributorIds ?? undefined,
kind: 'boundary',
trx,
});
}
@@ -522,7 +557,7 @@ export class PersistenceExtension implements Extension {
{ jobId: `embed-${page.id}`, delay: EMBED_DEBOUNCE_MS },
);
await this.enqueuePageHistory(page, lastUpdatedSource);
await this.enqueuePageHistory(page, documentName, lastUpdatedSource);
}
// #402 — report the serialized size for the store histogram's size_bucket.
@@ -554,6 +589,14 @@ export class PersistenceExtension implements Extension {
return; // unrelated / malformed stateless message
}
// #370 — explicit "save a version" (human Cmd+S / agent save tool). Edit
// rights are already enforced by the readOnly reject above (a reader can't
// create a version), exactly as intentional-clear requires.
if (message?.type === SAVE_VERSION_MESSAGE_TYPE) {
await this.handleSaveVersion(data);
return;
}
if (message?.type !== INTENTIONAL_CLEAR_MESSAGE_TYPE) return;
this.intentionalClear.set(
@@ -562,6 +605,160 @@ export class PersistenceExtension implements Extension {
);
}
/**
* #370 persist an intentional version from the live in-memory ydoc.
*
* One stateless path serves BOTH the human and the agent; the tier is derived
* SERVER-SIDE from the signed connection actor ('agent' 'agent', anything
* else 'manual'), so the version type cannot be spoofed by the client. We
* take the fresh ydoc from the collab process memory and run it through the
* EXISTING store path first (so pages.content/ydoc reflect the exact content
* being versioned a REST endpoint would race the up-to-10s-stale page row),
* then snapshot it into page_history with the intentional kind.
*
* Promote-not-dup: if the latest history row already holds this exact content
* and it is an autosave (idle/boundary/legacy-null), upgrade its kind in place
* instead of duplicating a heavy content row; if it is already 'manual', it is
* a no-op (the client shows an "already saved" toast). Otherwise a fresh
* version row is written, popping the aggregated contributors from Redis.
*/
private async handleSaveVersion(data: onStatelessPayload): Promise<void> {
const { connection, document, documentName } = data;
const context = connection?.context;
const pageId = getPageId(documentName);
// Unforgeable: 'agent' only for a signed agent connection, else 'manual'.
const kind: PageHistoryKind =
context?.actor === 'agent' ? 'agent' : 'manual';
// Flush the live ydoc through the normal store path so the page row + ydoc
// hold exactly what we are about to version (also fires the idle enqueue we
// supersede below, plus any source-transition boundary). onStoreDocument
// only needs document/documentName/context.
await this.onStoreDocument({
document,
documentName,
context,
} as onStoreDocumentPayload);
let result:
| { historyId: string; kind: PageHistoryKind; alreadySaved: boolean }
| undefined;
// #370 F8-twin — the contributor set popped from Redis (destructive SPOP)
// must be restored if the version row does not durably land. The inner
// try/catch below only covers a throw INSIDE the callback; but executeTx
// COMMITS after the callback, so a commit-abort (serialization/deadlock/
// connection drop — the transient class the epic retries in the processor)
// rejects OUTSIDE the callback, after saveHistory already ran and the SPOP
// already happened, while the INSERT rolls back. onStateless does NOT retry,
// so an unrestored pop is a one-shot irrecoverable attribution loss (the
// processor got exactly this fix: poppedForRestore + an outer catch). We
// track the popped set here (keyed by the page UUID it was popped by — never
// the doc-name id, which may be a slugId, #260) and restore it in the outer
// catch. addContributors is an idempotent Redis SADD, so a double-restore is
// harmless. versionedPageId is also reused below to remove the superseded
// idle job by its real jobId (page.id).
let poppedForRestore: string[] = [];
let versionedPageId: string | undefined;
try {
await executeTx(this.db, async (trx) => {
const page = await this.pageRepo.findById(pageId, {
withLock: true,
includeContent: true,
trx,
});
if (!page) return;
versionedPageId = page.id;
// Never version an effectively-empty page (mirrors the processor's
// first-history guard); there is nothing intentional to pin.
if (isEmptyParagraphDoc(page.content as any)) return;
const lastHistory = await this.pageHistoryRepo.findPageLastHistory(
page.id,
{ includeContent: true, trx },
);
if (
lastHistory &&
isDeepStrictEqual(lastHistory.content, page.content)
) {
// Content is already snapshotted. Promote-not-dup.
if (lastHistory.kind === 'manual') {
result = {
historyId: lastHistory.id,
kind: 'manual',
alreadySaved: true,
};
return;
}
await this.pageHistoryRepo.updateHistoryKind(
lastHistory.id,
kind,
trx,
);
result = { historyId: lastHistory.id, kind, alreadySaved: false };
return;
}
// Fresh version row. Pop the contributors aggregated since the last
// snapshot (SPOP); restore them if the write fails so they aren't lost.
const contributorIds = await this.collabHistory.popContributors(
page.id,
);
poppedForRestore = contributorIds;
try {
const saved = await this.pageHistoryRepo.saveHistory(page, {
contributorIds,
kind,
trx,
});
result = { historyId: saved.id, kind, alreadySaved: false };
} catch (err) {
await this.collabHistory.addContributors(page.id, contributorIds);
poppedForRestore = [];
throw err;
}
});
} catch (err) {
// A throw here means the tx did NOT commit (callback threw, or the commit
// itself failed and rolled back). If we popped contributors and the inner
// catch did not already restore them, restore now so attribution is not
// lost — onStateless has no retry to recover it. Restore by the page UUID
// the pop was keyed under (versionedPageId is always set before the pop).
if (poppedForRestore.length && versionedPageId) {
await this.collabHistory.addContributors(
versionedPageId,
poppedForRestore,
);
}
throw err;
}
// Housekeeping: this explicit version supersedes the page's pending idle
// autosnapshot, so cancel it and end the current idle burst so the next edit
// starts a fresh max-wait window. Remove the idle job by its REAL jobId
// (page.id UUID — computeHistoryJob arms it under page.id), not the raw
// doc-name id which may be a slugId for a `page.<slugId>` doc (#260), or the
// remove silently misses. The burst marker is keyed by documentName (like its
// sibling per-document maps), and is also cleaned in afterUnloadDocument.
if (versionedPageId) {
await this.historyQueue.remove(versionedPageId).catch(() => undefined);
}
this.idleBurstStart.delete(documentName);
if (result) {
document.broadcastStateless(
JSON.stringify({
type: 'version.saved',
historyId: result.historyId,
kind: result.kind,
alreadySaved: result.alreadySaved,
}),
);
}
}
async onChange(data: onChangePayload) {
const documentName = data.documentName;
const userId = data.context?.user?.id;
@@ -586,6 +783,10 @@ export class PersistenceExtension implements Extension {
this.contributors.delete(documentName);
this.agentTouched.delete(documentName);
this.intentionalClear.delete(documentName);
// #370 — drop the idle-burst marker with the other per-document maps so it
// cannot accumulate across the process lifetime for never-manually-saved
// pages. The pending idle job (if any) is a self-expiring BullMQ delayed job.
this.idleBurstStart.delete(documentName);
}
private consumeContributors(documentName: string): string[] {
@@ -617,19 +818,80 @@ export class PersistenceExtension implements Extension {
private async enqueuePageHistory(
page: Page,
documentName: string,
lastUpdatedSource: string,
): Promise<void> {
// Job id + delay arithmetic lives in the pure `computeHistoryJob` (see its
// doc comment for the agent-delay-0 / age-based-debounce invariants).
// #370 — trailing idle debounce with a max-wait ceiling. One pending idle
// job per page (jobId = page.id); on every store we remove the pending
// delayed job and re-add it, so the snapshot lands `delay` after edits go
// quiet rather than once per store (precedent: workspace.service.ts).
// remove() on a delayed job simply deletes it (0 if absent, no throw); if the
// job is already ACTIVE and the remove is a no-op, the add still de-dups and
// the processor's isDeepStrictEqual gate collapses the duplicate content.
//
// The FIRST arm of a burst records `burstStart`; computeHistoryJob shrinks
// the delay to the remaining max-wait budget from that point, so a continuous
// session cannot re-arm the trailing timer forever and starve the snapshot.
// A burst marker older than THIS TIER's max-wait means the previous idle job
// has already fired — start a fresh window instead of firing immediately on
// the next edit. Must use the SAME source-specific max-wait computeHistoryJob
// uses (agent 5m / user 10m): a hardcoded USER ceiling would leave an agent
// burst's marker stale for 5..10m, forcing delay=0 on every store in that
// window and writing one idle row per store — exactly the per-store bloat the
// debounce exists to prevent, on the continuous-agent path.
const maxWait =
lastUpdatedSource === 'agent' ? IDLE_MAX_WAIT_AGENT : IDLE_MAX_WAIT_USER;
const now = Date.now();
// Keyed by documentName (see the map declaration) so afterUnloadDocument can
// clean it; the queue jobId stays page.id (computeHistoryJob) as required.
let burstStart = this.idleBurstStart.get(documentName);
if (burstStart === undefined || now - burstStart >= maxWait) {
burstStart = now;
this.idleBurstStart.set(documentName, burstStart);
}
const { jobId, delay } = computeHistoryJob(
page,
lastUpdatedSource,
Date.now(),
burstStart,
now,
);
// remove-then-add trailing-debounce idiom, and its ONE race. We delete the
// pending delayed job and re-add it under the same jobId so the timer resets
// to the trailing edge of the burst. The race is the small window between
// these two awaits: if the delayed job's `delay` elapses in that gap it goes
// ACTIVE, and then:
// - remove() on an active/locked job is a no-op (BullMQ won't yank a job a
// worker holds), and our `.catch(() => undefined)` swallows that too; and
// - add() with a jobId that already exists (the now-active job's id) is
// DROPPED by BullMQ — a duplicate add is a no-op.
// So this store fails to re-arm the trailing job: the just-fired snapshot
// captured content up to the moment it went active, and THIS edit is left
// without a pending trailing job. It is bounded and self-healing — the NEXT
// store re-arms a fresh delayed job (the id is free again once the active job
// completes / removeOnComplete frees it), and the processor's
// isDeepStrictEqual gate collapses any content-identical duplicate. The only
// uncovered case is when the racing store was the LAST in the session: the
// tail edits made after the job went active get NO trailing snapshot until
// the next edit re-arms one. That is an acceptable safety-net gap (a manual
// Save, a source-transition boundary, or simply the next edit all still cover
// it), which is why the reviewer accepts documenting it here rather than
// adding a post-add "did the add actually arm a job?" re-check.
//
// NOTE — do NOT "unify" this with the neighbouring embed-debounce idiom
// (aiQueue.add of PAGE_CONTENT_UPDATED above): that one uses a STABLE jobId
// and NO remove(), relying purely on BullMQ coalescing a repeated add under
// the same id, because a re-embed only needs to eventually run once on the
// latest content and re-anchoring its delay on every keystroke is undesirable.
// THIS idiom deliberately removes-then-adds precisely to PUSH the delay back
// to the trailing edge on every store (a true debounce), which coalescing
// alone cannot do. Collapsing them would silently change the history cadence.
await this.historyQueue.remove(jobId).catch(() => undefined);
await this.historyQueue.add(
QueueJob.PAGE_HISTORY,
{ pageId: page.id } as IPageHistoryJob,
{ pageId: page.id, kind: 'idle' } as IPageHistoryJob,
{ jobId, delay },
);
}
@@ -66,6 +66,15 @@ describe('HistoryProcessor.process', () => {
notificationQueue = { add: jest.fn().mockResolvedValue(undefined) };
generalQueue = { add: jest.fn().mockResolvedValue(undefined) };
// #370 F3 — the processor now serializes its find+save under a page-row lock
// via executeTx. A db whose transaction().execute(fn) runs fn with a trx stub
// drives the real executeTx() helper without a database.
const db = {
transaction: () => ({
execute: (fn: (trx: any) => Promise<any>) => fn({ __trx: true }),
}),
};
// WorkerHost's constructor reads `this.worker`; passing repos positionally
// matches the constructor and avoids the Nest DI container.
proc = new HistoryProcessor(
@@ -73,6 +82,7 @@ describe('HistoryProcessor.process', () => {
pageRepo as any,
collabHistory as any,
watcherService as any,
db as any,
notificationQueue as any,
generalQueue as any,
);
@@ -126,15 +136,26 @@ describe('HistoryProcessor.process', () => {
await proc.process(buildJob());
expect(collabHistory.popContributors).toHaveBeenCalledWith(PAGE_ID);
// #370 F3/F9 — the snapshot decision runs under a page-row lock. Pin the lock
// structurally so a refactor that drops withLock/trx (silently reintroducing
// the TOCTOU double-insert) turns this red. The tx stub is { __trx: true }.
expect(pageRepo.findById).toHaveBeenCalledWith(
PAGE_ID,
expect.objectContaining({ withLock: true, trx: { __trx: true } }),
);
// #370 F7 — addPageWatchers MUST receive the trx, or its FK-check runs on a
// separate connection and self-deadlocks against our FOR UPDATE. Asserting
// the trx arg here is exactly what would have caught that regression.
expect(watcherService.addPageWatchers).toHaveBeenCalledWith(
['u1', 'u2'],
PAGE_ID,
SPACE_ID,
WORKSPACE_ID,
{ __trx: true },
);
expect(pageHistoryRepo.saveHistory).toHaveBeenCalledWith(
expect.objectContaining({ id: PAGE_ID }),
{ contributorIds: ['u1', 'u2'] },
{ contributorIds: ['u1', 'u2'], kind: 'idle', trx: { __trx: true } },
);
expect(generalQueue.add).toHaveBeenCalledWith(
QueueJob.PAGE_BACKLINKS,
@@ -186,6 +207,48 @@ describe('HistoryProcessor.process', () => {
]);
});
it('COMMIT failure (throw outside the tx callback) → contributors RESTORED', async () => {
// #370 F8 — a commit-time failure throws OUTSIDE the callback, so the inner
// try/catch does not run; the outer catch must restore the popped set (else a
// BullMQ retry writes an unattributed version). Use a db whose execute() runs
// the callback THEN throws, simulating a commit abort.
pageHistoryRepo.findPageLastHistory.mockResolvedValue({
content: { type: 'doc', content: [] },
});
const commitFail = {
transaction: () => ({
execute: async (fn: (trx: any) => Promise<any>) => {
await fn({ __trx: true }); // callback succeeds (saveHistory ok)
throw new Error('commit aborted'); // ...but the COMMIT fails
},
}),
};
const procCommitFail = new HistoryProcessor(
pageHistoryRepo as any,
pageRepo as any,
collabHistory as any,
watcherService as any,
commitFail as any,
notificationQueue as any,
generalQueue as any,
);
jest
.spyOn(procCommitFail['logger'], 'error')
.mockImplementation(() => undefined);
await expect(procCommitFail.process(buildJob())).rejects.toThrow(
'commit aborted',
);
// The inner catch did NOT run (save succeeded), so only the outer catch can
// restore — assert it did.
expect(collabHistory.addContributors).toHaveBeenCalledWith(PAGE_ID, [
'u1',
'u2',
]);
// And the post-snapshot queue work must NOT have run (we rethrew).
expect(generalQueue.add).not.toHaveBeenCalled();
});
it('backlinks + notification queue failures are swallowed (history still committed)', async () => {
pageHistoryRepo.findPageLastHistory.mockResolvedValue({
content: { type: 'doc', content: [] },
@@ -19,6 +19,9 @@ import { isDeepStrictEqual } from 'node:util';
import { CollabHistoryService } from '../services/collab-history.service';
import { WatcherService } from '../../core/watcher/watcher.service';
import { isEmptyParagraphDoc } from '../collaboration.util';
import { InjectKysely } from 'nestjs-kysely';
import { KyselyDB } from '@docmost/db/types/kysely.types';
import { executeTx } from '@docmost/db/utils';
@Processor(QueueName.HISTORY_QUEUE)
export class HistoryProcessor extends WorkerHost implements OnModuleDestroy {
@@ -29,6 +32,7 @@ export class HistoryProcessor extends WorkerHost implements OnModuleDestroy {
private readonly pageRepo: PageRepo,
private readonly collabHistory: CollabHistoryService,
private readonly watcherService: WatcherService,
@InjectKysely() private readonly db: KyselyDB,
@InjectQueue(QueueName.NOTIFICATION_QUEUE) private notificationQueue: Queue,
@InjectQueue(QueueName.GENERAL_QUEUE) private generalQueue: Queue,
) {
@@ -41,6 +45,9 @@ export class HistoryProcessor extends WorkerHost implements OnModuleDestroy {
try {
const { pageId } = job.data;
// Read the page WITHOUT a lock first, only to bail early on the two cheap
// no-write cases (page gone / empty first snapshot) without opening a
// transaction. The authoritative check-then-write happens locked below.
const page = await this.pageRepo.findById(pageId, {
includeContent: true,
});
@@ -51,40 +58,109 @@ export class HistoryProcessor extends WorkerHost implements OnModuleDestroy {
return;
}
const lastHistory = await this.pageHistoryRepo.findPageLastHistory(
pageId,
{ includeContent: true },
);
// #370 F3 — the snapshot decision (findPageLastHistory → saveHistory) must
// be serialized against manual-save/boundary writers, which run under a
// page-row lock in onStoreDocument. Without it, this processor and a
// concurrent manual-save each read the same lastHistory (MVCC), both see
// content != lastHistory, and both insert — producing two page_history rows
// with IDENTICAL content (one 'idle', one 'manual'), defeating
// promote-not-dup and the version-vs-autosave split. Taking the same
// page-row lock makes the second writer observe the first's committed row so
// the isDeepStrictEqual gate collapses the duplicate. Only the read+write
// is transacted; the post-snapshot queue work stays outside.
let contributorIds: string[] = [];
let snapshotWritten = false;
let lastHistoryContent: unknown;
// #370 F8 — the contributor set popped from Redis (destructive SPOP) must be
// restored if the snapshot does not durably land. The inner try/catch only
// covers a throw INSIDE the callback; a COMMIT failure (connection drop,
// serialization/deadlock abort on commit — the transient class the epic
// already retries) throws OUTSIDE it, rolling the snapshot back while the
// pop is already gone. We track the popped set here and restore it in the
// outer catch so a BullMQ retry re-attributes the version. addContributors
// is an idempotent Redis SADD, so a double-restore is harmless.
let poppedForRestore: string[] = [];
if (!lastHistory && isEmptyParagraphDoc(page.content as any)) {
this.logger.debug(
`Skipping first history for page ${pageId}: empty content`,
);
await this.collabHistory.clearContributors(pageId);
try {
await executeTx(this.db, async (trx) => {
const lockedPage = await this.pageRepo.findById(pageId, {
includeContent: true,
withLock: true,
trx,
});
if (!lockedPage) return;
const lastHistory = await this.pageHistoryRepo.findPageLastHistory(
pageId,
{ includeContent: true, trx },
);
lastHistoryContent = lastHistory?.content;
if (!lastHistory && isEmptyParagraphDoc(lockedPage.content as any)) {
this.logger.debug(
`Skipping first history for page ${pageId}: empty content`,
);
return;
}
if (
lastHistory &&
isDeepStrictEqual(lastHistory.content, lockedPage.content)
) {
return; // already snapshotted at this content — nothing to write
}
contributorIds = await this.collabHistory.popContributors(pageId);
poppedForRestore = contributorIds;
try {
// Pass `trx` so the watcher insert's FK check (FOR KEY SHARE on
// pages[pageId]) runs on the SAME connection that already holds the
// FOR UPDATE lock from findById — otherwise it takes the FK lock on a
// separate pool connection and self-deadlocks against our own tx.
await this.watcherService.addPageWatchers(
contributorIds,
pageId,
lockedPage.spaceId,
lockedPage.workspaceId,
trx,
);
// #370 — every job on this queue is a trailing idle-flush autosnapshot.
await this.pageHistoryRepo.saveHistory(lockedPage, {
contributorIds,
kind: job.data.kind ?? 'idle',
trx,
});
snapshotWritten = true;
this.logger.debug(`History created for page: ${pageId}`);
} catch (err) {
await this.collabHistory.addContributors(pageId, contributorIds);
poppedForRestore = [];
throw err;
}
});
} catch (err) {
// A throw here means the tx did NOT commit (callback threw, or the commit
// itself failed and rolled back). If we popped contributors and the inner
// catch did not already restore them, restore now so the retry keeps
// attribution. snapshotWritten is irrelevant: it is set before commit, so
// it can be true even when the commit rolled the snapshot back.
if (poppedForRestore.length) {
await this.collabHistory.addContributors(pageId, poppedForRestore);
}
throw err;
}
// No snapshot written (page vanished / empty-first / unchanged content) →
// clear the contributor set for the skip cases and stop.
if (!snapshotWritten) {
if (!lastHistoryContent && isEmptyParagraphDoc(page.content as any)) {
await this.collabHistory.clearContributors(pageId);
}
return;
}
if (
!lastHistory ||
!isDeepStrictEqual(lastHistory.content, page.content)
) {
const contributorIds = await this.collabHistory.popContributors(pageId);
try {
await this.watcherService.addPageWatchers(
contributorIds,
pageId,
page.spaceId,
page.workspaceId,
);
await this.pageHistoryRepo.saveHistory(page, { contributorIds });
this.logger.debug(`History created for page: ${pageId}`);
} catch (err) {
await this.collabHistory.addContributors(pageId, contributorIds);
throw err;
}
{
const mentions = extractMentions(page.content);
const pageMentions = extractPageMentions(mentions);
const internalLinkSlugIds = extractInternalLinkSlugIds(page.content);
@@ -102,7 +178,7 @@ export class HistoryProcessor extends WorkerHost implements OnModuleDestroy {
);
});
if (contributorIds.length > 0 && lastHistory?.content) {
if (contributorIds.length > 0 && lastHistoryContent) {
await this.notificationQueue
.add(QueueJob.PAGE_UPDATED, {
pageId,
@@ -1,42 +1,122 @@
import { Injectable, Logger, OnModuleDestroy } from '@nestjs/common';
/**
* In-memory run-stream registry (#184 phase 1.5). A durable agent run tees its
* SSE frames here (via `pipeUIMessageStreamToResponse({ consumeSseStream })`)
* so a LATE tab one that reloaded, or opened after the starter dropped can
* attach through `GET /ai-chat/runs/:chatId/stream`, replay the frames buffered
* so far, and then follow the live tail as a normal streamer.
* In-memory run-stream registry (#184 phase 1.5, step-aligned retention #491). A
* durable agent run tees its SSE frames here (via
* `pipeUIMessageStreamToResponse({ consumeSseStream })`) so a LATE tab one that
* reloaded, or opened after the starter dropped can attach through
* `GET /ai-chat/runs/:chatId/stream`, be handed the TAIL past the step it already
* has persisted, and then follow the live tail as a normal streamer.
*
* This is deliberately single-process and best-effort: it holds nothing the DB
* does not (the run + assistant row are the source of truth), so a process
* restart simply drops in-flight entries and the client falls back to its
* restore + degraded-poll path. The async `attach` return type is the seam for a
* future phase-2 cross-process backend (Redis) the interface does not change.
*
* #491 step-aligned retention (the OOM fix)
* The old registry buffered up to 32MB of raw SSE frames PER active run (V8 ~2×
* in memory) and, on attach, blasted the WHOLE buffer to the socket synchronously
* with no drain a handful of marathon runs on a 1GB container OOM'd. #491 caps
* the ring at a few MB (env-tunable, default 4MB) and keeps it there by ROTATING:
*
* - Every buffered frame is STAMPED with a step number at tee (see ingestFrame).
* Convention: the stamp of a frame is the number of `finish-step` parts seen
* BEFORE it (starting at 0). The finish-step frame itself carries the current
* value, THEN the counter increments. So a frame stamped `s` is the content of
* the (s+1)-th step 0-based step index `s` and the stamp aligns EXACTLY
* with `metadata.stepsPersisted`: a client whose persisted `stepsPersisted` is
* N has steps 0..N-1 on disk (and in its seed) and needs the tail `stamp >= N`.
*
* - The ring rotates ONLY on a CONFIRMED persist of step N
* (`confirmPersistedStep`), dropping frames with `stamp < N` (those steps are
* now on disk and a fresh client seed carries them). A NON-confirmed step is
* never rotated away, so a persist FAILURE just makes the ring cover MORE
* (auto-safe). This is the anti-inversion rule: a naive "rotate in .then()"
* that rotated after an UNwritten step would drop a step nobody has silent
* hole. Rotation is gated on a real, successful persist.
*
* - If the ring still exceeds its byte cap after rotation (a single fat step, or
* a lagging persist), the OLDEST frames are evicted to stay bounded. Evicting a
* not-yet-persisted frame opens a GAP: an attach whose N falls at or below an
* evicted step answers 204 and the client degrades to restore+poll. The gap is
* NOT sticky the coverage floor is recomputed from the ring, so a later
* persist that rotates past the holey steps clears it.
*
* attach numbering / coverage (the wire convention)
* The step marker N comes ONLY FROM THE CLIENT (a query param). The server never
* reads the row to derive N a server-side N from a stale seed would open a
* silent one-step hole. N is the client's persisted `stepsPersisted` (a COUNT):
* - the tail it needs = frames with `stamp >= N`;
* - coverage is OK `coverageFloor(entry) <= N`, where coverageFloor is the
* smallest step FULLY present in the ring (its smallest retained stamp, bumped
* by one when that leading step was only partially evicted by overflow). If
* `coverageFloor > N` the ring starts AFTER the client's frontier (a hole, or
* the client's seed simply lagged behind a rotation) 204 the client
* refetches (a larger N) and re-attaches.
* The N cutoff is applied in ALL branches, INCLUDING the finished-retained replay.
*
* same-tick invariants (unchanged, still load-bearing)
* invariant 1: only the matching run may mutate/observe an entry (runId check).
* invariant 2: retention deletes ONLY its own entry (a replacement may own the key).
* invariant 3: open() over a live entry mirrors the done-path (subscribers released).
* invariant 4: the tail SLICE + subscriber registration happen in ONE synchronous
* tick inside attach() no await between them so a concurrently
* ingested frame is EITHER in the snapshot (buffered before the sync
* block, and the just-added subscriber never sees it) OR fanned out to
* the paused subscriber's `pending` (ingested after) never both and
* never neither: no loss, no duplication. NOTE (#491): the controller
* now AWAITS the drain-respecting tail write BEFORE calling start(), so
* frames ingested during that await accumulate in `pending`; this is
* bounded by the subscriber cap (an overflow degrades start() to an
* end(), a 204-equivalent). It is the SYNCHRONOUS snapshot+registration
* not a same-tick start() that makes this correct.
* invariant 5: the controller wires close-cleanup BEFORE any write.
* invariant 6: no cross-run replay the `anchor` (the client's assistant row id)
* must match this run's assistant id, or a foreign run's transcript
* would be appended to the client's message.
*/
/** How long a finished entry is retained for late attach (replay + immediate end). */
export const RUN_STREAM_RETAIN_FINISHED_MS = 30_000;
/**
* Per-run replay buffer cap. Past this the buffer is dropped (attach -> 204, and
* the client falls back to its restore + degraded-poll path, #430).
*
* Raised from 4MB to 32MB (#430): marathon autonomous runs (11-25 min observed)
* stream far more than 4MB of SSE frames, so a live disconnect mid-run would find
* an already-overflowed buffer and could only degrade-poll instead of re-attaching
* to the live tail. 32MB comfortably covers those runs while staying bounded.
*
* Memory cost: this is the WORST-CASE retained size PER ACTIVE run (the buffer is
* freed on finish + retention, or dropped immediately on overflow). With the small
* number of concurrent autonomous runs a single workspace realistically has, 32MB
* each is an acceptable ceiling; the overflow->204->degraded-poll fallback remains
* the backstop for anything larger, so correctness never depends on this bound.
* DEFAULT per-run replay ring cap (#491, down from 32MB). SSE frames carry
* UNcompacted tool outputs + framing overhead (×1.52 vs the persisted parts), so
* a "2–3 large reads + reasoning" step routinely blows past 2MB; 4MB comfortably
* holds a step or two of TAIL, which is all a resuming client needs (steps below
* its persisted frontier come from the seed, not the ring). The ring stays bounded
* because it rotates on every confirmed persist; this cap is only the ceiling for
* the un-persisted tail between rotations. Env-tunable via
* AI_CHAT_RUN_STREAM_MAX_BUFFER_BYTES (bytes); a 0/invalid value falls back to this.
*/
export const RUN_STREAM_MAX_BUFFER_BYTES = 32 * 1024 * 1024;
export const AI_CHAT_RUN_STREAM_MAX_BUFFER_BYTES = 4 * 1024 * 1024;
// 2x the replay cap: a just-written full-replay burst alone can never trip the
// per-subscriber cap (see controller); only a genuinely stalled socket can.
export const SUBSCRIBER_MAX_BUFFERED_BYTES = 2 * RUN_STREAM_MAX_BUFFER_BYTES;
// 2× the ring cap: a just-written full-tail burst alone can never trip the
// per-subscriber cap (see controller); only a genuinely stalled socket can. This
// derivative relationship is preserved even when the ring cap is env-overridden.
export const SUBSCRIBER_MAX_BUFFERED_BYTES = 2 * AI_CHAT_RUN_STREAM_MAX_BUFFER_BYTES;
/**
* A finish-step boundary frame is exactly `data: {"type":"finish-step"...}\n\n`
* (verified empirically against ai@6.0.207 each UI-message-stream part is a
* single `data: {json}\n\n` event, never split across `data:` lines, and `type`
* is always the first key). A prefix match is cheaper than JSON.parse-per-frame
* and has no false positives: a literal `"type":"finish-step"` inside a text
* delta is JSON-escaped (`\"type\":...`), and the frame would start with
* `data: {"type":"text-delta"` anyway.
*/
const FINISH_STEP_FRAME_PREFIX = 'data: {"type":"finish-step"';
/** Resolve the ring cap from the environment, falling back to the default. */
function resolveMaxBufferBytes(): number {
const raw = process.env.AI_CHAT_RUN_STREAM_MAX_BUFFER_BYTES;
if (!raw) return AI_CHAT_RUN_STREAM_MAX_BUFFER_BYTES;
const parsed = Number(raw);
return Number.isFinite(parsed) && parsed > 0
? Math.floor(parsed)
: AI_CHAT_RUN_STREAM_MAX_BUFFER_BYTES;
}
export interface RunStreamCallbacks {
onFrame: (frame: string) => void;
@@ -44,6 +124,9 @@ export interface RunStreamCallbacks {
}
export interface RunStreamAttachment {
// The synthetic `start` frame (carrying { runId, chatId }) followed by the
// buffered TAIL filtered to `stamp >= N`. The controller writes these to the
// socket in chunks respecting drain, then calls start().
replay: string[];
finished: boolean;
start(): void; // drain pending frames (order preserved) and go live
@@ -53,14 +136,19 @@ export interface RunStreamAttachment {
interface Subscriber extends RunStreamCallbacks {
started: boolean;
pending: string[];
// Byte size of `pending`, capped at SUBSCRIBER_MAX_BUFFERED_BYTES. `start()` is
// called in the SAME tick as `attach()` today (see attach), so `pending` never
// holds more than one microtask of frames — but the async `attach` signature is
// a phase-2 seam: an await between attach and start would let a stalled paused
// subscriber buffer the WHOLE run here. The cap is the structural backstop.
// Byte size of `pending`, capped at the subscriber cap. `start()` is called in
// the SAME tick as `attach()` today, so `pending` never holds more than one
// microtask of frames — but the controller writes the (potentially large) tail
// respecting drain BEFORE start(), so a stalled socket can accumulate here; the
// cap is the structural backstop (an overflow degrades start() to an end()).
pendingBytes: number;
overflowed: boolean;
pendingEnd: boolean;
// The client's step frontier N: this subscriber only receives frames with
// `stamp >= minStamp` (the tail past what it already persisted). Live frames
// always satisfy this (their stamp is the current, highest step), so it only
// filters the rare out-of-order below-frontier frame.
minStamp: number;
}
interface Entry {
@@ -68,8 +156,20 @@ interface Entry {
// The persisted assistant row id of this run (set at bind; undefined if the
// seed failed). Used by the attach anchor check (invariant 6).
assistantMessageId?: string;
// Parallel arrays: frames[i] is the SSE string, stamps[i] its step number.
frames: string[];
stamps: number[];
bytes: number;
// The running step counter used to stamp the NEXT frame (number of finish-step
// frames seen so far).
currentStamp: number;
// The highest confirmed `stepsPersisted`: frames with stamp < persistedFloor are
// on disk (safe to drop, never re-buffered). Monotonic (confirmPersistedStep).
persistedFloor: number;
// The highest stamp EVICTED by an overflow (unsafe) drop, -1 if none. Used to
// detect a partially-evicted leading step when computing the coverage floor.
overflowThroughStamp: number;
// Sticky-for-logging only: at least one unsafe (overflow) eviction happened.
overflowed: boolean;
finished: boolean;
subscribers: Set<Subscriber>;
@@ -80,6 +180,10 @@ interface Entry {
export class AiChatStreamRegistryService implements OnModuleDestroy {
private readonly logger = new Logger(AiChatStreamRegistryService.name);
private readonly entries = new Map<string, Entry>(); // key: chatId
// Env-resolved caps (per instance) so a deployment can tune the ceiling without
// a code change. The subscriber cap keeps the documented 2× relationship.
readonly maxBufferBytes = resolveMaxBufferBytes();
readonly subscriberMaxBufferedBytes = 2 * this.maxBufferBytes;
/**
* Register a fresh entry at the START of a run (before any frame), so a tab
@@ -105,7 +209,11 @@ export class AiChatStreamRegistryService implements OnModuleDestroy {
this.entries.set(chatId, {
runId,
frames: [],
stamps: [],
bytes: 0,
currentStamp: 0,
persistedFloor: 0,
overflowThroughStamp: -1,
overflowed: false,
finished: false,
subscribers: new Set<Subscriber>(),
@@ -150,6 +258,34 @@ export class AiChatStreamRegistryService implements OnModuleDestroy {
void pump();
}
/**
* Confirm that step `stepsPersisted` (a COUNT: steps 0..stepsPersisted-1) is on
* disk for this run, and ROTATE the ring: drop the buffered frames of those
* now-persisted steps (stamp < stepsPersisted). This is the ONLY thing that
* rotates the ring, and it is called ONLY after a genuinely SUCCESSFUL per-step
* persist (see ai-chat.service updateStreaming). A failed persist never calls
* it, so the ring covers more (auto-safe). Identity-checked (invariant 1) and
* monotonic (a stale lower count is ignored).
*/
confirmPersistedStep(
chatId: string,
runId: string,
stepsPersisted: number,
): void {
const entry = this.entries.get(chatId);
if (!entry || entry.runId !== runId) return;
if (!Number.isFinite(stepsPersisted) || stepsPersisted <= entry.persistedFloor)
return;
entry.persistedFloor = stepsPersisted;
// Clean rotation: drop the persisted steps from the head. These frames are on
// disk + carried by a fresh client seed, so this NEVER opens a gap.
while (entry.frames.length > 0 && entry.stamps[0] < stepsPersisted) {
entry.bytes -= Buffer.byteLength(entry.frames[0]);
entry.frames.shift();
entry.stamps.shift();
}
}
/**
* Terminate a run's entry from the OUTER catch of the stream method (a failure
* before/while wiring the pipe, so `done` will never arrive). Identity-checked
@@ -162,36 +298,77 @@ export class AiChatStreamRegistryService implements OnModuleDestroy {
}
/**
* Attach to a run's stream. Async only for the phase-2 Redis seam the body
* runs synchronously so the replay snapshot and the subscriber registration
* happen in ONE tick with no await between them (invariant 4): a frame ingested
* concurrently cannot slip into the gap and be lost or duplicated.
* Attach to a run's stream from the client's step frontier `n` (its persisted
* `stepsPersisted`). Async only for the phase-2 Redis seam the body runs
* synchronously so the tail SLICE and the subscriber registration happen in ONE
* tick with no await between them (invariant 4).
*
* Returns null (-> the caller answers 204) when:
* - there is no entry, or it overflowed (replay is gone);
* - expect=live with an anchor that does not match this run's assistant id
* (invariant 6: a stripped tab must never replay a FOREIGN run's transcript);
* - the run finished and the caller did not expect a live tail.
* A finished run with expect=live yields a replay-only attachment (no
* subscriber registered). Otherwise a paused subscriber is registered and the
* caller replays `replay`, then calls start() to drain and go live.
* - there is no entry;
* - the `anchor` does not match this run's assistant id (invariant 6);
* - the ring does not cover the client's frontier (coverageFloor > n): a hole
* from overflow, or the client's seed simply lagged behind a rotation. The
* client then refetches (a larger n) and re-attaches.
*
* Otherwise the attachment's `replay` is a synthetic `start` frame (the run-fact
* on re-attach) followed by the buffered tail filtered to `stamp >= n`. For a
* FINISHED run this is replay-only (no subscriber) and ends after the replay
* with n = N_final that tail is just the run's `finish` frame, so the client
* closes the stream. For a LIVE run a paused subscriber is registered; the
* caller writes the replay (respecting drain) then calls start() to drain the
* pending frames and go live.
*/
async attach(
chatId: string,
expectLive: boolean,
anchor: string | undefined,
// The client's persisted step frontier. `null` = a NOT-tail-aware client (no
// `n` query param) — a legacy/parameterless tab that expects the old
// "finished -> 204 -> poll" contract; distinct from `0` (a tail-aware client
// with nothing persisted yet).
n: number | null,
cb: RunStreamCallbacks,
): Promise<RunStreamAttachment | null> {
const entry = this.entries.get(chatId);
if (!entry || entry.overflowed) return null;
if (!entry) return null;
// Invariant 6: cross-run replay is forbidden. Before bind, assistantMessageId
// is undefined and mismatches any anchor -> 204 -> client restore+poll path.
if (expectLive && anchor && entry.assistantMessageId !== anchor) return null;
if (entry.finished && !expectLive) return null;
if (entry.finished && expectLive) {
if (anchor && entry.assistantMessageId !== anchor) return null;
// #491 regression guard (#137/#161 dup): a NOT-tail-aware client (no `n`)
// resuming a FINISHED run must 204 and poll — the old `finished && !expectLive`
// gate. Without this, a missing `n` collapsing to frontier 0 would serve the
// WHOLE tail of a finished, NON-rotated run (coverageFloor 0), and a
// parameterless client that never stripped its transcript would APPEND that
// full replay onto the steps it already shows -> duplicated text. A tail-aware
// client (n present, incl. n=0) still gets the tail past its frontier.
if (entry.finished && n === null) return null;
// A finished entry with NOTHING in the ring (aborted before the first frame,
// or fully overflowed) has no tail to deliver -> 204 -> the client polls.
if (entry.finished && entry.frames.length === 0) return null;
// A LIVE run with no `n` (legacy parameterless) replays from step 0 (the old
// behavior); a tail-aware client resumes from its frontier.
const frontier = n ?? 0;
const floor = this.coverageFloor(entry);
if (floor > frontier) {
this.logger.warn(
`run-stream attach gap for run=${entry.runId}: coverageFloor=${floor} ` +
`> client frontier=${frontier} -> 204 (client refetches + re-attaches)`,
);
return null;
}
const startFrame = this.buildStartFrame(chatId, entry.runId);
const sliceTail = (): string[] => {
const out: string[] = [startFrame];
for (let i = 0; i < entry.frames.length; i++) {
if (entry.stamps[i] >= frontier) out.push(entry.frames[i]);
}
return out;
};
if (entry.finished) {
// Replay-only: the run is done, no subscriber is registered.
return {
replay: entry.frames.slice(),
replay: sliceTail(),
finished: true,
start: () => undefined,
unsubscribe: () => undefined,
@@ -206,15 +383,12 @@ export class AiChatStreamRegistryService implements OnModuleDestroy {
pendingBytes: 0,
overflowed: false,
pendingEnd: false,
minStamp: frontier,
};
// Register + snapshot in the SAME synchronous block (invariant 4). No await
// separates them, so a concurrently ingested frame cannot be lost/duplicated.
entry.subscribers.add(sub);
// Snapshot in the SAME synchronous block as the registration (invariant 4).
const replay = entry.frames.slice();
// CONTRACT: the caller MUST call start() in the SAME tick as this attach()
// returns — no await between them. While a subscriber is paused, every frame
// is buffered in sub.pending; a delayed start() lets a whole run accumulate
// there. The pendingBytes cap (see ingestFrame) is the structural backstop if
// that contract is ever broken (e.g. the phase-2 Redis await seam).
const replay = sliceTail();
return {
replay,
finished: false,
@@ -263,24 +437,83 @@ export class AiChatStreamRegistryService implements OnModuleDestroy {
this.entries.clear();
}
/** Buffer + fan-out a single frame. See invariant/overflow semantics inline. */
/** The synthetic `start` frame the tail is prefixed with the source of the
* run-fact (runId/chatId) on re-attach. A `start` frame does NOT reset the
* client's message parts (ai@6.0.207 createStreamingUIMessageState), so it is
* safe to prepend even when the sliced tail begins mid-message. */
private buildStartFrame(chatId: string, runId: string): string {
return `data: ${JSON.stringify({
type: 'start',
messageMetadata: { runId, chatId },
})}\n\n`;
}
/**
* The smallest step FULLY present in the ring: its smallest retained stamp, or
* (when the leading step was only partially evicted by an overflow) one past it.
* When the ring is empty it is the current step (only the live tail is coming).
* An attach at frontier `n` is covered coverageFloor <= n.
*/
private coverageFloor(entry: Entry): number {
// Empty ring: only the live tail is coming. The floor is the current step,
// but never below persistedFloor — a confirmed persist can rotate the ring
// empty while currentStamp still lags a beat behind on another connection, so
// max() keeps the invariant STRUCTURAL (a client with n = persistedFloor is
// always covered) rather than timing-dependent.
if (entry.frames.length === 0)
return Math.max(entry.currentStamp, entry.persistedFloor);
const min = entry.stamps[0];
return entry.overflowThroughStamp >= min ? min + 1 : min;
}
/**
* Buffer (step-stamped) + fan-out a single frame. The stamp is the number of
* finish-step frames seen BEFORE this one; a finish-step frame carries the
* current value and THEN increments the counter (so its stamp equals the 0-based
* index of the step it closes). Only frames at/above persistedFloor are buffered
* (already-persisted steps are on disk); the ring is then trimmed to the byte
* cap, an unsafe eviction opening a gap. Fan-out is always live (filtered per
* subscriber by its frontier).
*/
private ingestFrame(entry: Entry, frame: string): void {
entry.bytes += Buffer.byteLength(frame);
if (!entry.overflowed) {
const size = Buffer.byteLength(frame);
const stamp = entry.currentStamp;
if (frame.startsWith(FINISH_STEP_FRAME_PREFIX)) {
entry.currentStamp = stamp + 1;
}
// Buffer for replay only if this step is not already persisted+rotated away.
if (stamp >= entry.persistedFloor) {
entry.frames.push(frame);
if (entry.bytes > RUN_STREAM_MAX_BUFFER_BYTES) {
// The crossing frame was already counted AND (below) fanned out; only the
// replay buffer is dropped. After overflow no more frames are buffered,
// but live fan-out continues.
entry.overflowed = true;
entry.frames = [];
this.logger.warn(
`run-stream buffer overflow for run=${entry.runId}; ` +
`late attach will 204 until the run ends`,
);
entry.stamps.push(stamp);
entry.bytes += size;
// Enforce the ring cap. Evicting a not-yet-persisted frame (stamp >=
// persistedFloor) opens a GAP; a leftover persisted frame (< floor) is a
// safe drop. Keep evicting until the ring is back under the cap.
while (entry.bytes > this.maxBufferBytes && entry.frames.length > 0) {
const evStamp = entry.stamps[0];
entry.bytes -= Buffer.byteLength(entry.frames[0]);
entry.frames.shift();
entry.stamps.shift();
if (evStamp >= entry.persistedFloor) {
if (evStamp > entry.overflowThroughStamp)
entry.overflowThroughStamp = evStamp;
if (!entry.overflowed) {
entry.overflowed = true;
this.logger.warn(
`run-stream ring overflow for run=${entry.runId}: an un-persisted ` +
`step was evicted to stay under ${this.maxBufferBytes}B; a late ` +
`attach at an evicted step will 204 until a later persist confirms`,
);
}
}
}
}
// Fan out live, filtered to each subscriber's frontier (a subscriber only
// wants the tail past the step it already persisted).
for (const sub of entry.subscribers) {
if (stamp < sub.minStamp) continue;
if (sub.started) {
try {
sub.onFrame(frame);
@@ -289,12 +522,12 @@ export class AiChatStreamRegistryService implements OnModuleDestroy {
}
} else {
sub.pending.push(frame);
sub.pendingBytes += Buffer.byteLength(frame);
if (sub.pendingBytes > SUBSCRIBER_MAX_BUFFERED_BYTES) {
sub.pendingBytes += size;
if (sub.pendingBytes > this.subscriberMaxBufferedBytes) {
// The paused subscriber's buffer overflowed — only possible if start()
// was delayed past the same-tick contract (the phase-2 await seam).
// Drop it rather than buffer the whole run; on start() it degrades to an
// immediate end (a 204-equivalent) instead of replaying a partial.
// was delayed (the controller's drain-respecting tail write, or the
// phase-2 await seam). Drop it rather than buffer the whole run; on
// start() it degrades to an immediate end (a 204-equivalent).
sub.overflowed = true;
sub.pending = [];
entry.subscribers.delete(sub);
@@ -1,19 +1,27 @@
import {
AiChatStreamRegistryService,
RUN_STREAM_MAX_BUFFER_BYTES,
AI_CHAT_RUN_STREAM_MAX_BUFFER_BYTES,
RUN_STREAM_RETAIN_FINISHED_MS,
SUBSCRIBER_MAX_BUFFERED_BYTES,
RunStreamCallbacks,
} from './ai-chat-stream-registry.service';
/**
* Unit tests for the in-memory run-stream registry (#184 phase 1.5). The registry
* is the whole of the resumable-transport contract: replay ordering, paused ->
* live hand-off, overflow, retention, the anchor check (invariant 6), and the
* mirror-the-done-path replace semantics (invariant 3). Every enumerated case in
* the issue's task 1.5 has a test here.
* Unit tests for the in-memory run-stream registry (#184 phase 1.5, step-aligned
* retention #491). The registry is the whole of the resumable-transport contract:
* step-stamped retention, tail-only attach at the client's frontier N, the
* confirmed-persist ring rotation (and the anti-inversion rule), the memory bound,
* the overflow gap, paused -> live hand-off, retention, the anchor check
* (invariant 6), and the mirror-the-done-path replace semantics (invariant 3).
*/
// Real ai@6 UI-message-stream SSE frames are `data: {json}\n\n`, one part each.
const sse = (part: Record<string, unknown>): string =>
`data: ${JSON.stringify(part)}\n\n`;
const finishStep = (): string => sse({ type: 'finish-step' });
const textDelta = (id: string, delta: string): string =>
sse({ type: 'text-delta', id, delta });
const finish = (): string => sse({ type: 'finish' });
// A ReadableStream whose frames the test pushes explicitly, plus close/error.
function makePushStream(): {
stream: ReadableStream<string>;
@@ -58,6 +66,9 @@ function collector(): {
};
}
// The tail past the synthetic start frame (replay[0] is always the start frame).
const tail = (replay: string[]): string[] => replay.slice(1);
describe('AiChatStreamRegistryService', () => {
const CHAT = 'chat-1';
let registry: AiChatStreamRegistryService;
@@ -71,7 +82,21 @@ describe('AiChatStreamRegistryService', () => {
registry.onModuleDestroy();
});
it('replays frames in arrival order (live attach)', async () => {
it('prepends a synthetic start frame carrying { runId, chatId }', async () => {
registry.open(CHAT, 'run-1');
const src = makePushStream();
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
src.push('a');
await flush();
const c = collector();
const att = (await registry.attach(CHAT, 'assist-1', 0, c.cb))!;
const start = JSON.parse(att.replay[0].replace(/^data: /, '').trim());
expect(start.type).toBe('start');
expect(start.messageMetadata).toEqual({ runId: 'run-1', chatId: CHAT });
});
it('replays the buffered tail (from frontier 0) in arrival order (live attach)', async () => {
registry.open(CHAT, 'run-1');
const src = makePushStream();
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
@@ -81,13 +106,13 @@ describe('AiChatStreamRegistryService', () => {
await flush();
const c = collector();
const att = await registry.attach(CHAT, false, undefined, c.cb);
const att = await registry.attach(CHAT, 'assist-1', 0, c.cb);
expect(att).not.toBeNull();
expect(att!.replay).toEqual(['a', 'b', 'c']);
expect(tail(att!.replay)).toEqual(['a', 'b', 'c']);
expect(att!.finished).toBe(false);
});
it('late attach gets the full prefix as replay plus the live tail', async () => {
it('late attach gets the buffered prefix as tail plus the live tail', async () => {
registry.open(CHAT, 'run-1');
const src = makePushStream();
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
@@ -96,17 +121,16 @@ describe('AiChatStreamRegistryService', () => {
await flush();
const c = collector();
const att = (await registry.attach(CHAT, false, undefined, c.cb))!;
expect(att.replay).toEqual(['a', 'b']);
const att = (await registry.attach(CHAT, 'assist-1', 0, c.cb))!;
expect(tail(att.replay)).toEqual(['a', 'b']);
att.start();
// Live tail arrives after start().
src.push('c');
src.push('d');
await flush();
expect(c.frames).toEqual(['c', 'd']);
});
it('a paused subscriber receives frames buffered during pause in order, then live (no loss/reorder)', async () => {
it('a paused subscriber receives frames buffered during pause in order, then live', async () => {
registry.open(CHAT, 'run-1');
const src = makePushStream();
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
@@ -114,81 +138,45 @@ describe('AiChatStreamRegistryService', () => {
await flush();
const c = collector();
// Attach (paused). Frames that arrive BEFORE start() must queue, not drop.
const att = (await registry.attach(CHAT, false, undefined, c.cb))!;
expect(att.replay).toEqual(['a']);
const att = (await registry.attach(CHAT, 'assist-1', 0, c.cb))!;
expect(tail(att.replay)).toEqual(['a']);
src.push('b'); // arrives while paused -> pending
src.push('c');
await flush();
expect(c.frames).toEqual([]); // nothing delivered yet (paused)
att.start(); // drains pending in order
att.start();
expect(c.frames).toEqual(['b', 'c']);
src.push('d'); // now live
src.push('d');
await flush();
expect(c.frames).toEqual(['b', 'c', 'd']);
});
it('a run that finishes while a subscriber is paused ends it on start()', async () => {
registry.open(CHAT, 'run-1');
registry.bind(CHAT, 'run-1', 'assist-1', makePushStream().stream);
const c = collector();
const att = (await registry.attach(CHAT, false, undefined, c.cb))!;
// Terminate the run while the subscriber is still paused.
const att = (await registry.attach(CHAT, 'assist-1', 0, c.cb))!;
registry.abortEntry(CHAT, 'run-1');
expect(c.ended()).toBe(0); // paused: not ended yet
att.start();
expect(c.ended()).toBe(1); // start() drains + ends
});
it('finished + expect=live returns a replay WITHOUT registering a subscriber', async () => {
registry.open(CHAT, 'run-1');
const src = makePushStream();
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
src.push('a');
src.push('b');
src.close();
await flush();
const c = collector();
const att = (await registry.attach(CHAT, true, undefined, c.cb))!;
expect(att.finished).toBe(true);
expect(att.replay).toEqual(['a', 'b']);
// No subscriber registered: start()/unsubscribe are no-ops and the entry has
// zero subscribers.
const entry = (registry as any).entries.get(CHAT);
expect(entry.subscribers.size).toBe(0);
att.start();
expect(c.frames).toEqual([]);
});
it('finished WITHOUT expect=live returns null', async () => {
registry.open(CHAT, 'run-1');
const src = makePushStream();
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
src.push('a');
src.close();
await flush();
const c = collector();
expect(await registry.attach(CHAT, false, undefined, c.cb)).toBeNull();
});
it('anchor mismatch with expect=live returns null (and null before bind sets assistantMessageId)', async () => {
it('anchor mismatch returns null (and null before bind sets assistantMessageId)', async () => {
registry.open(CHAT, 'run-1');
const c = collector();
// Before bind: assistantMessageId is undefined -> mismatches any anchor.
expect(
await registry.attach(CHAT, true, 'assist-1', c.cb),
).toBeNull();
expect(await registry.attach(CHAT, 'assist-1', 0, c.cb)).toBeNull();
const src = makePushStream();
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
src.push('a');
await flush();
// Wrong anchor -> null (cross-run replay forbidden, invariant 6).
expect(await registry.attach(CHAT, true, 'other-id', c.cb)).toBeNull();
expect(await registry.attach(CHAT, 'other-id', 0, c.cb)).toBeNull();
});
it('matching anchor with expect=live attaches', async () => {
it('matching anchor attaches', async () => {
registry.open(CHAT, 'run-1');
const src = makePushStream();
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
@@ -196,97 +184,60 @@ describe('AiChatStreamRegistryService', () => {
await flush();
const c = collector();
const att = await registry.attach(CHAT, true, 'assist-1', c.cb);
const att = await registry.attach(CHAT, 'assist-1', 0, c.cb);
expect(att).not.toBeNull();
expect(att!.replay).toEqual(['a']);
expect(tail(att!.replay)).toEqual(['a']);
});
it('overflow: attach returns null, but the LIVE subscriber keeps receiving (incl. the crossing frame)', async () => {
it('a throwing onFrame ejects only that subscriber; the ingest loop stays alive', async () => {
registry.open(CHAT, 'run-1');
const src = makePushStream();
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
// A live (started) subscriber attached before the flood.
const bad = collector();
const badAtt = (await registry.attach(CHAT, 'assist-1', 0, {
onFrame: () => {
throw new Error('boom');
},
onEnd: bad.cb.onEnd,
}))!;
badAtt.start();
const good = collector();
const goodAtt = (await registry.attach(CHAT, 'assist-1', 0, good.cb))!;
goodAtt.start();
src.push('a'); // bad throws on this frame -> ejected
src.push('b'); // good still receives both
await flush();
const entry = (registry as any).entries.get(CHAT);
expect(entry.subscribers.size).toBe(1);
expect(good.frames).toEqual(['a', 'b']);
});
it('open() over a LIVE entry ends started subscribers once; a late done never touches the new entry (invariant 3)', async () => {
registry.open(CHAT, 'run-1');
const src = makePushStream();
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
src.push('a');
await flush();
const c = collector();
const att = (await registry.attach(CHAT, false, undefined, c.cb))!;
const att = (await registry.attach(CHAT, 'assist-1', 0, c.cb))!;
att.start();
// Cap-relative so it survives a buffer-cap change (#430): a quarter-cap frame
// means 5 frames comfortably exceed the replay cap; the last one crosses.
const chunk = 'x'.repeat(Math.floor(RUN_STREAM_MAX_BUFFER_BYTES / 4));
for (let i = 0; i < 5; i++) src.push(chunk + i);
await flush();
const entry = (registry as any).entries.get(CHAT);
expect(entry.overflowed).toBe(true);
expect(entry.bytes).toBeGreaterThan(RUN_STREAM_MAX_BUFFER_BYTES);
// The live subscriber received ALL 5 frames, including the crossing one.
expect(c.frames).toHaveLength(5);
expect(c.frames[4]).toBe(chunk + 4);
// A NEW attach after overflow gets null (replay buffer is gone).
const c2 = collector();
expect(await registry.attach(CHAT, false, undefined, c2.cb)).toBeNull();
});
it('a paused subscriber whose pending buffer overflows is dropped and ends on start(); other subscribers keep receiving', async () => {
registry.open(CHAT, 'run-1');
const src = makePushStream();
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
// A: paused (start() deliberately delayed to simulate the phase-2 await seam).
const a = collector();
const attA = (await registry.attach(CHAT, false, undefined, a.cb))!;
// B: live (started) — its delivery must be unaffected by A's overflow.
const b = collector();
const attB = (await registry.attach(CHAT, false, undefined, b.cb))!;
attB.start();
// Cap-relative so it survives a buffer-cap change (#430): a quarter-of-the-
// per-subscriber-cap frame means 5 frames exceed A's paused-pending cap while
// B streams every frame live.
const chunk = 'x'.repeat(Math.floor(SUBSCRIBER_MAX_BUFFERED_BYTES / 4));
for (let i = 0; i < 5; i++) src.push(chunk + i);
await flush();
const entry = (registry as any).entries.get(CHAT);
// A was dropped from the subscriber set on overflow; B (started) remains.
expect(entry.subscribers.size).toBe(1);
expect(a.frames).toEqual([]); // paused + overflowed: nothing was delivered
// B received every frame live (delivery unaffected by A's overflow).
expect(b.frames).toHaveLength(5);
// A's start() (arriving late) degrades to an immediate end, not a partial replay.
attA.start();
expect(a.frames).toEqual([]);
expect(a.ended()).toBe(1);
});
it('open() over a LIVE entry ends started subscribers exactly once and a late done does not touch the new entry (invariant 3)', async () => {
registry.open(CHAT, 'run-1');
const src = makePushStream();
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
src.push('a');
await flush();
const c = collector();
const att = (await registry.attach(CHAT, false, undefined, c.cb))!;
att.start(); // started subscriber on run-1
// run-2 starts on the same chat while run-1's tee is still reading.
registry.open(CHAT, 'run-2');
expect(c.ended()).toBe(1); // exactly one onEnd from the replace
expect(c.ended()).toBe(1);
const newEntry = (registry as any).entries.get(CHAT);
expect(newEntry.runId).toBe('run-2');
expect(newEntry.finished).toBe(false);
// The old tee now completes: its late done must NOT double-end nor delete the
// new entry.
src.push('b');
src.close();
await flush();
expect(c.ended()).toBe(1); // still exactly one
expect(c.ended()).toBe(1);
const still = (registry as any).entries.get(CHAT);
expect(still).toBe(newEntry);
expect(still.runId).toBe('run-2');
@@ -299,7 +250,6 @@ describe('AiChatStreamRegistryService', () => {
src.push('a');
await flush();
const entry = (registry as any).entries.get(CHAT);
// Frames were NOT ingested (bind bailed), assistantMessageId untouched.
expect(entry.frames).toEqual([]);
expect(entry.assistantMessageId).toBeUndefined();
});
@@ -310,32 +260,276 @@ describe('AiChatStreamRegistryService', () => {
const entry = (registry as any).entries.get(CHAT);
expect(entry.finished).toBe(false);
});
});
it('a throwing onFrame ejects only that subscriber; the ingest loop stays alive', async () => {
/**
* #491 step-stamped retention: the boundary detector, tail-only slicing at the
* client's frontier N, the confirmed-persist rotation (+ anti-inversion), the
* overflow gap, the memory bound, and the finished-retained tail. All observable
* against the REAL registry driven through open/bind/ingest.
*/
describe('AiChatStreamRegistryService step-aligned retention (#491)', () => {
const CHAT = 'chat-s';
let registry: AiChatStreamRegistryService;
beforeEach(() => {
registry = new AiChatStreamRegistryService();
jest.spyOn((registry as any).logger, 'warn').mockImplementation(() => {});
});
afterEach(() => registry.onModuleDestroy());
const entryOf = () => (registry as any).entries.get(CHAT);
it('stamps frames by finish-step count, aligned with stepsPersisted', async () => {
registry.open(CHAT, 'run-1');
const src = makePushStream();
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
// step 0 content, its finish-step, step 1 content, its finish-step, finish.
src.push(textDelta('t0', 'a')); // stamp 0
src.push(finishStep()); // stamp 0 (the finish-step frame carries the pre value)
src.push(textDelta('t1', 'b')); // stamp 1
src.push(finishStep()); // stamp 1
src.push(finish()); // stamp 2
await flush();
const e = entryOf();
expect(e.stamps).toEqual([0, 0, 1, 1, 2]);
expect(e.currentStamp).toBe(2);
});
const bad = collector();
const badAtt = (await registry.attach(CHAT, false, undefined, {
onFrame: () => {
throw new Error('boom');
},
onEnd: bad.cb.onEnd,
}))!;
badAtt.start();
it('does NOT treat a text delta that merely quotes "finish-step" as a boundary', async () => {
registry.open(CHAT, 'run-1');
const src = makePushStream();
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
// A model that literally types "type":"finish-step" — JSON-escaped in the frame.
src.push(textDelta('t0', '"type":"finish-step"'));
await flush();
expect(entryOf().currentStamp).toBe(0); // no false boundary
});
const good = collector();
const goodAtt = (await registry.attach(CHAT, false, undefined, good.cb))!;
goodAtt.start();
src.push('a'); // bad throws on this frame -> ejected
src.push('b'); // good still receives both
it('tail-only: attach at N slices frames with stamp >= N', async () => {
registry.open(CHAT, 'run-1');
const src = makePushStream();
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
src.push(textDelta('t0', 'a')); // 0
src.push(finishStep()); // 0
src.push(textDelta('t1', 'b')); // 1
src.push(finishStep()); // 1
src.push(textDelta('t2', 'c')); // 2 (in-progress)
await flush();
const entry = (registry as any).entries.get(CHAT);
expect(entry.subscribers.size).toBe(1); // bad ejected, good remains
expect(good.frames).toEqual(['a', 'b']);
const c = collector();
// Client persisted 2 steps -> wants the tail from step 2.
const att = (await registry.attach(CHAT, 'assist-1', 2, c.cb))!;
expect(tail(att.replay)).toEqual([textDelta('t2', 'c')]);
});
it('attach in the MIDDLE of a step (N between finish-steps) slices from that step', async () => {
registry.open(CHAT, 'run-1');
const src = makePushStream();
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
src.push(textDelta('t0', 'a')); // 0
src.push(finishStep()); // 0
src.push(textDelta('t1', 'b1')); // 1
src.push(textDelta('t1', 'b2')); // 1 (still step 1, no finish-step yet)
await flush();
const c = collector();
const att = (await registry.attach(CHAT, 'assist-1', 1, c.cb))!;
// Step 0's frames are dropped from the tail; the whole in-progress step 1 is kept.
expect(tail(att.replay)).toEqual([textDelta('t1', 'b1'), textDelta('t1', 'b2')]);
});
it('rotates the ring ONLY on a confirmed persist (drops stamp < N)', async () => {
registry.open(CHAT, 'run-1');
const src = makePushStream();
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
src.push(textDelta('t0', 'a')); // 0
src.push(finishStep()); // 0
src.push(textDelta('t1', 'b')); // 1
await flush();
expect(entryOf().stamps).toEqual([0, 0, 1]);
// Confirm step 0 persisted (stepsPersisted = 1) -> drop stamp < 1.
registry.confirmPersistedStep(CHAT, 'run-1', 1);
expect(entryOf().stamps).toEqual([1]);
expect(entryOf().persistedFloor).toBe(1);
});
it('persist FAILED but the ring still fits -> attach SUCCEEDS and the tail includes step N', async () => {
registry.open(CHAT, 'run-1');
const src = makePushStream();
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
src.push(textDelta('t0', 'a')); // 0
src.push(finishStep()); // 0
src.push(textDelta('t1', 'b')); // 1 (step 1's persist FAILED -> no confirm)
await flush();
// No confirmPersistedStep for step 1: the ring still holds step 1.
const c = collector();
// Client's last successful persist was step 0 -> stepsPersisted = 1.
const att = await registry.attach(CHAT, 'assist-1', 1, c.cb);
expect(att).not.toBeNull();
expect(tail(att!.replay)).toEqual([textDelta('t1', 'b')]); // includes step 1
});
it('persist failed AND the ring overflowed past N -> 204 (coverage gap)', async () => {
registry.open(CHAT, 'run-1');
const src = makePushStream();
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
// Step 0: a fat step that blows past the cap with NO persist confirmation.
const big = 'x'.repeat(Math.floor(AI_CHAT_RUN_STREAM_MAX_BUFFER_BYTES / 2));
src.push(textDelta('t0', big)); // 0
src.push(textDelta('t0', big)); // 0
src.push(textDelta('t0', big)); // 0 -> overflow evicts stamp-0 frames
await flush();
const e = entryOf();
expect(e.overflowed).toBe(true);
expect(e.bytes).toBeLessThanOrEqual(registry.maxBufferBytes);
// A client at frontier 0 falls at/below an evicted step -> gap -> null.
const c = collector();
expect(await registry.attach(CHAT, 'assist-1', 0, c.cb)).toBeNull();
});
it('stale N (client seed lagged behind a rotation) -> 204; after a refetch (larger N) -> success', async () => {
registry.open(CHAT, 'run-1');
const src = makePushStream();
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
src.push(textDelta('t0', 'a')); // 0
src.push(finishStep()); // 0
src.push(textDelta('t1', 'b')); // 1
src.push(finishStep()); // 1
src.push(textDelta('t2', 'c')); // 2
await flush();
// Server confirmed steps 0 and 1 -> rotate away stamp < 2.
registry.confirmPersistedStep(CHAT, 'run-1', 2);
expect(entryOf().stamps).toEqual([2]);
// A client whose seed still says stepsPersisted = 1 -> below minStamp -> 204.
const stale = collector();
expect(await registry.attach(CHAT, 'assist-1', 1, stale.cb)).toBeNull();
// It refetches (now stepsPersisted = 2) and re-attaches -> success.
const fresh = collector();
const att = await registry.attach(CHAT, 'assist-1', 2, fresh.cb);
expect(att).not.toBeNull();
expect(tail(att!.replay)).toEqual([textDelta('t2', 'c')]);
});
it('overflow gap CLEARS once a later persist rotates out the holey steps', async () => {
registry.open(CHAT, 'run-1');
const src = makePushStream();
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
const big = 'x'.repeat(Math.floor(AI_CHAT_RUN_STREAM_MAX_BUFFER_BYTES / 2));
src.push(textDelta('t0', big)); // 0
src.push(textDelta('t0', big)); // 0
src.push(finishStep()); // 0 (still stamp 0)
src.push(textDelta('t1', 'small')); // 1
src.push(finishStep()); // 1
src.push(textDelta('t2', 'c')); // 2
await flush();
expect(entryOf().overflowed).toBe(true);
// Late persist confirms steps 0..1 -> rotates out the holey step-0 frames.
registry.confirmPersistedStep(CHAT, 'run-1', 2);
// A client at frontier 2 is now cleanly covered (the hole was below it).
const c = collector();
const att = await registry.attach(CHAT, 'assist-1', 2, c.cb);
expect(att).not.toBeNull();
expect(tail(att!.replay)).toEqual([textDelta('t2', 'c')]);
});
it('finished-retained + N = N_final -> empty tail plus the finish frame', async () => {
registry.open(CHAT, 'run-1');
const src = makePushStream();
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
src.push(textDelta('t0', 'a')); // 0
src.push(finishStep()); // 0
src.push(finish()); // 1 (N_final = 1)
src.close();
await flush();
// The last step's per-step persist confirmed stepsPersisted = 1.
registry.confirmPersistedStep(CHAT, 'run-1', 1);
const c = collector();
const att = (await registry.attach(CHAT, 'assist-1', 1, c.cb))!;
expect(att.finished).toBe(true);
// Empty step tail; just the finish frame so the client's SDK closes the stream.
expect(tail(att.replay)).toEqual([finish()]);
// No subscriber registered for a finished run.
expect(entryOf().subscribers.size).toBe(0);
});
it('#491 regression (#137/#161 dup): a PARAMETERLESS attach (n=null) to a finished NON-rotated run -> 204, but n=0 still gets the tail', async () => {
// A finished, non-rotated run: frames present, coverageFloor 0. A missing `n`
// (null — a legacy/parameterless tab that never stripped its transcript) must
// 204 -> poll, NOT receive the whole tail it would append (duplicate). A
// tail-aware client (n=0 present) still resumes.
registry.open(CHAT, 'run-1');
const src = makePushStream();
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
src.push(textDelta('t0', 'a')); // 0
src.push(finishStep()); // 0
src.push(finish()); // 1
src.close();
await flush();
// NOT rotated (no confirmPersistedStep) -> stamps[0]=0, coverageFloor=0.
// MUTATION-VERIFY: revert the `finished && n === null -> null` gate (default n
// to 0) and the parameterless attach below serves the full tail instead of 204.
expect(await registry.attach(CHAT, 'assist-1', null, collector().cb)).toBeNull();
// A tail-aware client at frontier 0 IS served (the distinction: null != 0).
const tailAware = await registry.attach(CHAT, 'assist-1', 0, collector().cb);
expect(tailAware).not.toBeNull();
expect(tailAware!.finished).toBe(true);
});
it('confirmPersistedStep is monotonic and identity-checked', async () => {
registry.open(CHAT, 'run-1');
const src = makePushStream();
registry.bind(CHAT, 'run-1', 'assist-1', src.stream);
src.push(textDelta('t0', 'a'));
src.push(finishStep());
src.push(textDelta('t1', 'b'));
await flush();
registry.confirmPersistedStep(CHAT, 'run-1', 1);
expect(entryOf().persistedFloor).toBe(1);
// A stale lower count is ignored.
registry.confirmPersistedStep(CHAT, 'run-1', 0);
expect(entryOf().persistedFloor).toBe(1);
// A foreign runId is ignored.
registry.confirmPersistedStep(CHAT, 'WRONG', 5);
expect(entryOf().persistedFloor).toBe(1);
});
it('MEMORY BOUND: 5 parallel marathon runs each stream well past 32MB; each ring stays <= the cap', async () => {
const cap = registry.maxBufferBytes;
const chats = ['m0', 'm1', 'm2', 'm3', 'm4'];
const srcs = chats.map((chat) => {
registry.open(chat, `run-${chat}`);
const s = makePushStream();
registry.bind(chat, `run-${chat}`, `assist-${chat}`, s.stream);
return s;
});
// ~256KB frames; 160 per chat = 40MB streamed each, well past the old 32MB.
// Interleave a finish-step every 8 frames so steps advance realistically. No
// persist confirmation -> the ONLY thing keeping memory bounded is the cap.
const frame = 'y'.repeat(256 * 1024);
for (let batch = 0; batch < 20; batch++) {
for (let i = 0; i < 8; i++) {
for (const s of srcs) s.push(textDelta('t', frame));
}
for (const s of srcs) s.push(finishStep());
await flush(); // drain the pump so queues never hold a whole run
}
let total = 0;
for (const chat of chats) {
const e = (registry as any).entries.get(chat);
expect(e.bytes).toBeLessThanOrEqual(cap);
total += e.bytes;
}
// Total retained across all 5 runs is bounded by 5x the per-run cap — the old
// registry would have retained ~5x40MB = 200MB here.
expect(total).toBeLessThanOrEqual(cap * chats.length);
});
});
@@ -361,7 +555,7 @@ describe('AiChatStreamRegistryService retention timers', () => {
it('a finished entry is removed after the retention window', () => {
registry.open(CHAT, 'run-1');
registry.abortEntry(CHAT, 'run-1'); // finalize -> retention armed
registry.abortEntry(CHAT, 'run-1');
expect((registry as any).entries.get(CHAT)).toBeDefined();
jest.advanceTimersByTime(RUN_STREAM_RETAIN_FINISHED_MS + 1);
expect((registry as any).entries.get(CHAT)).toBeUndefined();
@@ -369,20 +563,18 @@ describe('AiChatStreamRegistryService retention timers', () => {
it('retention deletes ONLY its own entry (invariant 2)', () => {
registry.open(CHAT, 'run-1');
registry.abortEntry(CHAT, 'run-1'); // arm retention for entry A
// Simulate the race where the key was replaced without clearing A's timer.
registry.abortEntry(CHAT, 'run-1');
const sentinel = { marker: true };
(registry as any).entries.set(CHAT, sentinel);
jest.advanceTimersByTime(RUN_STREAM_RETAIN_FINISHED_MS + 1);
// A's timer saw entries.get(CHAT) !== A, so it did NOT delete the successor.
expect((registry as any).entries.get(CHAT)).toBe(sentinel);
});
it('open() over a retained entry clears its timer and the successor survives', () => {
registry.open(CHAT, 'run-1');
registry.abortEntry(CHAT, 'run-1'); // retained, timer armed
registry.abortEntry(CHAT, 'run-1');
const clearSpy = jest.spyOn(global, 'clearTimeout');
registry.open(CHAT, 'run-2'); // must clear run-1's retain timer
registry.open(CHAT, 'run-2');
expect(clearSpy).toHaveBeenCalled();
jest.advanceTimersByTime(RUN_STREAM_RETAIN_FINISHED_MS + 1);
const entry = (registry as any).entries.get(CHAT);
@@ -8,10 +8,12 @@ import { SUBSCRIBER_MAX_BUFFERED_BYTES } from './ai-chat-stream-registry.service
import type { User, Workspace } from '@docmost/db/types/entity.types';
/**
* Wiring spec for the #184 phase 1.5 attach endpoint
* Wiring spec for the #184 phase 1.5 attach endpoint (tail-only #491)
* (`GET /ai-chat/runs/:chatId/stream`). Owner-gated via assertOwnedChat; the
* registry is mocked so this exercises ONLY the controller's replay/live/204/
* cleanup wiring against a fake raw socket. Constructor order is (aiChatService,
* registry is mocked so this exercises ONLY the controller's tail-write/live/204/
* cleanup wiring against a fake raw socket. The attach signature is now
* `(chatId, anchor, n, cb)` the client hands its persisted step frontier `n`
* and its assistant row id `anchor`. Constructor order is (aiChatService,
* aiChatRunService, aiChatRepo, aiChatMessageRepo, aiTranscription, pageRepo,
* streamRegistry, environment).
*/
@@ -86,8 +88,8 @@ describe('AiChatController attach endpoint (#184 phase 1.5)', () => {
attach: jest.fn(
(
_chatId: string,
_live: boolean,
_anchor: string | undefined,
_n: number,
cb: RunStreamCallbacks,
) => {
capturedCb = cb;
@@ -156,7 +158,7 @@ describe('AiChatController attach endpoint (#184 phase 1.5)', () => {
expect(res.hijack).not.toHaveBeenCalled();
});
it('threads expect=live and anchor through to the registry', async () => {
it('threads anchor and the numeric frontier n through to the registry', async () => {
const { controller, streamRegistry } = makeController({
chat: owned,
attachment: null,
@@ -165,8 +167,8 @@ describe('AiChatController attach endpoint (#184 phase 1.5)', () => {
const { req } = makeReq();
await controller.attachRunStream(
'c1',
'live',
'anchor-1',
'2',
req,
res,
user,
@@ -174,13 +176,44 @@ describe('AiChatController attach endpoint (#184 phase 1.5)', () => {
);
expect(streamRegistry.attach).toHaveBeenCalledWith(
'c1',
true,
'anchor-1',
2, // parsed to a number
expect.anything(),
);
});
it('passes expect=false when the query is absent', async () => {
it('#491: an ABSENT/invalid n passes null (not 0) so a finished run 204s (not-tail-aware)', async () => {
// Distinguishing a MISSING `n` from `n=0` is the #137/#161 dup guard: a
// parameterless/legacy tab must be handed null (-> the registry 204s a finished
// run) rather than frontier 0 (which would serve a finished non-rotated run's
// whole tail). MUTATION-VERIFY: revert to `Number(n) || 0` and this asserts 0.
const { controller, streamRegistry } = makeController({
chat: owned,
attachment: null,
});
for (const bad of [undefined, '', 'abc']) {
streamRegistry.attach.mockClear();
const { res } = makeRawRes();
const { req } = makeReq();
await controller.attachRunStream(
'c1',
undefined,
bad,
req,
res,
user,
workspace,
);
expect(streamRegistry.attach).toHaveBeenCalledWith(
'c1',
undefined,
null,
expect.anything(),
);
}
});
it('#491: a PRESENT n=0 passes 0 (tail-aware, distinct from absent)', async () => {
const { controller, streamRegistry } = makeController({
chat: owned,
attachment: null,
@@ -190,7 +223,7 @@ describe('AiChatController attach endpoint (#184 phase 1.5)', () => {
await controller.attachRunStream(
'c1',
undefined,
undefined,
'0',
req,
res,
user,
@@ -198,8 +231,8 @@ describe('AiChatController attach endpoint (#184 phase 1.5)', () => {
);
expect(streamRegistry.attach).toHaveBeenCalledWith(
'c1',
false,
undefined,
0,
expect.anything(),
);
});
@@ -245,8 +278,8 @@ describe('AiChatController attach endpoint (#184 phase 1.5)', () => {
const { req } = makeReq();
await controller.attachRunStream(
'c1',
'live',
'a1',
'1',
req,
res,
user,
@@ -0,0 +1,108 @@
import { ForbiddenException } from '@nestjs/common';
import { AiChatController } from './ai-chat.controller';
import type { User, Workspace } from '@docmost/db/types/entity.types';
/**
* Wiring spec for the #491 delta-poll endpoint (`POST /ai-chat/messages/delta`).
* Owner-gated via assertOwnedChat (same gate as the other reads), NOT flag-gated.
* The run fact rides IN the delta response (no separate /run poll). Hand-rolled
* mocks no Nest graph, no DB. Constructor order: (aiChatService,
* aiChatRunService, aiChatRepo, aiChatMessageRepo, aiTranscription, pageRepo).
*/
describe('AiChatController POST /ai-chat/messages/delta (#491)', () => {
const user = { id: 'u1' } as User;
const workspace = { id: 'ws1' } as Workspace;
function makeController(opts: {
chat?: unknown;
delta?: { rows: unknown[]; cursor: string };
run?: unknown;
}) {
const aiChatRunService = {
getLatestForChat: jest.fn().mockResolvedValue(opts.run),
};
const aiChatRepo = {
findById: jest.fn().mockResolvedValue(opts.chat),
};
const aiChatMessageRepo = {
findByChatUpdatedAfter: jest
.fn()
.mockResolvedValue(opts.delta ?? { rows: [], cursor: 'C1' }),
};
const controller = new AiChatController(
{} as never,
aiChatRunService as never,
aiChatRepo as never,
aiChatMessageRepo as never,
{} as never,
{} as never,
);
return { controller, aiChatRunService, aiChatRepo, aiChatMessageRepo };
}
it('owner-gates: a chat the user does not own throws, never reaching the repo', async () => {
const { controller, aiChatMessageRepo, aiChatRunService } = makeController({
chat: { id: 'c1', creatorId: 'someone-else' },
});
await expect(
controller.getMessagesDelta({ chatId: 'c1' }, user, workspace),
).rejects.toBeInstanceOf(ForbiddenException);
expect(aiChatMessageRepo.findByChatUpdatedAfter).not.toHaveBeenCalled();
expect(aiChatRunService.getLatestForChat).not.toHaveBeenCalled();
});
it('returns { rows, cursor, run:{id,status} } with the run fact inlined', async () => {
const rows = [{ id: 'm1' }];
const { controller } = makeController({
chat: { id: 'c1', creatorId: 'u1' },
delta: { rows, cursor: 'C2' },
run: { id: 'r1', status: 'running', error: 'ignored', stepCount: 3 },
});
const res = await controller.getMessagesDelta(
{ chatId: 'c1', cursor: 'C1' },
user,
workspace,
);
expect(res).toEqual({
rows,
cursor: 'C2',
// ONLY id + status — never the whole run row.
run: { id: 'r1', status: 'running' },
});
});
it('run is null when the chat has never had a run', async () => {
const { controller } = makeController({
chat: { id: 'c1', creatorId: 'u1' },
run: undefined,
});
const res = await controller.getMessagesDelta(
{ chatId: 'c1' },
user,
workspace,
);
expect(res.run).toBeNull();
});
it('passes cursor through, defaulting a missing cursor to null (first poll)', async () => {
const { controller, aiChatMessageRepo } = makeController({
chat: { id: 'c1', creatorId: 'u1' },
});
await controller.getMessagesDelta({ chatId: 'c1' }, user, workspace);
expect(aiChatMessageRepo.findByChatUpdatedAfter).toHaveBeenCalledWith(
'c1',
'ws1',
null,
);
await controller.getMessagesDelta(
{ chatId: 'c1', cursor: 'CX' },
user,
workspace,
);
expect(aiChatMessageRepo.findByChatUpdatedAfter).toHaveBeenLastCalledWith(
'c1',
'ws1',
'CX',
);
});
});
@@ -35,6 +35,7 @@ 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';
@@ -43,6 +44,8 @@ import {
AiChatRunHooks,
AiChatService,
AiChatStreamBody,
rowHasInlineParts,
hydrateAssistantParts,
} from './ai-chat.service';
import { AiChatRunService } from './ai-chat-run.service';
import { AiTranscriptionService } from './ai-transcription.service';
@@ -51,6 +54,7 @@ import {
ChatIdDto,
ExportChatDto,
GeneratePageTitleDto,
GetChatDeltaDto,
GetChatMessagesDto,
GetRunDto,
RenameChatDto,
@@ -63,6 +67,47 @@ import {
SUBSCRIBER_MAX_BUFFERED_BYTES,
} from './ai-chat-stream-registry.service';
import { startSseHeartbeat } from './sse-resilience';
/**
* Write the attach TAIL to the hijacked socket in chunks that RESPECT drain
* (#491): each `write()` that returns false (the kernel buffer is full) is awaited
* on the next 'drain' before continuing. The old code wrote the whole buffer
* synchronously, which with the pre-#491 32MB ring spiked memory (half the
* OOM). Bails immediately if the socket ended/errored mid-write. Frames that the
* paused registry subscriber buffers while this awaits are delivered by start().
*/
async function writeTailRespectingDrain(
raw: {
write(chunk: string): boolean;
writableEnded?: boolean;
destroyed?: boolean;
once(event: string, cb: () => void): unknown;
removeListener?(event: string, cb: () => void): unknown;
},
frames: string[],
): Promise<void> {
for (const frame of frames) {
if (raw.writableEnded || raw.destroyed) return;
const ok = raw.write(frame);
if (!ok) {
// Kernel buffer full — wait for drain (or an early close/error) before the
// next chunk, so a slow reader never forces the whole tail into memory.
// Remove ALL three listeners once any fires, so a many-chunk tail with
// repeated backpressure never leaks (MaxListenersExceededWarning).
await new Promise<void>((resolve) => {
const finish = (): void => {
raw.removeListener?.('drain', finish);
raw.removeListener?.('close', finish);
raw.removeListener?.('error', finish);
resolve();
};
raw.once('drain', finish);
raw.once('close', finish);
raw.once('error', finish);
});
}
}
}
import { EnvironmentService } from '../../integrations/environment/environment.service';
/**
@@ -87,8 +132,39 @@ 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')
@@ -142,11 +218,60 @@ export class AiChatController {
@AuthWorkspace() workspace: Workspace,
) {
await this.assertOwnedChat(dto.chatId, user, workspace);
return this.aiChatMessageRepo.findByChat(
const page = await 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),
};
}
/**
* Delta poll (#491) the degraded-poll fallback's payload. Returns the chat's
* message rows changed since `cursor` (a DB-clock timestamp from the previous
* poll), a FRESH cursor, AND the current run fact `{ id, status } | null`. This
* replaces the old degraded poll that refetched ALL infinite-query pages (full
* parts) every 2.5s: the client seeds once and thereafter merges only the
* deltas by id (the overlap window guarantees repeats the merge is idempotent,
* see mergeById). The run fact rides IN the delta (a separate /run poll would
* double the poll QPS), so the client FSM gets the run's status on the same tick.
* Owner-gated via assertOwnedChat (same gate as the other read endpoints).
*/
@HttpCode(HttpStatus.OK)
@Post('messages/delta')
async getMessagesDelta(
@Body() dto: GetChatDeltaDto,
@AuthUser() user: User,
@AuthWorkspace() workspace: Workspace,
): Promise<{
rows: AiChatMessage[];
cursor: string;
run: { id: string; status: string } | null;
}> {
await this.assertOwnedChat(dto.chatId, user, workspace);
const { rows, cursor } =
await this.aiChatMessageRepo.findByChatUpdatedAfter(
dto.chatId,
workspace.id,
dto.cursor ?? null,
);
const run = await this.aiChatRunService.getLatestForChat(
dto.chatId,
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),
cursor,
run: run ? { id: run.id, status: run.status } : null,
};
}
/**
@@ -165,8 +290,10 @@ export class AiChatController {
@AuthWorkspace() workspace: Workspace,
): Promise<{ markdown: string }> {
const chat = await this.assertOwnedChat(dto.chatId, user, workspace);
const rows = await this.aiChatMessageRepo.findAllByChat(
dto.chatId,
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.
workspace.id,
);
const markdown = buildChatMarkdown({
@@ -206,7 +333,13 @@ export class AiChatController {
workspace.id,
)
: undefined;
return { run, message: message ?? null };
// #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 };
}
/**
@@ -249,19 +382,25 @@ export class AiChatController {
}
/**
* Attach to a chat's live run stream (#184 phase 1.5). A late/reloaded tab
* replays the frames buffered so far and then follows the live tail as a normal
* streamer. Owner-gated via assertOwnedChat (same gate as getRun). When there is
* nothing to resume no entry, a finished run without expect=live, an
* overflowed buffer, or an anchor that pins a DIFFERENT run the endpoint
* answers 204, the ONLY "nothing to resume" signal the AI SDK's reconnect
* accepts (it maps 204 to a silent no-op). With AI_CHAT_RESUMABLE_STREAM off the
* registry is never populated, so attach always 204s.
* Attach to a chat's live run stream from the client's step frontier (#184 phase
* 1.5, tail-only #491). A late/reloaded tab hands the server the step count it
* has PERSISTED (`n` = the seeded row's `metadata.stepsPersisted`) and its
* assistant row id (`anchor`); the registry answers with the TAIL past step `n`
* (a synthetic `start` frame + the buffered frames stamped >= n) and then the
* live tail. Owner-gated via assertOwnedChat (same gate as getRun). When there
* is nothing to resume no entry, a ring that does not cover the client's
* frontier (overflow gap, or the client's seed lagged a rotation), or an anchor
* that pins a DIFFERENT run (invariant 6) the endpoint answers 204, the ONLY
* "nothing to resume" signal the AI SDK's reconnect accepts (it maps 204 to a
* silent no-op); the client then refetches (a larger n) and re-attaches. With
* AI_CHAT_RESUMABLE_STREAM off the registry is never populated, so attach always
* 204s.
*
* `expect=live` opts into replaying a finished-but-retained run (safe only when
* the client stripped the streaming tail); `anchor` is the client's assistant
* row id, which must match this run's (invariant 6) or a foreign run's
* transcript would be replayed into the store.
* The step marker `n` comes ONLY from the client the server never reads the
* row to derive it, because a server-side n from a stale seed would open a
* silent one-step hole. The tail is written to the socket in CHUNKS respecting
* drain (writeTailRespectingDrain): the old code synchronously blasted the whole
* buffer, which with the old 32MB cap was half the OOM.
*/
@SkipTransform()
@UseGuards(JwtAuthGuard, UserThrottlerGuard)
@@ -269,39 +408,49 @@ export class AiChatController {
@Get('runs/:chatId/stream')
async attachRunStream(
@Param('chatId', new ParseUUIDPipe()) chatId: string,
@Query('expect') expect: string | undefined,
@Query('anchor') anchor: string | undefined,
@Query('n') n: string | undefined,
@Req() req: FastifyRequest,
@Res() res: FastifyReply,
@AuthUser() user: User,
@AuthWorkspace() workspace: Workspace,
): Promise<void> {
await this.assertOwnedChat(chatId, user, workspace); // same gate as getRun
// The client's persisted step frontier. #491: distinguish a MISSING/invalid `n`
// (null — a NOT-tail-aware, legacy/parameterless tab expecting the old
// "finished -> 204 -> poll" contract) from `n=0` (a tail-aware client with
// nothing persisted yet). Passing 0 for a missing `n` would serve a finished,
// non-rotated run's WHOLE tail and a parameterless client would append it onto
// the steps it already shows -> #137/#161 duplicate. null makes the registry
// 204 such a finished run (see attach); a tail-aware n=0 still resumes.
const frontier: number | null =
n === undefined || n === '' || !Number.isFinite(Number(n))
? null
: Math.max(0, Number(n));
// The per-subscriber backpressure cap tracks the (env-tunable) ring cap.
const subscriberCap =
this.streamRegistry?.subscriberMaxBufferedBytes ??
SUBSCRIBER_MAX_BUFFERED_BYTES;
let stopHeartbeat: () => void = () => undefined;
const attachment = await this.streamRegistry?.attach(
chatId,
expect === 'live',
anchor,
{
onFrame: (frame) => {
// Backpressure guard: 2x the replay cap, so the initial replay burst
// alone can never trip it; only a genuinely stalled socket can.
try {
if (res.raw.writableLength > SUBSCRIBER_MAX_BUFFERED_BYTES) {
res.raw.destroy(); // 'close' fires -> unsubscribe below
return;
}
if (!res.raw.writableEnded) res.raw.write(frame);
} catch {
res.raw.destroy();
const attachment = await this.streamRegistry?.attach(chatId, anchor, frontier, {
onFrame: (frame) => {
// Backpressure guard: 2x the ring cap, so the initial tail burst alone
// can never trip it; only a genuinely stalled socket can.
try {
if (res.raw.writableLength > subscriberCap) {
res.raw.destroy(); // 'close' fires -> unsubscribe below
return;
}
},
onEnd: () => {
stopHeartbeat();
if (!res.raw.writableEnded) res.raw.end();
},
if (!res.raw.writableEnded) res.raw.write(frame);
} catch {
res.raw.destroy();
}
},
);
onEnd: () => {
stopHeartbeat();
if (!res.raw.writableEnded) res.raw.end();
},
});
if (!attachment) {
res.status(204).send(); // the ONLY "nothing to resume" signal the SDK accepts
return;
@@ -330,13 +479,16 @@ export class AiChatController {
// deliberately NO Connection/Keep-Alive (hop-by-hop; Safari/HTTP2)
});
res.raw.flushHeaders?.();
for (const frame of attachment.replay) res.raw.write(frame);
// Write the tail in chunks respecting drain (not a synchronous blast, which
// was half the OOM). Frames the paused subscriber buffers meanwhile are
// drained by start() below; its cap is the backstop for a stalled socket.
await writeTailRespectingDrain(res.raw, attachment.replay);
if (attachment.finished) {
res.raw.end();
if (!res.raw.writableEnded) res.raw.end();
return;
}
stopHeartbeat = startSseHeartbeat(res.raw, 15_000);
attachment.start(); // drain pending accumulated during replay, go live
attachment.start(); // drain pending accumulated during the tail write, go live
} catch {
attachment.unsubscribe();
stopHeartbeat();
+252 -21
View File
@@ -22,6 +22,7 @@ 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';
@@ -189,10 +190,11 @@ 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 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.
// 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).
//
// 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.
@@ -517,6 +519,12 @@ 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
@@ -1113,8 +1121,34 @@ 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 }> = [
...oldHistory.map(rowToUiMessage),
...replayHistory.map(rowToUiMessage),
{
id: 'pending-user',
role: 'user',
@@ -1153,7 +1187,9 @@ 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` (the aborted assistant row replays via findRecent).
// 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.
// Append the new user turn (shape-only) so index -2 is the prior assistant.
const interrupted = isInterruptResume(
[...oldHistory, { role: 'user', status: null, metadata: null }],
@@ -1410,10 +1446,11 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
const baseTools = { ...external.tools, ...docmostTools };
// Deferred tool loading state (#332), scoped to THIS streaming loop:
// - `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.
// - `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.
// - `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
@@ -1543,30 +1580,79 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
// Per-step (non-terminal) update: persist the finished steps the moment a
// step ends. Tolerant — a failed update is logged and swallowed so it never
// throws into the stream. Keeps status 'streaming'.
const updateStreaming = async (): Promise<void> => {
if (!assistantId) return;
//
// #491: it now SIGNALS its outcome — the persisted `stepsPersisted` count on
// a CONFIRMED write, or null when it was skipped/failed. The caller rotates
// the run-stream registry ring ONLY on a non-null return (a confirmed
// persist), so a failed persist never rotates away a step nobody has (the
// classic inversion bug); a failure just makes the ring cover more.
const updateStreaming = async (): Promise<number | null> => {
if (!assistantId) return null;
// Cheap short-circuit once the turn is finalized (see `finalized` below).
// The AUTHORITATIVE guard is `onlyIfStreaming` on the UPDATE: a late
// fire-and-forget step update could still be in flight on another pool
// connection when finalize runs, so the SQL `WHERE status='streaming'`
// (not this flag) is what prevents it clobbering the terminal row.
if (finalized) return;
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;
try {
await this.aiChatMessageRepo.update(
assistantId,
workspace.id,
flushAssistant(capturedSteps, '', 'streaming', {
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,
}),
{ onlyIfStreaming: true },
);
});
await this.aiChatMessageRepo.update(
assistantId,
workspace.id,
flushed,
{ onlyIfStreaming: true },
);
}
return stepsPersisted;
} catch (err) {
this.logger.warn(
`Failed to update streaming assistant row: ${
err instanceof Error ? err.message : 'unknown error'
}`,
);
return null;
}
};
@@ -1742,7 +1828,24 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
// this point still recovers the step. Not awaited here (never block the
// stream), but SERIALIZED via stepUpdateChain so the writes commit in
// step order; updateStreaming is error-tolerant (logs + swallows).
stepUpdateChain = stepUpdateChain.then(() => updateStreaming());
// #491: on a CONFIRMED persist, rotate the run-stream registry ring to
// drop the now-on-disk steps (stamp < stepsPersisted). Gated on the
// resumable flag (same as open/bind) and identity-checked in the
// registry; a null return (skipped/failed) rotates NOTHING (auto-safe).
stepUpdateChain = stepUpdateChain.then(async () => {
const persisted = await updateStreaming();
if (
persisted != null &&
runId &&
this.environment?.isAiChatResumableStreamEnabled?.()
) {
this.streamRegistry?.confirmPersistedStep(
chatId,
runId,
persisted,
);
}
});
// #184: persist the run's progress (finished-step count). Fire-and-
// forget; the hook swallows its own errors.
if (runId) runHooks?.onStep?.(runId, capturedSteps.length);
@@ -2721,6 +2824,122 @@ 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 rowparts 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 +
@@ -2891,6 +3110,18 @@ export function flushAssistant(
// `parts`). Old rows have no marker and the legacy { output } shape; a
// dual-shape query branches on this. Old rows are deliberately NOT migrated.
toolTraceVersion: 2,
// #491 STEP MARKER: the number of FINISHED steps whose parts are in THIS row,
// written by the SAME flush that builds `parts` (atomically — they are both
// derived from `finished`, so the marker can NEVER disagree with the persisted
// parts). This is the step-alignment anchor the resume stack builds on:
// - the registry rotates its retention ring only on a CONFIRMED persist of
// step N (commit 3);
// - attach slices the tail at "step > N" from the client's persisted seed.
// It is NOT `run.stepCount`: recordStep is fire-and-forget and NOT atomic with
// the parts write, so stepCount could race ahead of the persisted parts
// (seed↔marker drift). The in-progress trailing text (an error/abort partial,
// or a mid-stream flush) is NOT a finished step and is excluded from the count.
stepsPersisted: finished.length,
};
// finishReason: prefer an explicit one; else derive a sensible value from the
// terminal status (so onError/onAbort records keep their historical reason).
@@ -0,0 +1,65 @@
import { flushAssistant } from './ai-chat.service';
/**
* #491 STEP MARKER `metadata.stepsPersisted` is written by the SAME flush that
* builds `metadata.parts`, so the marker can never disagree with the persisted
* parts (the step-alignment anchor the resume stack builds on). These are
* PROPERTY tests: they assert the marker tracks the number of FINISHED steps for
* every flush shape.
*/
// A finished step carrying one line of text and one tool call/result.
function step(i: number) {
return {
text: `step ${i}`,
toolCalls: [
{ toolCallId: `c${i}`, toolName: 'getPage', input: { id: `p${i}` } },
],
toolResults: [
{ toolCallId: `c${i}`, toolName: 'getPage', output: { title: `T${i}` } },
],
};
}
describe('flushAssistant step marker (#491)', () => {
it('seed (no steps) → stepsPersisted 0', () => {
const f = flushAssistant([], '', 'streaming');
expect(f.metadata.stepsPersisted).toBe(0);
});
it('PROPERTY: stepsPersisted equals the number of FINISHED steps, for any N', () => {
for (let n = 0; n <= 6; n++) {
const steps = Array.from({ length: n }, (_, i) => step(i));
const f = flushAssistant(steps, '', 'streaming');
expect(f.metadata.stepsPersisted).toBe(n);
// ...and the parts actually contain those N steps' text (marker agrees with
// the persisted parts — the atomicity the whole design relies on).
const parts = f.metadata.parts as Array<Record<string, unknown>>;
const textParts = parts.filter((p) => p.type === 'text');
expect(textParts).toHaveLength(n);
}
});
it('an in-progress trailing partial does NOT increment the marker', () => {
// 2 finished steps + a partial (not-yet-finished) trailing text: the marker
// counts only the CONFIRMED step boundaries, not the partial.
const f = flushAssistant([step(0), step(1)], 'partial third step', 'error', {
error: 'boom',
});
expect(f.metadata.stepsPersisted).toBe(2);
// The partial text IS persisted in parts (so the user sees it), but it is not a
// counted step.
const parts = f.metadata.parts as Array<Record<string, unknown>>;
expect(parts[parts.length - 1]).toEqual({
type: 'text',
text: 'partial third step',
});
});
it('terminal completed flush counts all finished steps', () => {
const f = flushAssistant([step(0), step(1), step(2)], '', 'completed', {
finishReason: 'stop',
});
expect(f.metadata.stepsPersisted).toBe(3);
});
});
@@ -1,4 +1,10 @@
import { IsOptional, IsString, MaxLength, MinLength } from 'class-validator';
import {
IsISO8601,
IsOptional,
IsString,
MaxLength,
MinLength,
} from 'class-validator';
/** Identify a chat by id (workspace-scoped on the server). */
export class ChatIdDto {
@@ -37,6 +43,24 @@ export class GetChatMessagesDto {
cursor?: string;
}
/**
* Delta poll (#491): pull the chat's rows changed since `cursor` (a DB-clock
* timestamp from the previous poll) plus the current run fact the degraded-poll
* fallback's payload, replacing the full infinite-query refetch. Omit `cursor` on
* the first poll (returns just a fresh cursor to start the chain).
*/
export class GetChatDeltaDto {
@IsString()
chatId: string;
// ISO-8601 timestamp echoed from the previous poll's response. Validated as
// ISO-8601 (not a bare string): a malformed cursor would otherwise reach the
// `::timestamptz` cast in findByChatUpdatedAfter and 500 instead of a clean 400.
@IsOptional()
@IsISO8601()
cursor?: string;
}
/** Resolve the chat bound to a document (the page's most-recent owned chat). */
export class BoundChatDto {
@IsString()
@@ -0,0 +1,193 @@
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();
});
@@ -0,0 +1,201 @@
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 };
}
}
@@ -0,0 +1,19 @@
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 {}
@@ -0,0 +1,299 @@
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();
});
});
@@ -0,0 +1,191 @@
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 23 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
}`,
);
});
}
}
@@ -0,0 +1,51 @@
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;
}
@@ -0,0 +1,6 @@
import { IsUUID } from 'class-validator';
export class RevokeApiKeyDto {
@IsUUID()
id: string;
}
@@ -20,3 +20,71 @@ 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);
});
});
+13 -1
View File
@@ -207,8 +207,20 @@ export class AuthController {
async collabToken(
@AuthUser() user: User,
@AuthWorkspace() workspace: Workspace,
@Req() req: FastifyRequest,
) {
return this.authService.getCollabToken(user, workspace.id);
// 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);
}
@SkipThrottle({ [AUTH_THROTTLER]: true })
+5 -1
View File
@@ -5,9 +5,13 @@ 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({
imports: [TokenModule, WorkspaceModule],
// 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],
controllers: [AuthController],
providers: [AuthService, SignupService, JwtStrategy],
exports: [SignupService, AuthService],
@@ -33,6 +33,17 @@ 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;
@@ -44,6 +55,13 @@ 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 = {
@@ -375,10 +375,20 @@ export class AuthService {
}
}
async getCollabToken(user: User, workspaceId: string) {
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 },
) {
const token = await this.tokenService.generateCollabToken(
user,
workspaceId,
undefined,
apiKey,
);
return { token };
}
@@ -1,4 +1,5 @@
import { ForbiddenException, UnauthorizedException } from '@nestjs/common';
import * as jwt from 'jsonwebtoken';
import { TokenService } from './token.service';
import { JwtType } from '../dto/jwt-payload';
@@ -213,4 +214,175 @@ 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';
}
@@ -4,7 +4,6 @@ import {
UnauthorizedException,
} from '@nestjs/common';
import { JwtService } from '@nestjs/jwt';
import type { StringValue } from 'ms';
import { EnvironmentService } from '../../../integrations/environment/environment.service';
import {
JwtApiKeyPayload,
@@ -62,6 +61,12 @@ 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();
@@ -71,6 +76,10 @@ 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 }
: {}),
@@ -123,9 +132,8 @@ export class TokenService {
apiKeyId: string;
user: User;
workspaceId: string;
expiresIn?: StringValue | number;
}): Promise<string> {
const { apiKeyId, user, workspaceId, expiresIn } = opts;
const { apiKeyId, user, workspaceId } = opts;
if (isUserDisabled(user)) {
throw new ForbiddenException();
}
@@ -137,7 +145,32 @@ export class TokenService {
type: JwtType.API_KEY,
};
return this.jwtService.sign(payload, expiresIn ? { expiresIn } : {});
// 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;
}
async generatePdfRenderToken(
@@ -177,4 +210,31 @@ 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;
}
}
@@ -22,7 +22,8 @@ describe('JwtStrategy — provenance derivation', () => {
const userSessionRepo: any = { findActiveById: jest.fn() };
const sessionActivityService: any = { trackActivity: jest.fn() };
const environmentService: any = { getAppSecret: () => 'test-secret' };
const moduleRef: any = {};
// ACCESS-path tests never touch the api-key seam; a bare stub suffices.
const apiKeyService: any = { validate: jest.fn() };
const strategy = new JwtStrategy(
userRepo,
@@ -30,7 +31,7 @@ describe('JwtStrategy — provenance derivation', () => {
userSessionRepo,
sessionActivityService,
environmentService,
moduleRef,
apiKeyService,
);
return { strategy, userRepo };
}
@@ -122,25 +123,29 @@ describe('JwtStrategy — provenance derivation', () => {
});
/**
* Provenance derivation on the API-KEY path (jwt.strategy.validateApiKey, #486).
* Provenance derivation on the API-KEY path (jwt.strategy.validateApiKey, #486
* + #501).
*
* 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.validateApiKey: isAgent -> 'agent',
* SERVER-SIDE user returned by ApiKeyService.validate: isAgent -> 'agent',
* otherwise 'user'; aiChatId is always null (an API key has no ai_chats row).
*
* The enterprise ApiKeyService is not bundled in the OSS build, so the strategy
* loads it through an overridable `resolveApiKeyService` seam that we stub here.
* #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.
*/
describe('JwtStrategy — API-key provenance derivation (#486)', () => {
function makeApiKeyStrategy(validateApiKeyImpl: (p: any) => Promise<any>) {
describe('JwtStrategy — API-key provenance derivation (#486/#501)', () => {
function makeApiKeyStrategy(validateImpl: (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 moduleRef: any = {};
const validate = jest.fn(validateImpl);
const apiKeyService: any = { validate };
const strategy = new JwtStrategy(
userRepo,
@@ -148,14 +153,9 @@ describe('JwtStrategy — API-key provenance derivation (#486)', () => {
userSessionRepo,
sessionActivityService,
environmentService,
moduleRef,
apiKeyService,
);
// 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 };
return { strategy, validate };
}
const makeReq = () => ({ raw: {} as Record<string, any> });
@@ -166,22 +166,23 @@ describe('JwtStrategy — API-key provenance derivation (#486)', () => {
type: JwtType.API_KEY,
});
it("stamps actor='agent' for an is_agent API key (from the validated user)", async () => {
it("stamps actor='agent' + authType/apiKeyId for an is_agent API key", async () => {
const validated = {
user: { id: 'svc-1', isAgent: true },
workspace: { id: 'ws-1' },
};
const { strategy, validateApiKey } = makeApiKeyStrategy(
async () => validated,
);
const { strategy, validate } = makeApiKeyStrategy(async () => validated);
const req = makeReq();
const result = await strategy.validate(req, apiKeyPayload() as any);
expect(validateApiKey).toHaveBeenCalledTimes(1);
expect(validate).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);
});
@@ -197,25 +198,19 @@ describe('JwtStrategy — API-key provenance derivation (#486)', () => {
expect(req.raw.actor).toBe('user');
expect(req.raw.aiChatId).toBeNull();
expect(req.raw.authType).toBe('api_key');
});
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);
it('propagates a validate() rejection and stamps nothing', async () => {
const { strategy } = makeApiKeyStrategy(async () => {
throw new UnauthorizedException();
});
const req = makeReq();
await expect(
strategy.validate(req, apiKeyPayload() as any),
).rejects.toThrow(UnauthorizedException);
expect(req.raw.actor).toBeUndefined();
expect(req.raw.authType).toBeUndefined();
});
});
@@ -1,4 +1,4 @@
import { Injectable, Logger, UnauthorizedException } from '@nestjs/common';
import { Injectable, UnauthorizedException } from '@nestjs/common';
import { PassportStrategy } from '@nestjs/passport';
import { Strategy } from 'passport-jwt';
import { EnvironmentService } from '../../../integrations/environment/environment.service';
@@ -9,20 +9,18 @@ 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 moduleRef: ModuleRef,
private readonly apiKeyService: ApiKeyService,
) {
super({
jwtFromRequest: (req: FastifyRequest) => {
@@ -102,12 +100,17 @@ export class JwtStrategy extends PassportStrategy(Strategy, 'jwt') {
}
private async validateApiKey(req: any, payload: JwtApiKeyPayload) {
const apiKeyService = this.resolveApiKeyService();
if (!apiKeyService) {
throw new UnauthorizedException('Enterprise API Key module missing');
}
// 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 result = await apiKeyService.validateApiKey(payload);
// 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;
// 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
@@ -119,32 +122,10 @@ 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 as any)?.user, null);
const provenance = resolveProvenance(result.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;
}
}
}
+2
View File
@@ -6,6 +6,7 @@ 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';
@@ -29,6 +30,7 @@ import { ClsMiddleware } from 'nestjs-cls';
imports: [
UserModule,
AuthModule,
ApiKeyModule,
WorkspaceModule,
PageModule,
AttachmentModule,
@@ -32,11 +32,13 @@ 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';
@@ -124,11 +126,13 @@ import { firstSqlToken } from '../integrations/metrics/metrics.constants';
AiChatRepo,
AiChatMessageRepo,
AiChatRunRepo,
AiChatRunStepRepo,
AiChatPageSnapshotRepo,
AiProviderCredentialsRepo,
AiMcpServerRepo,
AiAgentRoleRepo,
PageEmbeddingRepo,
ApiKeyRepo,
PageListener,
],
exports: [
@@ -159,11 +163,13 @@ import { firstSqlToken } from '../integrations/metrics/metrics.constants';
AiChatRepo,
AiChatMessageRepo,
AiChatRunRepo,
AiChatRunStepRepo,
AiChatPageSnapshotRepo,
AiProviderCredentialsRepo,
AiMcpServerRepo,
AiAgentRoleRepo,
PageEmbeddingRepo,
ApiKeyRepo,
],
})
export class DatabaseModule implements OnApplicationBootstrap {
@@ -0,0 +1,27 @@
import { type Kysely } from 'kysely';
/**
* #370 page-versioning intentionality tier on a history snapshot.
*
* Adds `page_history.kind`, the three-tier "how intentional was this snapshot"
* marker that lets versions (intentional points) be told apart from autosaves:
* - 'manual' a human explicitly saved a version (Cmd+S / Save button)
* - 'agent' the AI agent explicitly saved a version
* - 'idle' trailing idle-flush autosnapshot (safety net)
* - 'boundary' autosnapshot pinned on a source transition (useragentgit)
*
* Nullable with NO default (mirrors last_updated_source in the agent-provenance
* migration): legacy rows predate the marker and read back as `null`, which the
* client renders as a plain autosave. Stored as a short varchar to stay
* forward-compatible without an enum migration.
*/
export async function up(db: Kysely<any>): Promise<void> {
await db.schema
.alterTable('page_history')
.addColumn('kind', 'varchar(20)', (col) => col)
.execute();
}
export async function down(db: Kysely<any>): Promise<void> {
await db.schema.alterTable('page_history').dropColumn('kind').execute();
}
@@ -0,0 +1,70 @@
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();
}
@@ -0,0 +1,280 @@
import { randomUUID } from 'crypto';
import { CamelCasePlugin, Kysely, sql } from 'kysely';
import { PostgresJSDialect } from 'kysely-postgres-js';
// NOT a default import: the project tsconfig is `module: commonjs` with NO
// esModuleInterop, so `import postgres from 'postgres'` compiles to
// `postgres_1.default(...)` and the CJS `postgres` export has no `.default` —
// it threw in beforeAll, was swallowed as "DB unreachable", and SILENTLY voided
// all six tests. Mirror the working integration harness (test/integration/db.ts).
import * as postgres from 'postgres';
import { AiChatMessageRepo } from './ai-chat-message.repo';
import { AiChatRunRepo } from './ai-chat-run.repo';
/**
* #491 delta-poll OBSERVABLE-PROPERTY tests against a LIVE Postgres (the local
* gitmost test DB, docker `gitmost-test-pg` on :5432), not "rows through a mock"
* (a mock cannot observe the DB clock nor the overlap-window race the very
* things that matter here). Drives the REAL repo methods (`findByChatUpdatedAfter`,
* the now()-stamped `update`) and asserts:
* 1. delta-relevant writes stamp `updatedAt` from the DB clock, not the app clock
* (proven by faking the process clock far into the future and observing the
* stamp stays on real DB time);
* 2. the poll returns only rows changed after the cursor, ordered, with a fresh
* DB-clock cursor;
* 3. the "committed late but stamped earlier than the cursor" RACE is caught by
* the overlap window (a naive `updatedAt > cursor` would MISS it);
* 4. the overlap GUARANTEES repeats across close polls the contract behind the
* client's idempotent merge (mergeById).
*
* INTEGRATION lane (`*.int-spec.ts`): runs under `test:int`, whose global-setup
* DROPS + RE-CREATES + MIGRATES `docmost_test`, so the real `ai_chat_messages` /
* `ai_chat_runs` tables EXIST here. (It was previously a `.spec.ts` defaulting to
* the UNmigrated dev `docmost`; in the CI unit lane where `WAL_TEST_DATABASE_URL`
* is unset and only `test:int` migrates that meant 5/6 ERROR
* `relation "ai_chat_messages" does not exist`, silently voiding coverage of the
* risky cursor/overlap logic. Renaming to `.int-spec.ts` + defaulting the DSN to
* `docmost_test` fixes the CI fidelity.)
*
* FK triggers are bypassed (`session_replication_role = replica`) so synthetic
* chat/workspace ids need no parent fixtures; a single pooled connection (max 1)
* keeps that session setting for every query. SKIPS cleanly when the DB is
* unreachable so a DB-less CI never breaks.
*/
const CONN =
process.env.WAL_TEST_DATABASE_URL ??
process.env.TEST_DATABASE_URL ??
'postgresql://docmost:docmost_dev_pw@localhost:5432/docmost_test';
let db: Kysely<any>;
let sqlClient: ReturnType<typeof postgres>;
let msgRepo: AiChatMessageRepo;
let runRepo: AiChatRunRepo;
let reachable = false;
const workspaceId = randomUUID();
const chatId = randomUUID();
beforeAll(async () => {
try {
sqlClient = postgres(CONN, { max: 1, onnotice: () => {} });
db = new Kysely<any>({
dialect: new PostgresJSDialect({ postgres: sqlClient }),
plugins: [new CamelCasePlugin()],
});
// Single connection keeps this session-scoped bypass for the whole suite.
await sql`set session_replication_role = replica`.execute(db);
await sql`select 1`.execute(db);
reachable = true;
} catch (err) {
reachable = false;
// A genuine connection failure (ECONNREFUSED etc.) is a legitimate skip on a
// DB-less CI. A PROGRAMMING error (bad import, typo, driver misuse) must NOT
// masquerade as "DB unreachable" and silently void the whole suite (that is
// exactly the bug that hid this spec's zero coverage) — rethrow it so the
// suite fails LOUDLY.
const msg = String((err as Error)?.message ?? err);
if (
!/ECONNREFUSED|ENOTFOUND|ETIMEDOUT|EHOSTUNREACH|connect|terminating|password|authentication|role .* does not exist|database .* does not exist/i.test(
msg,
)
) {
throw err;
}
}
msgRepo = new AiChatMessageRepo(db as never);
runRepo = new AiChatRunRepo(db as never);
});
afterAll(async () => {
if (db) {
try {
await db
.deleteFrom('aiChatMessages')
.where('workspaceId', '=', workspaceId)
.execute();
await db
.deleteFrom('aiChatRuns')
.where('workspaceId', '=', workspaceId)
.execute();
} catch {
/* best-effort cleanup */
}
await db.destroy();
}
});
afterEach(() => {
jest.useRealTimers();
});
async function seedMessage(overrides: Record<string, unknown> = {}) {
return msgRepo.insert({
chatId,
workspaceId,
userId: null as never,
role: 'assistant',
content: 'x',
status: 'streaming',
...overrides,
} as never);
}
async function dbNow(): Promise<string> {
const r = await sql<{ now: Date }>`select now() as now`.execute(db);
return r.rows[0].now.toISOString();
}
// Fake ONLY the Date object (so in-process `new Date()`/`Date.now()` jump), while
// leaving every TIMER function real. Faking timers wholesale freezes postgres.js's
// internal connection/query timers, so the awaited DB round-trip would hang the
// test (and the afterAll cleanup) at the jest 5s cap. With Date-only faking the
// query resolves normally, and we still prove the stamp is the DB clock (not the
// skewed process clock).
function fakeDateOnly(iso: string): void {
jest.useFakeTimers({
doNotFake: [
'hrtime',
'nextTick',
'performance',
'queueMicrotask',
'requestAnimationFrame',
'cancelAnimationFrame',
'requestIdleCallback',
'cancelIdleCallback',
'setImmediate',
'clearImmediate',
'setInterval',
'clearInterval',
'setTimeout',
'clearTimeout',
],
now: new Date(iso),
});
}
const maybe = (name: string, fn: () => Promise<void>) =>
it(name, async () => {
if (!reachable) {
console.warn(`SKIP (${name}): test DB unreachable at ${CONN}`);
return;
}
await fn();
});
describe('AiChatMessageRepo.findByChatUpdatedAfter (#491 delta poll)', () => {
maybe('null cursor returns no rows and a fresh DB-clock cursor', async () => {
const before = await dbNow();
const { rows, cursor } = await msgRepo.findByChatUpdatedAfter(
chatId,
workspaceId,
null,
);
expect(rows).toEqual([]);
expect(new Date(cursor).getTime()).toBeGreaterThanOrEqual(
new Date(before).getTime(),
);
});
maybe('returns only rows changed after the cursor', async () => {
const { cursor: c0 } = await msgRepo.findByChatUpdatedAfter(
chatId,
workspaceId,
null,
);
const m = await seedMessage();
const { rows, cursor: c1 } = await msgRepo.findByChatUpdatedAfter(
chatId,
workspaceId,
c0,
);
expect(rows.map((r) => r.id)).toContain(m.id);
// Cursor is monotonic (advances).
expect(new Date(c1).getTime()).toBeGreaterThanOrEqual(
new Date(c0).getTime(),
);
});
maybe(
'RACE: a row stamped BEFORE the cursor but seen after is caught by the overlap',
async () => {
// Cursor taken now; then a row appears whose updatedAt is 2s in the PAST
// (committed late on another connection but stamped earlier). A naive
// `updatedAt > cursor` would MISS it; the 5s overlap window catches it.
const cursor = await dbNow();
const m = await seedMessage();
await sql`update ai_chat_messages set updated_at = now() - interval '2 seconds' where id = ${m.id}`.execute(
db,
);
const { rows } = await msgRepo.findByChatUpdatedAfter(
chatId,
workspaceId,
cursor,
);
expect(rows.map((r) => r.id)).toContain(m.id);
},
);
maybe(
'overlap GUARANTEES repeats across close polls (idempotent-merge contract)',
async () => {
const { cursor: c0 } = await msgRepo.findByChatUpdatedAfter(
chatId,
workspaceId,
null,
);
const m = await seedMessage();
const first = await msgRepo.findByChatUpdatedAfter(
chatId,
workspaceId,
c0,
);
expect(first.rows.map((r) => r.id)).toContain(m.id);
// Immediately re-poll with the JUST-returned cursor: the row is still within
// the overlap window, so it is returned AGAIN — the client MUST dedupe by id.
const second = await msgRepo.findByChatUpdatedAfter(
chatId,
workspaceId,
first.cursor,
);
expect(second.rows.map((r) => r.id)).toContain(m.id);
},
);
maybe(
'update() stamps updatedAt from the DB clock, not the app clock',
async () => {
const m = await seedMessage();
// Skew the PROCESS clock ~73 years into the future (Date only). If the stamp
// came from `new Date()` the row would read year 2099; sql now() keeps it on
// DB time.
fakeDateOnly('2099-01-01T00:00:00Z');
const updated = await msgRepo.update(m.id, workspaceId, {
content: 'y',
});
jest.useRealTimers();
expect(updated).toBeDefined();
expect(new Date(updated!.updatedAt).getFullYear()).toBeLessThan(2099);
},
);
maybe(
'run update() also stamps updatedAt from the DB clock',
async () => {
const run = await runRepo.insert({
chatId,
workspaceId,
createdBy: null as never,
trigger: 'user',
status: 'running',
stepCount: 0,
} as never);
fakeDateOnly('2099-01-01T00:00:00Z');
const updated = await runRepo.update(run.id, workspaceId, {
stepCount: 1,
});
jest.useRealTimers();
expect(updated).toBeDefined();
expect(new Date(updated!.updatedAt).getFullYear()).toBeLessThan(2099);
},
);
});
@@ -25,6 +25,20 @@ const SWEEP_STREAMING_STALE_MS = 10 * 60 * 1000; // 10 minutes
// into memory; far above any realistic transcript length.
const FIND_ALL_BY_CHAT_LIMIT = 5000;
// Delta-poll overlap (#491): the poll query reaches this far BEHIND the client's
// echoed cursor, so a row that committed with an `updatedAt` marginally before the
// previous cursor was taken (on another autocommit connection) is still caught.
// Sized well above realistic single-row commit skew; the client merge is
// idempotent by id (mergeById), so the guaranteed repeats the overlap produces are
// harmless.
export const DELTA_POLL_OVERLAP_SECONDS = 5;
// Hard cap on rows one delta poll returns — a safety bound (a poll should carry a
// handful of just-changed rows, never a whole transcript). Ordered by (updatedAt,
// id) asc, so on the pathological overflow the OLDEST changes win and the newest
// are picked up by the next poll (its cursor did not advance past them).
export const DELTA_POLL_MAX_ROWS = 500;
@Injectable()
export class AiChatMessageRepo {
private readonly logger = new Logger(AiChatMessageRepo.name);
@@ -139,6 +153,72 @@ export class AiChatMessageRepo {
.executeTakeFirst();
}
/**
* Delta read (#491) for the degraded poll: the chat's messages whose row
* changed AFTER the client's `cursor`, plus a FRESH cursor taken from the DB
* clock. Replaces the old "refetch ALL infinite-query pages every 2.5s with
* full parts" poll the client seeds once (findByChat) and thereafter pulls
* only the deltas and merges them by id (mergeById).
*
* Cursor: a DB-clock timestamp (now()) the client echoes back each poll. All
* delta-relevant writes stamp `updatedAt` with now() (see `update` /
* `finalizeOwner`), so this is a SINGLE monotonic axis. The query overlaps the
* cursor by DELTA_POLL_OVERLAP_SECONDS to catch a row committed with an
* `updatedAt` marginally BEFORE the previous cursor was taken on another
* connection (single-row autocommit UPDATEs; no long transactions). The overlap
* GUARANTEES occasional REPEATS, so the client merge MUST be idempotent by id.
*
* `cursor === null` (first poll after the full seed) returns NO rows there is
* nothing "new" relative to a just-loaded seed only the fresh cursor to start
* the delta chain. The fresh cursor is read AFTER the rows, so it is >= every
* returned row's `updatedAt` (they were read strictly earlier) a row that
* commits between the rows-read and the cursor-read is at most
* DELTA_POLL_OVERLAP_SECONDS behind the returned cursor, so the next poll's
* overlap window always re-includes it (no miss).
*/
async findByChatUpdatedAfter(
chatId: string,
workspaceId: string,
cursor: string | null,
): Promise<{ rows: AiChatMessage[]; cursor: string }> {
if (cursor === null) {
const nowRow = await sql<{ now: Date }>`select now() as now`.execute(
this.db,
);
return { rows: [], cursor: nowRow.rows[0].now.toISOString() };
}
// Overlap the client cursor by DELTA_POLL_OVERLAP_SECONDS, computed in SQL off
// the echoed cursor so the whole comparison stays on the DB clock.
const rows = await this.db
.selectFrom('aiChatMessages')
.select(this.baseFields)
.where('chatId', '=', chatId)
.where('workspaceId', '=', workspaceId)
.where('deletedAt', 'is', null)
.where(
'updatedAt',
'>',
sql<Date>`${cursor}::timestamptz - make_interval(secs => ${DELTA_POLL_OVERLAP_SECONDS})`,
)
.orderBy('updatedAt', 'asc')
.orderBy('id', 'asc')
.limit(DELTA_POLL_MAX_ROWS)
.execute();
// When the page filled (pathological overflow), DO NOT advance the cursor to
// now(): that would skip the changed rows past the cap that this poll did not
// return. Resume from the last returned row's updatedAt instead (the next
// poll's overlap re-includes ties by id). In the normal case the fresh DB-clock
// now() is the cursor.
if (rows.length === DELTA_POLL_MAX_ROWS) {
return {
rows,
cursor: rows[rows.length - 1].updatedAt.toISOString(),
};
}
const nowRow = await sql<{ now: Date }>`select now() as now`.execute(this.db);
return { rows, cursor: nowRow.rows[0].now.toISOString() };
}
async insert(
insertable: InsertableAiChatMessage,
trx?: KyselyTransaction,
@@ -172,7 +252,13 @@ export class AiChatMessageRepo {
const db = dbOrTx(this.db, opts?.trx);
let query = db
.updateTable('aiChatMessages')
.set({ ...(patch as Record<string, unknown>), updatedAt: new Date() })
// #491: stamp `updatedAt` from the DB clock (sql now()), NOT the app clock
// (new Date()). The delta-poll cursor (findByChatUpdatedAfter) is a single
// DB-clock axis; a per-step 'streaming' UPDATE stamped with the app clock
// would be a SECOND, skewed clock source and could leave a row's updatedAt
// just under a cursor taken from now() on another connection — an
// independent source of delta MISSES. All delta-relevant writes use now().
.set({ ...(patch as Record<string, unknown>), updatedAt: sql`now()` })
.where('id', '=', id)
.where('workspaceId', '=', workspaceId);
// Concurrency guard (#183 review): a per-step 'streaming' update must NEVER
@@ -214,7 +300,9 @@ export class AiChatMessageRepo {
const db = dbOrTx(this.db, trx);
return db
.updateTable('aiChatMessages')
.set({ ...(patch as Record<string, unknown>), updatedAt: new Date() })
// #491: DB-clock stamp (see `update`) — this terminal write flips the row's
// status, which the delta poll must observe on the shared now() cursor axis.
.set({ ...(patch as Record<string, unknown>), updatedAt: sql`now()` })
.where('id', '=', id)
.where('workspaceId', '=', workspaceId)
.where((eb) =>
@@ -249,7 +337,9 @@ export class AiChatMessageRepo {
.set({
status,
metadata: sql`coalesce(metadata, '{}'::jsonb) || jsonb_build_object('finalizeFailed', true)`,
updatedAt: new Date(),
// #491: DB-clock stamp (see `update`) so a reconcile status flip lands on
// the same now() cursor axis the delta poll reads.
updatedAt: sql`now()`,
})
.where('id', '=', id)
.where('workspaceId', '=', workspaceId)
@@ -307,7 +397,9 @@ export class AiChatMessageRepo {
.set({
status: 'aborted',
metadata: sql`coalesce(m.metadata, '{}'::jsonb) || jsonb_build_object('finalizeFailed', true)`,
updatedAt: new Date(),
// #491: DB-clock stamp (see `update`). The staleness WHERE below stays on
// the app clock — a >minutes window makes the ms-scale skew irrelevant.
updatedAt: sql`now()`,
})
.where('m.status', '=', 'streaming')
.where('m.updatedAt', '<', staleBefore)
@@ -351,7 +443,8 @@ export class AiChatMessageRepo {
.set({
status: 'aborted',
metadata: sql`coalesce(metadata, '{}'::jsonb) || jsonb_build_object('finalizeFailed', true)`,
updatedAt: new Date(),
// #491: DB-clock stamp (see `update`). Staleness WHERE stays app-clock.
updatedAt: sql`now()`,
})
.where('status', '=', 'streaming')
.where('updatedAt', '<', staleBefore)
@@ -0,0 +1,95 @@
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;
}
}
@@ -62,10 +62,17 @@ describe('AiChatRunRepo.sweepRunning', () => {
// ...but a fresh 'running' run (updatedAt = now) must NOT be skipped: no
// updatedAt predicate at all on the boot path.
expect(rec.wheres.some(([col]) => col === 'updatedAt')).toBe(false);
// It flips to 'aborted' and stamps finishedAt.
expect(rec.set).toEqual(
expect.objectContaining({ status: 'aborted', finishedAt: expect.any(Date) }),
);
// It flips to 'aborted' and stamps finishedAt + updatedAt. #491: the stamps
// are now DB-clock `sql now()` expressions (raw builders), NOT app-clock
// `new Date()`, so the run row shares the delta poll's single now() cursor axis
// — assert they are present and are the sql raw-builder objects (not a Date,
// not undefined).
expect(rec.set?.status).toBe('aborted');
for (const stamp of ['finishedAt', 'updatedAt'] as const) {
expect(rec.set?.[stamp]).toBeDefined();
expect(rec.set?.[stamp]).not.toBeInstanceOf(Date);
expect(typeof rec.set?.[stamp]).toBe('object');
}
});
it('phase-2 path: an explicit staleMs reintroduces the updatedAt window', async () => {
@@ -136,7 +136,11 @@ export class AiChatRunRepo {
const db = dbOrTx(this.db, trx);
return db
.updateTable('aiChatRuns')
.set({ ...(patch as Record<string, unknown>), updatedAt: new Date() })
// #491: DB-clock stamp (sql now()) so the run row shares the delta poll's
// single now() cursor axis with the assistant message rows — a run-status
// change (the run fact the delta carries) must never sit on a skewed app
// clock relative to the message updatedAt cursor.
.set({ ...(patch as Record<string, unknown>), updatedAt: sql`now()` })
.where('id', '=', id)
.where('workspaceId', '=', workspaceId)
.returning(this.baseFields)
@@ -162,14 +166,15 @@ export class AiChatRunRepo {
trx?: KyselyTransaction,
): Promise<AiChatRun | undefined> {
const db = dbOrTx(this.db, trx);
const now = new Date();
return db
.updateTable('aiChatRuns')
.set({
status: patch.status,
error: patch.error,
finishedAt: now,
updatedAt: now,
// #491: DB-clock stamps (finished_at + updated_at) so the terminal run
// fact lands on the delta poll's now() cursor axis.
finishedAt: sql`now()`,
updatedAt: sql`now()`,
})
.where('id', '=', id)
.where('workspaceId', '=', workspaceId)
@@ -192,7 +197,8 @@ export class AiChatRunRepo {
const db = dbOrTx(this.db, trx);
return db
.updateTable('aiChatRuns')
.set({ stopRequestedAt: new Date(), updatedAt: new Date() })
// #491: DB-clock stamps (see `update`).
.set({ stopRequestedAt: sql`now()`, updatedAt: sql`now()` })
.where('id', '=', id)
.where('workspaceId', '=', workspaceId)
.where('status', 'in', ACTIVE_RUN_STATUSES as unknown as string[])
@@ -249,13 +255,14 @@ export class AiChatRunRepo {
trx?: KyselyTransaction,
): Promise<number> {
const db = dbOrTx(this.db, trx);
const now = new Date();
let query = db
.updateTable('aiChatRuns')
.set({
status: 'aborted',
finishedAt: now,
updatedAt: now,
// #491: DB-clock stamps (see `update`). The staleness WHERE below stays on
// the app clock — a >minutes window makes the ms-scale skew irrelevant.
finishedAt: sql`now()`,
updatedAt: sql`now()`,
error: sql`coalesce(error, ${'Run interrupted by a server restart.'})`,
})
.where('status', 'in', ACTIVE_RUN_STATUSES as unknown as string[]);
@@ -263,7 +270,7 @@ export class AiChatRunRepo {
// sibling replica's live run is never aborted. Omitted on the phase-1 boot
// sweep -> unconditional.
if (typeof opts.staleMs === 'number') {
const staleBefore = new Date(now.getTime() - opts.staleMs);
const staleBefore = new Date(Date.now() - opts.staleMs);
query = query.where('updatedAt', '<', staleBefore);
}
const rows = await query.returning('id').execute();
@@ -0,0 +1,110 @@
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();
}
}
@@ -13,6 +13,7 @@ import { jsonArrayFrom, jsonObjectFrom } from 'kysely/helpers/postgres';
import { ExpressionBuilder, sql } from 'kysely';
import { DB } from '@docmost/db/types/db';
import { resolveAgentProvenance } from '../agent-provenance';
import { PageHistoryKind } from '../../../collaboration/constants';
/**
* Role-resolution subquery for a page-history row's bound AI chat (#300). Joins
@@ -46,6 +47,9 @@ export class PageHistoryRepo {
'lastUpdatedById',
'lastUpdatedSource',
'lastUpdatedAiChatId',
// #370 — intentionality tier ('manual' | 'agent' | 'idle' | 'boundary');
// null on legacy rows (= autosave). Selected so callers can read/promote it.
'kind',
'contributorIds',
'spaceId',
'workspaceId',
@@ -85,9 +89,15 @@ export class PageHistoryRepo {
async saveHistory(
page: Page,
opts?: { contributorIds?: string[]; trx?: KyselyTransaction },
): Promise<void> {
await this.insertPageHistory(
opts?: {
contributorIds?: string[];
// #370 — intentionality tier for this snapshot. Omitted → null (legacy
// autosave semantics). Callers derive it server-side, never from a client.
kind?: PageHistoryKind;
trx?: KyselyTransaction;
},
): Promise<PageHistory> {
return await this.insertPageHistory(
{
pageId: page.id,
slugId: page.slugId,
@@ -99,6 +109,7 @@ export class PageHistoryRepo {
// Copy the provenance marker off the page row, as for lastUpdatedById.
lastUpdatedSource: page.lastUpdatedSource,
lastUpdatedAiChatId: page.lastUpdatedAiChatId,
kind: opts?.kind ?? null,
contributorIds: opts?.contributorIds,
spaceId: page.spaceId,
workspaceId: page.workspaceId,
@@ -107,6 +118,25 @@ export class PageHistoryRepo {
);
}
/**
* #370 promote an existing snapshot's intentionality tier in place. Used by
* the manual-save "promote-not-dup" path: when the latest history row already
* holds the exact content being versioned, we upgrade its `kind` instead of
* duplicating a heavy content row.
*/
async updateHistoryKind(
pageHistoryId: string,
kind: PageHistoryKind,
trx?: KyselyTransaction,
): Promise<void> {
const db = dbOrTx(this.db, trx);
await db
.updateTable('pageHistory')
.set({ kind })
.where('id', '=', pageHistoryId)
.execute();
}
async findPageHistoryByPageId(pageId: string, pagination: PaginationOptions) {
const query = this.db
.selectFrom('pageHistory')
+18
View File
@@ -280,6 +280,7 @@ export interface PageHistory {
createdAt: Generated<Timestamp>;
icon: string | null;
id: Generated<string>;
kind: string | null;
lastUpdatedAiChatId: string | null;
lastUpdatedById: string | null;
lastUpdatedSource: string | null;
@@ -691,6 +692,22 @@ 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
@@ -728,6 +745,7 @@ export interface DB {
aiChats: AiChats;
aiChatMessages: AiChatMessages;
aiChatRuns: AiChatRuns;
aiChatRunSteps: AiChatRunSteps;
aiChatPageSnapshots: AiChatPageSnapshots;
apiKeys: ApiKeys;
attachments: Attachments;
@@ -4,6 +4,7 @@ import {
AiChats,
AiChatMessages,
AiChatRuns,
AiChatRunSteps,
AiChatPageSnapshots,
Attachments,
Comments,
@@ -64,6 +65,12 @@ 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.
@@ -70,6 +70,23 @@ 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;
@@ -103,6 +103,15 @@ 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'])
@@ -0,0 +1,131 @@
// 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);
});
});
@@ -8,6 +8,7 @@ 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';
@@ -849,7 +850,12 @@ export class ImportAttachmentService {
): Promise<Buffer> {
try {
const drawioContent = await fs.readFile(drawioPath, 'utf-8');
const drawioBase64 = Buffer.from(drawioContent).toString('base64');
// 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);
let imageElement = '';
// If we have a PNG, include it in the SVG
@@ -875,7 +881,7 @@ export class ImportAttachmentService {
width="600"
height="400"
viewBox="0 0 600 400"
content="${drawioBase64}">${imageElement}</svg>`;
content="${drawioEscaped}">${imageElement}</svg>`;
return Buffer.from(svgContent, 'utf-8');
} catch (error) {
@@ -884,6 +890,24 @@ 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;
@@ -304,37 +304,102 @@ export function clientIp(req: ClientIpRequest): string {
return 'unknown';
}
// 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<{
// 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;
sub?: string;
email?: string;
workspaceId?: string;
sessionId?: string;
}> {
return (token: string) => tokenService.verifyJwt(token, JwtType.ACCESS);
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,
});
}
// Minimal shapes for the Bearer revocation/disabled check. Kept structural so
@@ -728,18 +793,24 @@ export async function resolveMcpSessionConfig(
};
}
// --- 2) fallback A: Bearer access-JWT (user-supplied token) ---
// --- 2) fallback A: Bearer JWT (user-supplied ACCESS or agent API_KEY) ---
const bearer = extractBearer(authHeader);
if (bearer) {
let payload: { sub?: string; email?: string };
try {
payload = await deps.verifyAccessJwt(bearer);
} catch (err) {
const message =
err instanceof Error && err.message
? err.message
: 'Invalid or expired token';
throw new UnauthorizedException(message);
// 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;
}
return {
config: { apiUrl, getToken: async () => bearer },
@@ -115,6 +115,7 @@ 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.
@@ -4,13 +4,16 @@ 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 access-JWT verification for the token fallback).
// 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).
@Module({
imports: [DatabaseModule, AuthModule, TokenModule],
imports: [DatabaseModule, AuthModule, TokenModule, ApiKeyModule],
controllers: [McpController],
providers: [McpService],
})
@@ -6,9 +6,10 @@ import {
isCredentialsFailure,
isInitializeRequestBody,
verifyBearerAccess,
verifyMcpBearer,
bindMcpBearerVerifier,
sharedTokenMatches,
clientIp,
bindAccessJwtVerifier,
extractBearer,
decideBasicGate,
mapAuthResultToResponse,
@@ -524,13 +525,28 @@ describe('resolveMcpSessionConfig', () => {
expect(resolved.identity).toBe('bearer:user-9');
});
it('Bearer invalid -> specific 401 from verifyAccessJwt', async () => {
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.
const verifyAccessJwt = jest
.fn()
.mockRejectedValue(new UnauthorizedException('jwt expired'));
await expect(
resolveMcpSessionConfig('Bearer expired', makeDeps({ verifyAccessJwt })),
).rejects.toThrow('jwt expired');
).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');
});
it('no creds + env service account configured -> service-account config', async () => {
@@ -1001,48 +1017,104 @@ describe('clientIp (XFF-fallback precedence, item 5)', () => {
});
});
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
describe('bindMcpBearerVerifier pins the {ACCESS, API_KEY} allowlist (#501)', () => {
it('calls verifyJwtOneOf with exactly [ACCESS, API_KEY]', async () => {
const verifyJwtOneOf = jest
.fn()
.mockResolvedValue({ sub: 'user-1', workspaceId: 'ws-1' });
const verify = bindAccessJwtVerifier({ verifyJwt });
.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']);
});
});
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');
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,
});
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 -> 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' });
});
// 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',
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',
workspaceId: 'ws-1',
sessionId: 'sess-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 });
await verifyMcpBearer('tok', accessDeps({ verifyJwtOneOf }));
expect(verifyJwtOneOf).toHaveBeenCalledTimes(1);
});
});
@@ -1198,6 +1270,7 @@ describe('McpService.onModuleDestroy — CollabSession teardown (#486)', () => {
{} as any,
{} as any,
{} as any,
{} as any,
);
}
+34 -25
View File
@@ -14,16 +14,17 @@ 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 { JwtPayload } from '../../core/auth/dto/jwt-payload';
import { JwtApiKeyPayload } 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,
verifyBearerAccess,
verifyMcpBearer,
isInitializeRequestBody,
sharedTokenMatches,
clientIp,
bindAccessJwtVerifier,
bindMcpBearerVerifier,
decideBasicGate,
mapAuthResultToResponse,
DocmostMcpConfig,
@@ -105,6 +106,10 @@ 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,
) {
@@ -194,37 +199,41 @@ export class McpService implements OnModuleDestroy {
}
}
// 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.
// 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).
private async verifyMcpBearer(
token: string,
): Promise<{ sub?: string; email?: string }> {
// 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.
// 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.
const instanceWorkspace = await this.workspaceRepo.findFirst();
// 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>,
// 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),
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),
});
}
@@ -20,6 +20,10 @@ export interface IStripeSeatsSyncJob {
export interface IPageHistoryJob {
pageId: string;
// #370 — intentionality tier the worker stamps on the snapshot. All jobs on
// this queue are trailing idle-flush autosnapshots, so this is 'idle' (absent
// → treated as 'idle' by the processor).
kind?: 'idle';
}
/**
@@ -0,0 +1,412 @@
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);
});
@@ -0,0 +1,173 @@
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);
});
@@ -30,11 +30,13 @@ import {
* tees the SSE frames into it via `consumeSseStream` while stamping the DB row id
* via `generateMessageId` (both gated on runId + the resumable flag).
*
* Proven here: a finished run's replay is the full frame sequence incl `[DONE]`
* with `start.messageId` == the seeded DB row id; the anchor check (invariant 6);
* an attach opened BEFORE the first frame follows the live stream from frame 0; an
* explicit stop surfaces `{"type":"abort"}` + `[DONE]` + end to the subscriber;
* and the legacy (non-run) path tees nothing.
* Proven here (tail-only #491): a finished run attached at its persisted frontier
* N_final delivers only the TAIL past N (a synthetic `start` carrying the run-fact
* + the terminal `finish`/`[DONE]`) the step content below N lives in the seeded
* DB row, NOT the ring; the anchor check (invariant 6); an attach opened BEFORE the
* first frame follows the live stream from frame 0; an explicit stop surfaces
* `{"type":"abort"}` + `[DONE]` + end to the subscriber; and the legacy (non-run)
* path tees nothing.
*/
const sleep = (ms: number) => new Promise((r) => setTimeout(r, ms));
@@ -133,14 +135,16 @@ function liveSink(): {
};
}
// The SSE `start` frame carries the message id; pull it out of a `data: {...}`.
function parseStartMessageId(frames: string[]): string | undefined {
// Parse the first `start` frame's JSON out of a `data: {...}` sequence.
function parseStartFrame(
frames: string[],
): { messageId?: string; messageMetadata?: any } | undefined {
for (const f of frames) {
const m = /^data: (\{.*\})\s*$/m.exec(f.trim());
if (!m) continue;
try {
const json = JSON.parse(m[1]);
if (json.type === 'start') return json.messageId;
if (json.type === 'start') return json;
} catch {
/* not this frame */
}
@@ -270,7 +274,7 @@ describe('AiChatService run-stream attach [integration]', () => {
await destroyTestDb();
});
it('run-wrapped: replay is the full frame sequence incl [DONE], start.messageId == the seeded DB row id', async () => {
it('run-wrapped, tail-only: a finished run at N_final delivers the run-fact start + finish/[DONE]; the step content lives in the seeded row', async () => {
const chatId = (await createChat(db, { workspaceId, creatorId: userId })).id;
const registry = new AiChatStreamRegistryService();
const runService = new AiChatRunService(runRepo, {
@@ -300,27 +304,37 @@ describe('AiChatService run-stream attach [integration]', () => {
);
});
const rowId = await assistantRowId(chatId);
// The client reads its persisted step frontier N from the seeded row.
const row: any = await msgRepo.findById(rowId, workspaceId);
const nFinal = row.metadata.stepsPersisted as number;
expect(nFinal).toBe(1); // a single finished step
// The step content is in the SEEDED row (parts/content), not the ring.
expect(JSON.stringify(row.metadata.parts)).toContain('Hello');
// Finished-run replay with expect=live + the correct anchor.
// Attach at N_final with the correct anchor: the tail past step 1 is just
// the terminal frames; step 0's 'Hello' is BELOW the frontier (seeded).
const sink = liveSink();
const att = await registry.attach(chatId, true, rowId, sink.cb);
const att = await registry.attach(chatId, rowId, nFinal, sink.cb);
expect(att).not.toBeNull();
expect(att!.finished).toBe(true);
// The tee captured frames (consumeSseStream was wired).
expect(att!.replay.length).toBeGreaterThan(0);
// generateMessageId stamped the DB row id onto the streamed start frame.
expect(parseStartMessageId(att!.replay)).toBe(rowId);
// The full sequence includes the streamed text and the terminal marker.
const joined = att!.replay.join('');
expect(joined).toContain('Hello');
// The synthetic start frame carries the run-fact (runId/chatId), the source
// of the run-fact on re-attach.
const start = parseStartFrame(att!.replay);
expect(start?.messageMetadata).toMatchObject({
runId: box.runId,
chatId,
});
// The terminal marker is delivered so the client's SDK closes the stream.
expect(att!.replay.some((f) => f.includes('[DONE]'))).toBe(true);
// 'Hello' (step 0, below the frontier) is NOT re-streamed — it is seeded.
expect(att!.replay.some((f) => f.includes('Hello'))).toBe(false);
} finally {
registry.onModuleDestroy();
await cleanup();
}
});
it('anchor mismatch with expect=live returns null (invariant 6)', async () => {
it('anchor mismatch returns null (invariant 6)', async () => {
const chatId = (await createChat(db, { workspaceId, creatorId: userId })).id;
const registry = new AiChatStreamRegistryService();
const runService = new AiChatRunService(runRepo, {
@@ -347,7 +361,7 @@ describe('AiChatService run-stream attach [integration]', () => {
const sink = liveSink();
// A foreign anchor must NOT replay this run's transcript.
expect(
await registry.attach(chatId, true, 'a-different-run-row', sink.cb),
await registry.attach(chatId, 'a-different-run-row', 1, sink.cb),
).toBeNull();
} finally {
registry.onModuleDestroy();
@@ -376,8 +390,11 @@ describe('AiChatService run-stream attach [integration]', () => {
try {
// Attach while the entry exists (opened at begin) but before any frame.
const sink = liveSink();
const att = (await registry.attach(chatId, false, undefined, sink.cb))!;
expect(att.replay).toEqual([]); // nothing streamed yet -> replay from 0
const att = (await registry.attach(chatId, undefined, 0, sink.cb))!;
// Nothing streamed yet -> the tail is just the synthetic start frame; the
// whole live stream (start..DONE) follows via onFrame after start().
expect(att.replay).toHaveLength(1);
expect(att.replay[0]).toContain('"type":"start"');
att.start(); // go live (drains nothing, then follows)
// Now emit the whole turn.
@@ -448,7 +465,7 @@ describe('AiChatService run-stream attach [integration]', () => {
});
try {
const sink = liveSink();
const att = (await registry.attach(chatId, false, undefined, sink.cb))!;
const att = (await registry.attach(chatId, undefined, 0, sink.cb))!;
att.start();
// Give streamText a beat to begin consuming the partial output.
@@ -526,7 +543,9 @@ describe('AiChatService run-stream attach [integration]', () => {
expect(entry).toBeDefined();
expect(entry.finished).toBe(true);
const sink = liveSink();
expect(await registry.attach(chatId, false, undefined, sink.cb)).toBeNull();
// Finished with an EMPTY ring (aborted before any frame) -> null -> the
// client degrades to poll instead of hanging on an empty stream.
expect(await registry.attach(chatId, undefined, 0, sink.cb)).toBeNull();
} finally {
registry.onModuleDestroy();
await cleanup();
@@ -556,8 +575,8 @@ describe('AiChatService run-stream attach [integration]', () => {
});
const sink = liveSink();
// No entry was ever opened; attach always yields null.
expect(await registry.attach(chatId, false, undefined, sink.cb)).toBeNull();
expect(await registry.attach(chatId, true, 'anything', sink.cb)).toBeNull();
expect(await registry.attach(chatId, undefined, 0, sink.cb)).toBeNull();
expect(await registry.attach(chatId, 'anything', 1, sink.cb)).toBeNull();
} finally {
registry.onModuleDestroy();
await cleanup();
@@ -0,0 +1,169 @@
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);
});
@@ -0,0 +1,163 @@
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,
);
});
});
@@ -0,0 +1,63 @@
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);
});
});
@@ -322,20 +322,21 @@ describe('AiChatService.stream [integration]', () => {
});
/**
* #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.
* #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.
*/
describe('deferred tool loading ON — per-turn activation, no leak (#332)', () => {
describe('deferred tool loading ON — cross-turn activation persistence (#332 + #490)', () => {
// 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.
@@ -451,7 +452,7 @@ describe('AiChatService.stream [integration]', () => {
} as any);
}
it('activates a deferred tool for the SAME turn, and a NEW turn starts cold (no leak)', async () => {
it('activates a deferred tool for the SAME turn, and a NEW turn SEEDS it from persisted chat metadata (#490)', async () => {
const chatId = (await createChat(db, { workspaceId, creatorId: userId })).id;
// --- Turn 1: loadTools(createPage) on step 1, then answer on step 2. ---
@@ -474,7 +475,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: must start cold again. ---
// --- Turn 2 on the SAME chat: seeds the persisted activation (#490). ---
const model2 = new MockLanguageModelV3({
doStream: async () => ({ stream: successStream() }),
} as any);
@@ -485,9 +486,10 @@ describe('AiChatService.stream [integration]', () => {
const nextTurnFirstStep = toolNames(model2.doStreamCalls[0]);
expect(nextTurnFirstStep).toContain('loadTools');
// 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');
// #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');
});
});
});
@@ -0,0 +1,162 @@
import { randomUUID } from 'node:crypto';
import { Queue, Worker } from 'bullmq';
import { PersistenceExtension } from '../../src/collaboration/extensions/persistence.extension';
/**
* #370 integration property of the idle-snapshot pipeline against REAL BullMQ.
*
* This is deliberately NOT a unit test of computeHistoryJob (that lives in
* compute-history-job.spec.ts). The point here is the OBSERVABLE end-to-end
* behaviour of the production `enqueuePageHistory` remove-then-add debounce
* driving a real Redis-backed delayed queue + worker (the #431#439 class: a
* locally-correct function whose queue/timer property was never exercised):
*
* - a CONTINUOUS burst of stores lasting several caps yields periodic idle
* snapshots at least one per max-wait cap, NOT one-per-store; and
* - an INTERMITTENT burst (a few stores, then quiet) yields exactly ONE
* trailing snapshot.
*
* We shrink the idle windows to milliseconds (jest.mock of collaboration
* constants) so real BullMQ delayed jobs actually promote within the test
* fake timers cannot advance Redis's own delayed-set clock, so the intervals
* must be real but tiny. The production method under test is called verbatim.
*/
// NOTE: jest.mock is hoisted above the module's const initializers, so its
// factory cannot close over MAX_WAIT_MS/INTERVAL_MS — the literals are inlined
// here and MUST stay in sync with the consts below (a single source of truth is
// impossible across the hoist boundary).
jest.mock('../../src/collaboration/constants', () => {
const actual = jest.requireActual('../../src/collaboration/constants');
return {
...actual,
IDLE_MAX_WAIT_USER: 300,
IDLE_MAX_WAIT_AGENT: 300,
IDLE_INTERVAL_USER: 1000,
IDLE_INTERVAL_AGENT: 1000,
};
});
// Mirrors the mocked IDLE_MAX_WAIT_* above (IDLE_INTERVAL_* is 1000 > this, so
// the max-wait ceiling is what actually governs the trailing delay).
const MAX_WAIT_MS = 300;
const REDIS_CONNECTION = {
host: process.env.TEST_REDIS_HOST ?? '127.0.0.1',
port: Number(process.env.TEST_REDIS_PORT ?? 6379),
};
const sleep = (ms: number) => new Promise((r) => setTimeout(r, ms));
describe('#370 idle-snapshot pipeline (real BullMQ)', () => {
let queue: Queue;
let worker: Worker;
let extension: PersistenceExtension;
// Every processed snapshot, tagged by pageId so the two scenarios stay isolated.
const processed: Array<{ pageId: string; kind: string; at: number }> = [];
const queueName = `history-idle-int-${randomUUID()}`;
beforeAll(async () => {
queue = new Queue(queueName, {
connection: REDIS_CONNECTION,
// Mirror the production default (BullModule.forRoot removeOnComplete): the
// enqueue idiom relies on the jobId being freed once a job completes so the
// next burst can re-arm the same id.
defaultJobOptions: { removeOnComplete: true, removeOnFail: true },
});
await queue.waitUntilReady();
worker = new Worker(
queueName,
async (job) => {
processed.push({
pageId: job.data?.pageId,
kind: job.data?.kind,
at: Date.now(),
});
},
{ connection: REDIS_CONNECTION },
);
await worker.waitUntilReady();
// Construct the real extension; only historyQueue (5th ctor arg) and the
// internal idleBurstStart map are exercised by enqueuePageHistory, so the
// other collaborators can be null — the constructor only assigns fields.
extension = new PersistenceExtension(
null as any, // pageRepo
null as any, // pageHistoryRepo
null as any, // db
null as any, // aiQueue
queue as any, // historyQueue
null as any, // notificationQueue
null as any, // collabHistory
null as any, // transclusionService
);
});
afterAll(async () => {
// Force-close and fully drain so no BullMQ background activity (delayed-set
// polling, blocking BRPOPLPUSH) bleeds into later suites in this single
// shared jest worker (maxWorkers: 1).
await worker?.close(true).catch(() => undefined);
await queue?.obliterate({ force: true }).catch(() => undefined);
await queue?.close();
// Let the redis sockets settle before the next suite starts.
await sleep(150);
});
const arm = (pageId: string) =>
(extension as any).enqueuePageHistory({ id: pageId }, 'user');
it('continuous burst over several caps → periodic idle snapshots (≥1 per cap, not one-per-store)', async () => {
const pageId = randomUUID();
const runMs = 6 * MAX_WAIT_MS; // ~6 caps of unbroken editing
// Store cadence that does NOT evenly divide the cap: real hocuspocus stores
// are not aligned to cap boundaries, so a boundary job promotes in the gap
// before the next store's remove(). A cap-aligned cadence would instead land
// a store exactly on every boundary and lose the snapshot to the documented
// remove-vs-active race — an artefact of the test clock, not the pipeline.
const stepMs = 70;
const stores = Math.floor(runMs / stepMs);
const start = Date.now();
let count = 0;
while (Date.now() - start < runMs) {
await arm(pageId);
count++;
await sleep(stepMs);
}
// Let the final armed job flush.
await sleep(2 * MAX_WAIT_MS);
const snaps = processed.filter((p) => p.pageId === pageId);
// Every autosnapshot is an idle-kind row.
expect(snaps.every((s) => s.kind === 'idle')).toBe(true);
// Periodic: at least one per cap over a multi-cap burst (lower-bounded loosely
// to stay robust; the property is "fires at least every cap", not a single
// trailing snapshot).
expect(snaps.length).toBeGreaterThanOrEqual(3);
// But NOT one-per-store: ~`stores` stores were issued; the debounce must
// collapse them to a small multiple of the cap count, nowhere near per-store.
expect(snaps.length).toBeLessThanOrEqual(Math.ceil(stores / 2));
});
it('intermittent burst (a few stores, then quiet) → exactly ONE trailing snapshot', async () => {
const pageId = randomUUID();
// A short burst well within a single cap window, then silence.
await arm(pageId);
await sleep(40);
await arm(pageId);
await sleep(40);
await arm(pageId);
// Wait comfortably past the cap so the single pending trailing job fires.
await sleep(4 * MAX_WAIT_MS);
const snaps = processed.filter((p) => p.pageId === pageId);
expect(snaps).toHaveLength(1);
expect(snaps[0].kind).toBe('idle');
});
});
+1 -1
View File
@@ -387,7 +387,7 @@ bucketByDay(sessions, tz):
факты ниже — ground truth, можно дозапросить файлы через gitea MCP по указанному SHA):
- `page_history.kind``varchar(20)`, NULLABLE, БЕЗ дефолта (migration
`20260705T120000-page-history-kind.ts`). Домен: `manual`/`agent`/`idle`/`boundary`;
`20260707T120000-page-history-kind.ts`). Домен: `manual`/`agent`/`idle`/`boundary`;
legacy `null` = автосейв (`collaboration/constants.ts`, `PageHistoryKind`).
- `kind` УЖЕ включён в `PageHistoryRepo.baseFields` (`page-history.repo.ts`) — читается всеми
выборками истории. `saveHistory({kind})` и `updateHistoryKind(id, kind)` существуют.
+36 -9
View File
@@ -178,10 +178,11 @@ function sliceModel(xml: string): string | null {
// --- decode chain ----------------------------------------------------------
/**
* 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.
* 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.
*/
export function extractContentAttr(svg: string): string {
const { doc, error } = parseXml(svg);
@@ -195,12 +196,23 @@ 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.
// 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.
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");
@@ -307,9 +319,16 @@ 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="${base64(drawioFile)}">${inner}</svg>
* <svg xmlns= xmlns:xlink= content="${xmlEscape(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,
@@ -318,14 +337,14 @@ export function buildDrawioSvg(
title = "Page-1",
): string {
const file = encodeDrawioFile(modelXml, title);
const base64 = Buffer.from(file, "utf-8").toString("base64");
const content = xmlEscape(file);
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="${base64}">${inner}</svg>`
`content="${content}">${inner}</svg>`
);
}
@@ -334,7 +353,15 @@ function xmlEscape(s: string): string {
.replace(/&/g, "&amp;")
.replace(/</g, "&lt;")
.replace(/>/g, "&gt;")
.replace(/"/g, "&quot;");
.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;");
}
// --- normalization + hash --------------------------------------------------
+123 -2
View File
@@ -51,6 +51,14 @@ 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", () => {
@@ -335,16 +343,20 @@ 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
// base64-encoded inside content=, and escaped inside the file XML).
// entity-escaped inside content=, doubly so where the file XML already escaped
// them inside name="...").
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 = Buffer.from(contentMatch[1], "base64").toString("utf-8");
const fileXml = unescapeAttr(contentMatch[1]);
const nameMatch = /<diagram id="[^"]*" name="([^"]*)">/.exec(fileXml);
assert.ok(nameMatch, "the diagram name attribute is intact and quote-safe");
const decodedTitle = nameMatch[1]
@@ -358,3 +370,112 @@ 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);
});