Compare commits

..

21 Commits

Author SHA1 Message Date
agent_coder 9f0e880e0b perf(mcp): checkNewComments — параллелизм с капом (#490)
checkNewComments делал O(N) последовательных REST-вызовов listComments по страницам
working set — большой space линеен по round-trip'ам. Теперь per-page фетчи идут с
ограниченным параллелизмом (cap 6, середина полосы 5–8): независимые чтения не ждут
друг друга, но и не заваливают сервер/сокеты.

mapWithConcurrency — крошечный пул без зависимости от p-limit: N воркеров тянут
следующий индекс с общего курсора. Порядок результатов сохраняется (по входному
порядку страниц), поэтому вывод детерминирован независимо от того, какой фетч
завершился первым. Серверный batch-эндпоинт «comments updated since T по space» —
опционально, отдельным заходом.

Тест (mock-HTTP): 13 страниц, задержанный /api/comments — maxInFlight > 1 и <= 6
(последовательная реализация дала бы 1), порядок результатов = порядок обхода.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-11 11:35:58 +03:00
agent_coder d35956b9e9 perf(ai-chat): snapshotOpenPage fast-path (#490)
snapshotOpenPage делал полный экспорт Markdown + upsert каждый ход. Fast-path: если
снапшот уже существует на ТЕКУЩЕЙ версии страницы (тот же instant updated_at), его
контент уже актуален — пропускаем экспорт+upsert целиком. Ход, не тронувший
открытую страницу (частый случай), больше не делает работы по снапшоту.

Зеркалит read-side fast-path в detectPageChange (sameInstant): оба доверяют, что
правка страницы двигает updated_at. Когда агент/человек ПРАВИЛ страницу этим ходом,
updated_at продвинулся → не совпадает → экспортируем как раньше (правки агента
запекаются в снапшот, инвариант #274 сохранён).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-11 11:30:46 +03:00
agent_coder c0e39a395c perf(ai-chat): deferred-активация тулов в metadata чата (#490)
Активированный set сбрасывался каждый ход → модель заново гоняла loadTools, чтобы
переактивировать те же тулы (лишний round-trip на каждом ходу). Теперь набор
персистится в metadata чата и сидируется на следующем ходу.

- Миграция: jsonb-колонка metadata на ai_chats (default '{}'); db.d.ts дополнен
  вручную (AiChats.metadata: Generated<Json>).
- seedActivatedTools(metadata, validDeferredNames): читает сохранённый набор,
  ПЕРЕСЕКАЯ с актуальными validDeferredNames — смена allowlist/ролей не воскресит
  несуществующий тул (иначе prepareAgentStep получил бы фантомное активное имя).
  Сид только при deferredEnabled.
- Персист на завершении хода (once-guard, во всех терминальных ветках рядом со
  snapshotTurnEnd): детерминированно отсортированный набор, merge в существующий
  bag (другие ключи сохраняются), запись пропускается если ничего нового не
  активировано (обычный ход не даёт лишней записи).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-11 11:27:15 +03:00
agent_coder c144abcb83 feat(ai-chat): токен-бюджет реплея истории + реактивная ветка (#490)
Вся персистентная история реплеится провайдеру КАЖДЫЙ ход, поэтому длинный чат
рано или поздно упирается в контекстное окно и получает провайдерский 400 на
каждом ходу — навсегда (чат «кирпичится»). Бюджетер ограничивает РЕПЛЕЙ (никогда
не мутирует персист — в БД остаётся полная запись), детерминированно и byte-stable
(обрезанный префикс идентичен от хода к ходу → дружелюбно к prompt-cache).

Единый оценщик chars/2.5 (кириллица; chars/4 занижает вдвое) вынесен в shared-пакет
packages/token-estimate; клиентский count-stream-tokens.ts переведён на него ТЕМ ЖЕ
коммитом (два расходящихся оценщика = «бейдж 60%, а бюджет уже режет»).

history-budget.ts (чистый, покрыт тестами):
- resolveReplayBudget(raw): min(100k, 0.7×window) при заданном окне; флэт 100k при
  незаданном (именно эти инсталляции ловят терминальный overflow — warn-лог); 0 =
  явный off-switch. Читается СЫРОЙ chatContextWindow, т.к. parsePositiveInt схлопывает
  0 и unset в undefined (новое поле ResolvedAiConfig.chatContextWindowRaw).
- trimHistoryForReplay: первичный сигнал — провайдерский факт metadata.contextTokens
  прошлого хода; chars-оценка — дельта/раскройка/фолбэк. Порядок: обрезка tool-outputs
  старых ходов (head+tail+маркер) → механическое схлопывание старейших ходов
  (конкатенация, НЕ LLM) → текущий + последние N ходов всегда полные. Пейринг
  tool-call/result сохраняется (схлопывание убирает ОБЕ части).
- isContextOverflowError: классификация провайдерского 400 (статус + паттерны).

Реактивная ветка: превентивная оценка не даёт инварианта (первый переполняющий ход
не имеет usage). onError классифицирует context-overflow → пишет различимую причину
и штампует metadata.replayOverflow; следующий ход бюджетер режет агрессивно
(0.5×), что и раскирпичивает чат. Наблюдаемость: metadata.replayTrimmedToTokens.

ПРИМЕЧАНИЕ по реактивной ветке (форк, требует решения ревьюера): истинный in-turn
re-pipe (перезапуск streamText в тот же ответ) архитектурно несовместим с текущим
пайпом — pipeUIMessageStreamToResponse пишет writeHead СИНХРОННО (подтверждено в
ai@6.0.207), а suite ожидает await stream() c моком, не дёргающим колбэки, — так что
отложенный пайп/ожидание сигнала повесит тесты. Поэтому реализована реактивная
рекавери «классификация → штамп → агрессивный ре-трим на следующем ходу», что даёт
тот же инвариант (чат не кирпичится) без рискованного рефактора стрима.

Тесты (наблюдаемые свойства): объём записи через дельту pg_current_wal_lsn() на живой
gitmost-test-pg вокруг 50-шагового прогона (несжимаемые payload'ы) — trace-колонка
v1=140МБ → v2=0.04МБ (в 3206× меньше), полная строка 289МБ → 140МБ (−51%); dual-shape
не нужен здесь; «окно не задано → бюджет применяется»; реактивная классификация на
реальном 400-шейпе; parity клиент/сервер оценщика.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-11 11:21:27 +03:00
agent_coder 71517552b7 perf(ai-chat): кэш compactToolOutput по identity шага (#490)
compactToolOutput делает JSON.stringify каждого output на КАЖДОМ flush. Т.к.
onStepFinish на шаге N перестраивает всю assistant-строку по всем N накопленным
шагам, а каждый output — 50–200 KB, это O(N²) stringify за ход.

Мемоизация по identity шага: finished-шаг в capturedSteps неизменен и держит
стабильную ссылку между flush'ами, поэтому его parts (и дорогой stringify output)
строятся ровно раз за ход. buildStepParts вынесен в чистую функцию; assistantParts
принимает опциональный StepPartsCache (WeakMap<step, parts>), flushAssistant
пробрасывает его, stream() заводит один WeakMap на ход и передаёт во все flush'и.
Промах кэша (или его отсутствие в тестах/легаси-вызовах) просто пересобирает —
байтового расхождения нет.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-11 10:43:35 +03:00
agent_coder 3f217f4835 perf(ai-chat): формат трейса tool_calls v2 — outputs только в parts (#490)
Каждый tool-output хранился ДВАЖДЫ: в metadata.parts (assistantParts) И в
tool_calls (serializeSteps). При 50-шаговом ране с outputs по 50–200 KB это
127–510 МБ записи в Postgres за ход (+WAL/TOAST/dead tuples), т.к. onStepFinish
переписывает всю строку. Копия в parts — та, что реально реплеится модели и
рендерится UI/markdown-экспортом, так что копия в трейсе была чистым дублем.

Новый формат элементов tool_calls (v2), парно на каждый вызов:
  {toolName, input}                      — вызов
  {toolName, ok: true}                   — успех (БЕЗ output)
  {toolName, error, kind: 'thrown'}      — брошенный tool-error
  {toolName, error, kind: 'interrupted'} — прерван mid-step (abort/restart)

kind обязателен: синтетический «Tool call did not complete.» при прерывании иначе
неотличим от реального hard-fail и загрязняет error-rate. Различие структурное
(errorsById-хит против синтетической ветки), НЕ per-tool классификатор — soft-
маркеры в трейс не выносятся (остаются в metadata.parts).

metadata.toolTraceVersion: 2 — маркер эры; старые строки НЕ мигрируются
(перезапись гигантских jsonb — тот самый WAL-чарн). serializeSteps пейрит
результаты/ошибки по toolCallId (как assistantParts); общая константа
TOOL_CALL_INCOMPLETE_TEXT держит текст реплея и трейса в синхроне.

docs/reading-ai-logs.md переписан dual-shape: ветвление по toolTraceVersion,
soft-анализ v2 через metadata.parts, правило «не сравнивать агрегаты через границу
эр». UI action-log и markdown-экспорт читают только parts — не затронуты.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-11 10:39:42 +03:00
agent_coder 791a707eb6 docs(ai-chat): задокументировать поверхность #489 — env-var, устаревшая заметка, CHANGELOG
F1 (ревью #508):
- .env.example: новая AI_MCP_SSE_BODY_TIMEOUT_MS (дефолт 600000, 10 мин) —
  bodyTimeout SSE-транспорта external-MCP; idle между тул-вызовами легитимен.
- .env.example: поправлена ставшая ложной заметка про AI_MCP_STREAM_TIMEOUT_MS
  (SSE-idle-между-вызовами больше не режется им — с #489 это отдельная var;
  1-мин silence остаётся только для HTTP/headers и одиночного залипшего вызова).
- CHANGELOG [Unreleased]/Fixed: битый part больше не 500-ит каждый ход + не
  плодит дубль user-строки; MCP-транспорт-дропы восстанавливаются in-run
  (readOnly ретраится раз, write никогда) + новая env-var.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-11 09:11:32 +03:00
agent_coder 60205398bb test(ai-chat): застабить findAllByChat в lifecycle-моке под convert-before-insert (#489)
Коммит 1 (#489) прогоняет загрузку истории + convertToModelMessages ДО insert
user-строки (convert-before-insert, чтобы ретрай не плодил дубли), поэтому
stream() теперь зовёт реальный метод репо findAllByChat перед insert. Хэнд-роллед
мок в тесте «exception after beginRun → settled to error» стабил только insert,
и тест ловил «findAllByChat is not a function» вместо «insert boom».

Порядок инварианта НЕ изменился: и findAllByChat, и insert идут ПОСЛЕ beginRun
(reconcileChat между ними — best-effort, глотает ошибку), так что бросок по-
прежнему происходит ПОСЛЕ начала рана и обязан сеттлить ран в error. Застабил
findAllByChat → [] (реальный репо-метод, см. ai-chat.controller.ts), тест снова
доходит до insert boom и проверяет settle-to-error.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-11 08:25:29 +03:00
agent_coder 38a09c5ca1 fix(ai-chat): write-class-ретрай только для доверенного внутреннего Docmost-сервера (#489 ревью)
MEDIUM (реальная брешь): SHARED_TOOL_WRITE_CLASS — имена ТУЛОВ Docmost, но
mergeNamespaced применял мапу к тулам ЛЮБОГО стороннего MCP-сервера по совпадению
rawName. Сторонний WRITE-тул, названный как Docmost-read (getPage/listPages/…),
наследовал readOnly → авто-ретрай при транспортной ошибке → double-apply (класс
#435). Гарантия «unknown third-party → write → не ретраится» была ЛОЖНОЙ при
коллизии имён.

Фикс: write-class мапа применяется ТОЛЬКО к серверу, про который известно, что
это внутренний Docmost-MCP (isInternalDocmostServer). Сейчас это ВСЕГДА false —
в этом пути нет встроенного/доверенного Docmost-сервера: все ai_mcp_servers суть
сторонние admin-конфиги, а собственные тулы Docmost идут отдельным in-app путём
(docmostTools), не через mcp-clients. Значит НИ ОДИН сторонний тул не получает
readOnly по коллизии и не авто-ретраится (undefined → трактуется как write).
Мапа грузится лениво только если есть доверенный сервер (иначе ESM-импорт
пропускается). isInternalDocmostServer — метод-сим, флипается при появлении
доверенного сервера (kind/isBuiltin-колонка или сконфигурённый self-MCP URL).

LOW: reconnect не наблюдает composed abort во время 5-сек handshake — задокумен-
тировано (окно поздней отмены ≤5s, сокет закрывается на turn-end); проброс
composed в общий CAS-дедуплицированный reconnect намеренно НЕ сделан (отменил бы
реконнект, нужный конкурентному живому вызову).

Тест: сторонний WRITE-тул с именем getPage при транспортной ошибке НЕ ретраится;
mutation-verify — форсирование trusted делает тест красным (чужой getPage ретраится).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-11 08:01:33 +03:00
agent_coder 1b05224b27 fix(ai-chat): MCP-кэш — in-run восстановление, ретраи только для readOnly (#489)
Кэш внешних MCP-клиентов не делал health-check/reconnect-on-error: после
разрыва SSE-транспорта (bodyTimeout режет при тишине МЕЖДУ вызовами) отдавал
труп до конца TTL, а модель жгла шаги на ретраях ВНУТРИ текущего рана.

- Новое поле writeClass ('readOnly'|'write') на КАЖДОМ SHARED_TOOL_SPEC +
  registration-time assert (и compile-time через satisfies). Все 48 спеков
  расклассифицированы: чтения → readOnly, любая мутация страницы/коммента/шэра/
  диаграммы → write. Экспортированы SHARED_TOOL_WRITE_CLASS + isRetryableWriteClass.
- Per-run обёртка восстановления транспорта: при транспортной ошибке readOnly-тул
  реконнектит свой сервер и ретраит РОВНО 1 раз ВНУТРИ рана; write-тул НЕ
  авторетраится (indeterminate — «могло примениться, проверь», класс инцидента
  #435). CAS-своп байндинга по identity (проигравший конкурентный вызов ретраит
  на текущем клиенте, не минтит второй). Лизы не освобождаются mid-run — ран
  копит set (старая+новая) и релизит на turn-end.
- Проверка abortSignal ПЕРЕД ретраем И ПЕРЕД чеканкой свежего клиента; per-call
  cap покрывает оба attempt'а + connect.
- Классификация транспортной ошибки по РЕАЛЬНЫМ шейпам undici (SocketError/
  BodyTimeoutError, cause-цепочка), не по мок-ошибкам.
- Отдельный, поднятый bodyTimeout для SSE-транспорта MCP (тишина между вызовами
  легальна) — DEFAULT 10 мин, AI_MCP_SSE_BODY_TIMEOUT_MS.

write-class map грузится в mcp-clients лениво через dynamic import (пакет ESM),
type-only импорт — без static require ESM из commonjs.

Тесты на РЕАЛЬНЫХ error-шейпах: «повтор после обрыва ВНУТРИ рана получает живой
клиент», «write-тул не авторетраится», «ретрай после Stop не происходит»,
+ writeClass-контракт в mcp node --test.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-11 08:01:33 +03:00
agent_coder b50b32bf64 fix(ai-chat): валидация клиентских parts + insert после конвертации (#489)
Юзер-сообщение персистилось в metadata БЕЗ валидации и реплеилось через
convertToModelMessages на каждом ходе; insert user-строки шёл ДО конвертации.
Битая структура parts (например, null-элемент в массиве) → конвертация кидает
→ 500 на КАЖДОМ ходе навсегда, а каждый ретрай добавлял дубль user-строки.

- Санитизация parts при приёме: whitelist { text }, прочее (в т.ч. tool-part
  в input-available) отбрасывается с warn — не попадает в metadata.
- convertToModelMessages прогоняется ДО insert'а user-строки (ретрай не плодит
  дубли); при падении на СТАРОЙ истории — per-row конвертация изолирует битую
  строку и деградирует её до plain-text с маркером «[tool context omitted]»
  (молчаливая потеря tool-контекста недопустима).

Тесты против РЕАЛЬНОГО convertToModelMessages (null-part реально кидает):
unit трёх веток + сервис-регресс «чат с битым сообщением в истории работает,
маркер доходит до модели, одна user-строка».

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-11 08:01:33 +03:00
agent_coder 1d704d4ec5 test(mcp): покрыть write-safe-point «после Stop новая запись не стартует» (#487, F4)
paginate-abort-safepoint покрывал только READ-safe-point; WRITE-шов
(collaboration.mutatePageContent / context.mutateLiveContentUnlocked —
signal?.throwIfAborted() перед session.mutate) был без теста. Добавляю
мок-collab юнит через тест-сим __setCollabProviderFactory (фейковый провайдер
с мгновенным onSynced): предварительно abort-нутый сигнал → оба шва реджектят
ДО session.mutate, трансформ не вызывается; контрольные кейсы (живой сигнал)
подтверждают, что session.mutate достижим.

MUTATION-VERIFY: снятие throwIfAborted в обоих швах делает красными ровно два
abort-кейса, контрольные остаются зелёными.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-11 07:49:30 +03:00
agent_coder 2cb07ef7fb test(ai-chat): покрыть оркестраторы reconcile()/reconcileChat() (#487, F3)
Клаузы (a-d) тестировались по отдельности, но сам оркестратор reconcile()
(ПОРЯДОК клауз + пер-клаузная try/catch-изоляция «одна упавшая не блокирует
остальные») не проверялся, а reconcileChat() (старт каждого хода) не был
покрыт вовсе. Добавляю: тест порядка a→b→c→d; тест изоляции (клауза b кидает
— c и d всё равно отрабатывают, reconcile не реджектит); тест reconcileChat
(скоуп по чату, bound 50, failed→error / прочее→aborted).

MUTATION-VERIFY: снятие try/catch у клаузы (b) делает тест изоляции красным
(исключение пробрасывается, c+d пропускаются).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-11 07:49:30 +03:00
agent_coder 9ba97b6d2b docs(ai-chat): задокументировать наблюдаемую поверхность #487 (F2)
- CHANGELOG [Unreleased]: серверный supersede + три кода (SUPERSEDE_INVALID /
  SUPERSEDE_TARGET_MISMATCH / SUPERSEDE_TIMEOUT) в Added; смена поведения
  (легаси вторая вкладка теперь → 409 A_RUN_ALREADY_ACTIVE вместо второго
  параллельного стрима; каждый ход — ран в обоих режимах) в Changed.
- .env.example: три новых AI_CHAT_* (SUPERSEDE_TIMEOUT_MS / RECONCILE_INTERVAL_MS
  / INAPP_TOOL_CALL_CAP_MS, дефолты 10000/120000/120000, связь cap↔staleness).
- AGENTS.md:458: ран теперь универсален (оба режима), флаг autonomousRuns =
  только disconnect-семантика; плюс упоминание supersede.
- ai-chat.service.ts:1136: устаревший коммент про «legacy → socket signal» →
  оба режима на run-signal, guard runId защищает лишь no-handle fallback.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-11 07:49:14 +03:00
agent_coder 0b9de2c25e fix(ai-chat): owner-gate stream() до supersede — закрыть утечку чужого runId (#487, F1)
stream() был единственным эндпоинтом ai-chat без assertOwnedChat: участник
того же воркспейса (НЕ владелец чата) мог послать POST /stream
{chatId:<чужой>, supersede:{runId:<любой>}} и (а) выудить активный runId
жертвы из ответа 409 SUPERSEDE_TARGET_MISMATCH, затем (б) requestStop чужого
рана. Добавляю owner-check в начале stream() (когда есть body.chatId), ровно
как в /stop и соседях — pre-hijack, чистый 403. Это заодно закрывает и
негейченную кросс-юзерную запись через тот же stream() (база #500).

Тест: не-владелец POST /stream с чужим chatId → ForbiddenException,
runId не утёк, supersede/requestStop не вызваны. MUTATION-VERIFY: снятие
assertOwnedChat делает тест красным (возвращается путь утечки).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-11 07:49:02 +03:00
agent_coder 123cba7de1 fix(mcp): убрать залипший маркер конфликта в context.ts (артефакт ребейза)
После ребейза в src/client/context.ts:209 остался одиночный маркер
`>>>>>>> 917c4064` без парного открытия — он попал в коммит и ломал сборку
всего пакета mcp (TS1185 Merge conflict marker), из-за чего не запускались
никакие node --test. Код выше маркера цел; удаляю только строку маркера.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-11 07:48:52 +03:00
agent_coder 001eb1c923 docs(mcp): точный коммент toolAbortSignal — set-and-leave, не restore (#487, ревью nit)
Внутреннее ревью: docstring поля/метода toolAbortSignal утверждал «restores the
prior value on unwind», но wrapInAppToolWithCap намеренно НЕ восстанавливает
сигнал (set-and-leave) — именно это заставляет брошенного проигравшего race
бросить на следующем safe-point; корректность держится на свежем клиенте на ход +
перезаписи следующим вызовом. Комментарии приведены в соответствие с механизмом.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-11 07:26:01 +03:00
agent_coder 002c931c6b feat(ai-chat): ретраи finalizeAssistant + owner-write приоритет + двусторонний reconcile (#487)
Раньше finalizeAssistant ставил finalized=true ДО записи и не ретраил → один
неудачный UPDATE = строка вечно 'streaming'; свип был boot-only; run-свип
безусловный — асимметрия «run succeeded / message streaming навсегда».

- finalizeAssistant: bounded-ретраи; once-гейт закрывается ТОЛЬКО после успешной
  записи; возвращает ok. Правило owner-write: терминальная запись owner'а
  условна на status='streaming' OR metadata.finalizeFailed (repo.finalizeOwner) —
  перетирает reconcile-штамп, но не проставленный терминал. ВСЕ status-only
  штампы reconcile (stampTerminalIfStreaming, sweepStreaming) пишут строго
  onlyIfStreaming И мёржат metadata.finalizeFailed:true (иначе поздний owner-write
  не перетрёт).
- Порядок: попытка message-finalize → ран финализируется ВСЕГДА; при провале
  message onFinish помечает ран 'error' (не 'completed'). Ран не гейтится на
  message.
- Периодический reconcile-джоб (setInterval, env-tunable) клаузами по порядку:
  (a) пере-драйв зомби; (b) message streaming + ран терминален → штамп по статусу
  рана (succeeded-ран + зависшая строка → 'aborted'+finalizeFailed, НЕ
  'completed'-empty); (c) run running + НЕТ entry И НЕТ zombie + staleness →
  aborted (гейт «нет entry» первичный, staleness от last-progress updatedAt,
  X=max(2×per-call cap,15мин)+boot-warn); (d) message streaming + возраст>X + нет
  активной run-строки → aborted (двойной гейт). isInterruptResume исключает
  finalizeFailed-строки. Оппортунистический одно-чатовый reconcile при старте хода
  (best-effort, не фейлит ход). sweepStreaming boot-only → периодический.

Тесты (реальная БД): owner finalizeOwner чистит finalizeFailed; штамп не перетирает
терминал; поздний owner-write перетирает aborted-штамп; клаузы b/c/d (+живой entry
не трогать, двойной гейт d); «убить БД на finish → после восстановления ни строка,
ни ран не застряли». Юниты: finalizeFailed исключает interrupt-resume;
reconcileStaleRuns «нет entry».

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-11 07:26:01 +03:00
agent_coder 307ad49324 feat(ai-chat): единый гейт конкурентности + серверный supersede (CAS) (#487)
Раньше в legacy-режиме (дефолт!) гейта конкурентности НЕ было вообще — проверка
409 сидела внутри if (autonomousRuns && chatId): два таба = два параллельных
стрима на один чат (интерливинг истории, падения convertToModelMessages).

Серверная часть (клиентская лестница ретраев — отдельной клиентской итерацией
FSM, см. ниже):
- Run-строка ОБЯЗАТЕЛЬНА в ОБОИХ режимах: legacy тоже проходит beginRun/
  finalizeRun (гейт партиального уникального индекса + runId в start-метаданных).
  Различие режимов сведено к семантике abort: legacy onClose при дисконнекте
  зовёт requestStop(runId) вместо аборта сокет-сигнала (которого streamText уже
  не потребляет). Второй таб в ЛЮБОМ режиме → 409 A_RUN_ALREADY_ACTIVE.
- Supersede CAS POST /stream { supersede: { runId: X } } (AiChatRunService
  .supersede): валидация X.chatId===body.chatId (иначе 400 SUPERSEDE_INVALID);
  нет активного рана → degrade в обычный send; активный ≠ X → 409
  SUPERSEDE_TARGET_MISMATCH + текущий runId; активный = X → requestStop →
  awaitSettled (таймаут W=10c) → «записано/сдался» (сдался → settleZombie
  применяет intended условным UPDATE) → ready; таймаут → 409 SUPERSEDE_TIMEOUT,
  ничего не персистится. W обоснован race-on-abort'ом коммита 1; DB-brownout →
  TIMEOUT штатен, W не увеличивать (env-tunable).
- Задокументированы ограничения: нет квиесценции сайд-эффектов (в промпт нового
  рана добавлена строка SUPERSEDE_NOTE «предыдущий ран прерван, его последние
  операции могли примениться с задержкой»); кража слота между освобождением и
  beginRun (→ MISMATCH с новым runId, бэкстоп — уникальный индекс).

Тесты: два параллельных старта (оба режима) → строго 409 второму; supersede при
живом долгом ТУЛЕ (не UPDATE-задержке) settle'ится быстро; все CAS-ветки; дубль
supersede-POST → degrade; зомби-путь через settleZombie; HTTP-маппинг ветвей.
Кейс «beginRun ok → insert user-msg упал → ран settled, слот свободен» покрыт
lifecycle-спекой против реального safety-net catch.

КЛИЕНТ (deferred, точно flagged): удаление лестницы SUPERSEDE_RETRY_DELAYS_MS/
isRunAlreadyActive/supersedeRetryRef и переход на supersede-в-body требуют
адопции runId из start-метаданных (сейчас клиент runId не трекает вовсе) — это
и есть та «итерация клиентского FSM», которую данный серверный контракт
разблокирует. Не трогаю фрагильный клиентский FSM без возможности E2E-валидации
в браузере, чтобы не внести непроверяемую регрессию.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-11 07:26:01 +03:00
agent_coder 1031ed1ae8 feat(ai-chat): зомби-механика finalizeRun + settledPromise + условные терминальные UPDATE (#487)
Раньше give-up-путь finalizeRun ВОССТАНАВЛИВАЛ entry (зомби неотличим от живого
рана), а терминальный runRepo.update был БЕЗУСЛОВНЫМ (последний писатель перетирал
терминальный статус).

- Все терминальные UPDATE ранов теперь УСЛОВНЫЕ: новый repo.finalizeIfActive
  пишет только пока строка pending|running (зеркально onlyIfStreaming у сообщений)
  → двойной settle схлопывается в benign no-op, терминальный статус не перетереть.
- При исчерпании ретраев finalizeRun НЕ восстанавливает entry, а оставляет
  ЗОМБИ-запись { terminalWriteFailed, intended:{status,error} } в отдельной map;
  settleZombie пере-драйвит intended условным UPDATE (зовётся reconcile/supersede/
  boot sweep). Зомби удаляется после успешного UPDATE или обнаружения строки уже
  терминальной.
- Per-run settledPromise в ОТДЕЛЬНОЙ map runId→deferred (переживает active.delete),
  создаётся в beginRun, резолвится ровно один раз исходом (записано/сдался); поздний
  подписчик через ранее взятую ссылку получает резолвленный; peekSettled: live
  deferred → зомби-синтез → undefined (читать строку). Обе map bounded.
- Задокументирована потеря (single-process): рестарт до пере-драйва → boot sweep
  пишет 'aborted' поверх фактического intended — неустранимо в phase 1.

Тесты: юниты give-up-пути (зомби вместо restore + notifier terminalWriteFailed +
settleZombie), двойного settle, позднего подписчика; интеграционные против реальной
Б, что условный finalizeIfActive не перетирает терминал и двойной settle схлопывается.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-11 07:26:01 +03:00
agent_coder 3b6634c6be fix(ai-chat): in-app тулы — race-on-abort + safe-points + per-call cap (#487)
In-app тул-обёртки отбрасывали второй аргумент options с abortSignal — это был
единственный класс тулов без отмены и без wall-clock cap. Контент-мутации идут
через collab-WS (mutatePageContent), не через axios, поэтому «прокинуть signal в
axios» главный write-путь не покрывал.

- Переиспользован race-паттерн wrapToolWithCallTimeout: каждый in-app тул гонится
  против композитного сигнала (Stop + per-call cap); на abort — реджект
  немедленно, проигравший промис отбрасывается (латентность мс, не сетевой
  teardown — от этого зависит таймаут supersede в коммите 3).
- Safe-point проверки сигнала между последовательными вызовами paginateAll и
  пре-коммитная проверка в mutatePageContent (+ реентрантный близнец) через
  DocmostClient.setToolAbortSignal, который обёртка публикует перед каждым вызовом.
- Per-call cap покрывает весь вызов, env-tunable (AI_CHAT_INAPP_TOOL_CALL_CAP_MS,
  дефолт 2 мин).
- Задокументировано ограничение (#487): abort между вызовами многовызовного
  write-тула оставляет частично применённую операцию — отмена гарантирует «новый
  вызов не стартует», не «запись не доехала».

Тесты (честное свойство «после Stop не стартует новый HTTP/WS-вызов»): юнит
wrapInAppToolWithCap (реджект-на-abort, отсутствие новых вызовов, per-call cap,
публикация сигнала) + mock-HTTP тест реального paginateAll safe-point.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-11 07:26:01 +03:00
51 changed files with 6343 additions and 524 deletions
+41 -3
View File
@@ -225,11 +225,26 @@ MCP_DOCMOST_PASSWORD=
# Silence timeout (ms) for EXTERNAL-MCP transport ONLY (not the chat provider).
# Tighter than AI_STREAM_TIMEOUT_MS so a byte-silent/hung MCP server is broken in
# ~1 min instead of 15. Note it also cuts a legitimately long but byte-silent
# single tool call (a slow crawl that emits nothing until done) and an SSE
# transport idling >1 min BETWEEN tool calls. Default 60000 (1 min).
# ~1 min instead of 15. It cuts a legitimately long but byte-silent single tool
# call (a slow crawl that emits nothing until done) on the HTTP (streamable)
# transport, which opens a fresh request per call. The SSE transport — one
# long-lived body across many calls — is NO LONGER governed by this timeout
# (as of #489): its idle-BETWEEN-calls window has its own, raised bodyTimeout,
# AI_MCP_SSE_BODY_TIMEOUT_MS below. Default 60000 (1 min).
# AI_MCP_STREAM_TIMEOUT_MS=60000
# bodyTimeout (ms) for the EXTERNAL-MCP SSE transport ONLY (#489). The SSE
# transport holds ONE response body open across many tool calls, so undici's
# bodyTimeout (time between body bytes) counts the LEGITIMATE silence BETWEEN the
# model's tool calls, not just a hung single call. At the tight 1-min silence
# timeout above, a normal >1-min gap between calls would break the SSE socket and
# the cache would serve a dead client until TTL — so the SSE transport gets its
# OWN, RAISED bodyTimeout. A single stuck call is still bounded by the per-call
# cap (AI_MCP_CALL_TIMEOUT_MS), and a socket that does break is healed by the
# in-run transport-error retry. The HTTP (streamable) transport keeps the tight
# timeout. Default 600000 (10 min).
# AI_MCP_SSE_BODY_TIMEOUT_MS=600000
# Total wall-clock cap (ms) for ONE external MCP tool call (app-level, not
# transport). Aborts a tool that keeps the socket warm (SSE heartbeats / trickle)
# but never returns a result — which the silence timeout above never breaks.
@@ -288,6 +303,29 @@ MCP_DOCMOST_PASSWORD=
# registry is process-local).
# AI_CHAT_RESUMABLE_STREAM=false
# --- Run lifecycle tunables (#487) ---
# These govern the universal run machinery (every turn is now a first-class run,
# both modes) and rarely need changing.
#
# How long a server-side SUPERSEDE ("interrupt and send now") waits for the target
# run to settle after issuing Stop before it degrades to a 409 SUPERSEDE_TIMEOUT
# (nothing sent, the composer keeps the user's text). 10s is generous under a
# healthy DB; do NOT raise it to paper over a slow DB — a SUPERSEDE_TIMEOUT is the
# honest signal. Default 10000 (10s).
# AI_CHAT_SUPERSEDE_TIMEOUT_MS=10000
#
# How often the periodic bidirectional reconcile job runs (heals runs/messages
# left dangling by a crash or a lost terminal write). Default 120000 (2 min).
# AI_CHAT_RECONCILE_INTERVAL_MS=120000
#
# Wall-clock cap for a SINGLE in-app tool call (a long paginated read, or a content
# write whose collab commit hangs) — the per-call half of the composite abort
# signal every in-app tool is wrapped with (the other half is the turn's Stop).
# The reconcile staleness floor is derived as max(2 x this cap, 15min), so a very
# high value delays stale-run recovery (the server boot-warns above 30min). Default
# 120000 (2 min).
# AI_CHAT_INAPP_TOOL_CALL_CAP_MS=120000
# --- Anonymous public-share AI assistant ---
# Opt-in per workspace (AI settings -> "public share assistant"; off by default).
# When enabled, anonymous visitors of a published share can ask an AI about that
+4
View File
@@ -29,6 +29,10 @@ packages/mcp/build/
# is a build artifact like build/ — never committed, always fresh.
packages/mcp/src/registry-stamp.generated.ts
# token-estimate compiled output (#490; built in CI/Docker via `pnpm build` /
# the server `pretest`, never committed, so src/ and prod can never diverge).
packages/token-estimate/dist/
# Logs
logs
*.log
+1 -1
View File
@@ -455,7 +455,7 @@ The API server is a Fastify app with a global `/api` prefix (`main.ts` excludes
- `core/ai-chat/tools/` — the agent's ~40 read+write tools. Every tool runs under the **calling user's** CASL permissions via a per-user loopback access token (`docmost-client.loader.ts`), so the agent can never exceed what the user could do. Only **reversible** operations are exposed (page history + trash; no permanent delete). Agent edits get an "AI agent" provenance badge in page history (`20260616T130000-agent-provenance` migration).
- `core/ai-chat/embedding/` — RAG indexer + a BullMQ consumer on `AI_QUEUE` that embeds pages into `page_embeddings` (vector search), complementing Postgres full-text search. Pages are (re)indexed on edit; `AI_EMBEDDING_TIMEOUT_MS` bounds a hung embeddings endpoint.
- `core/ai-chat/external-mcp/` — admins can attach external MCP servers (e.g. Tavily) to give the agent web access. **`ssrf-guard.ts` validates outbound MCP URLs against SSRF** — keep that guard in the path when touching external-MCP connection logic.
- `core/ai-chat/ai-chat-run.service.ts` + `ai_chat_runs`**detached/autonomous agent runs** (`#184`), behind the per-workspace `settings.ai.autonomousRuns` flag (off by default). When on, a turn becomes a server-side RUN that survives a browser disconnect; only an explicit `POST /ai-chat/stop` ends it, and a client reconnects/live-follows via `POST /ai-chat/run`. **DEPLOY CONSTRAINT — single-instance only in phase 1:** Stop and the AbortController that backs it are process-local, so a Stop only aborts a run executing on the **same** replica that owns it (cross-instance pub/sub stop is phase 2). Do **not** enable `autonomousRuns` on a horizontally-scaled deployment (multiple replicas behind a load balancer, or Docmost cloud `CLOUD=true`) — run a single instance instead. The server logs a startup WARNING when it detects a multi-instance deployment (`CLOUD=true`) so the constraint is visible. The startup sweep settles any run left dangling by a restart.
- `core/ai-chat/ai-chat-run.service.ts` + `ai_chat_runs`**every agent turn is now a first-class server-side RUN** (`#184`, universalized in `#487`): its lifecycle is tracked in `ai_chat_runs` in **both** modes, and the single-active-run-per-chat concurrency gate is enforced universally (a legacy second tab now gets a clean `409 A_RUN_ALREADY_ACTIVE` instead of a second parallel stream that interleaved history). The per-workspace `settings.ai.autonomousRuns` flag (off by default) **no longer gates whether a turn is a run** — it now controls **only the browser-disconnect semantics**: when ON the run is *detached* (a disconnect leaves it executing server-side; only an explicit `POST /ai-chat/stop` ends it, and a client reconnects/live-follows via `POST /ai-chat/run`); when OFF (legacy) a disconnect ends the turn by stopping its run via the run's stop lever. `#487` also adds a server-side **supersede** CAS ("interrupt and send now") to `POST /ai-chat/stream` (`supersede: { runId }`): it atomically stops the chat's currently-active run and waits for it to settle before the new turn claims the slot, returning `SUPERSEDE_INVALID` / `SUPERSEDE_TARGET_MISMATCH` / `SUPERSEDE_TIMEOUT` on the non-proceed branches. **DEPLOY CONSTRAINT — single-instance only in phase 1:** Stop and the AbortController that backs it are process-local, so a Stop only aborts a run executing on the **same** replica that owns it (cross-instance pub/sub stop is phase 2). Do **not** enable `autonomousRuns` on a horizontally-scaled deployment (multiple replicas behind a load balancer, or Docmost cloud `CLOUD=true`) — run a single instance instead. The server logs a startup WARNING when it detects a multi-instance deployment (`CLOUD=true`) so the constraint is visible. The startup sweep settles any run left dangling by a restart.
### Client structure
Vite SPA. Code is organized by feature under `apps/client/src/features/*` (mirrors the server domains: `page`, `space`, `comment`, `ai-chat`, `editor`, …). Conventions:
+39
View File
@@ -202,6 +202,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
dangling by a restart. Phase 1 is single-instance-only (cross-instance Stop is
not yet reliable); the server warns at startup on a horizontally-scaled
deployment. (#184)
- **Server-side "interrupt and send now" (supersede) for AI chat.** `POST
/ai-chat/stream` now accepts a `supersede: { runId }` field: when the user sends
a new message while a run is active, the server atomically stops that run and
waits for it to settle before the new turn claims the chat's single run slot,
instead of the send being rejected as concurrent. The compare-and-set surfaces
three codes on its non-proceed branches — `SUPERSEDE_INVALID` (the targeted run
is malformed / belongs to another chat), `SUPERSEDE_TARGET_MISMATCH` (a
different run is now active; carries the current `activeRunId`), and
`SUPERSEDE_TIMEOUT` (the previous run did not stop within the settle window, so
nothing was sent and the composer keeps the text). Tunable via
`AI_CHAT_SUPERSEDE_TIMEOUT_MS` (default 10s). (#487)
- **Out-of-band page transfer via an in-RAM blob sandbox (`stash_page`).** A
new MCP tool serializes a whole page (its full ProseMirror JSON, with every
internal image/file mirrored) into an ephemeral in-RAM blob and returns only
@@ -282,6 +293,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
### Changed
- **Every AI-chat turn is now a first-class server-side run, and one run per chat
is enforced in both modes.** The run machinery from `#184` was universalized: a
turn is tracked in `ai_chat_runs` and gated by the single-active-run-per-chat
index regardless of the `settings.ai.autonomousRuns` flag. **Behavior change:**
a second tab (or a double-submit) that starts a turn while one is already active
on the chat is now rejected up front with `409 A_RUN_ALREADY_ACTIVE` (carrying
the `activeRunId`); previously, on the legacy path, it opened a second parallel
stream on the same chat that interleaved history. The `autonomousRuns` flag no
longer controls whether a turn is a run — it now governs **only** the
browser-disconnect semantics (ON = detached/survives a disconnect; OFF = a
disconnect stops the run). (#487)
- **Client markdown paste/copy and AI-chat rendering now go through the canonical
converter.** Pasting markdown into the editor, "Copy as markdown", the AI title
generator, and the AI-chat markdown renderer all now use
@@ -314,6 +336,23 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
### Fixed
- **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
the persisted history or the model context), and the turn is converted BEFORE
the user row is inserted, so a mid-flight failure cannot leave a duplicate
user row that a retry then compounds. A single part that still fails to convert
degrades to a `[tool context omitted]` marker on that one row instead of
bricking the whole chat. (#489)
- **A transport drop to an external MCP server now heals within the same turn.**
On an undici transport error, a read-only MCP tool reconnects its server and
retries once within the run; a write is never auto-retried (it may already have
applied). One flapping server no longer nulls the shared client cache, so other
servers' cached clients are untouched. The SSE transport also gets a raised
body-timeout so a legitimate >1-min idle between the model's tool calls no
longer breaks a long-lived SSE socket (new `AI_MCP_SSE_BODY_TIMEOUT_MS`, default
10 min; see `.env.example`). (#489)
- **The server no longer runs out of heap during long autonomous agent runs.** A
new pnpm patch on `ai@6.0.134` stops the SDK from building a cumulative
snapshot of the ENTIRE turn text on every streamed text-delta when no output
+1
View File
@@ -22,6 +22,7 @@
"@casl/react": "5.0.1",
"@docmost/editor-ext": "workspace:*",
"@docmost/prosemirror-markdown": "workspace:*",
"@docmost/token-estimate": "workspace:*",
"@excalidraw/excalidraw": "0.18.0-3a5ef40",
"@mantine/core": "8.3.18",
"@mantine/dates": "8.3.18",
@@ -6,10 +6,13 @@ describe("estimateTokens", () => {
expect(estimateTokens("")).toBe(0);
});
it("ceils chars/4 so any non-empty text is at least 1 token", () => {
// #490: migrated onto the shared @docmost/token-estimate module (chars/2.5, up
// from the old client-only chars/4) so the client counter and the server replay
// budgeter can never diverge.
it("ceils chars/2.5 so any non-empty text is at least 1 token", () => {
expect(estimateTokens("a")).toBe(1);
expect(estimateTokens("abcd")).toBe(1);
expect(estimateTokens("abcde")).toBe(2);
expect(estimateTokens("12345678")).toBe(2);
expect(estimateTokens("ab")).toBe(1);
expect(estimateTokens("abcde")).toBe(2); // 5 / 2.5 = 2
expect(estimateTokens("x".repeat(10))).toBe(4); // 10 / 2.5 = 4
});
});
@@ -2,18 +2,10 @@
* Rough client-side token estimation for AI-chat UI affordances.
*
* No provider streams exact per-token usage mid-stream, so any in-flight figure
* is a CLIENT ESTIMATE (chars/≈4 heuristic). Pure + unit-testable: it never runs
* a real BPE tokenizer (that would be O(n²) on the hot path, bloat the bundle,
* and be wrong for Gemini/Ollama anyway). Used by the in-body reasoning counter
* ("Thinking · N tokens").
* is a CLIENT ESTIMATE. This re-exports the SHARED estimator from
* `@docmost/token-estimate` (chars/2.5) so the in-body counter and the server's
* replay budgeter use the SAME heuristic — two divergent estimators would mean
* "the badge shows 60%" while "the budgeter already trimmed" (#490). Used by the
* in-body reasoning counter ("Thinking · N tokens").
*/
/**
* Rough token estimate for a piece of text using the standard chars/≈4 heuristic.
* Returns 0 for empty/whitespace-free-of-content input, and ceils so any
* non-empty text counts as at least one token.
*/
export function estimateTokens(text: string): number {
if (!text) return 0;
return Math.ceil(text.length / 4);
}
export { estimateTokens } from "@docmost/token-estimate";
+3 -1
View File
@@ -23,7 +23,7 @@
"migration:reset": "tsx src/database/migrate.ts down-to NO_MIGRATIONS",
"migration:codegen": "kysely-codegen --dialect=postgres --camel-case --env-file=../../.env --out-file=./src/database/types/db.d.ts",
"lint": "eslint \"{src,apps,libs,test}/**/*.ts\" --fix",
"pretest": "pnpm --filter @docmost/editor-ext build && pnpm --filter @docmost/prosemirror-markdown build",
"pretest": "pnpm --filter @docmost/editor-ext build && pnpm --filter @docmost/prosemirror-markdown build && pnpm --filter @docmost/token-estimate build",
"test": "jest",
"test:int": "jest --config test/jest-integration.json",
"test:watch": "jest --watch",
@@ -44,6 +44,7 @@
"@docmost/mcp": "workspace:*",
"@docmost/pdf-inspector": "1.9.6",
"@docmost/prosemirror-markdown": "workspace:*",
"@docmost/token-estimate": "workspace:*",
"@fastify/compress": "^9.0.0",
"@fastify/cookie": "^11.0.2",
"@fastify/multipart": "^10.0.0",
@@ -206,6 +207,7 @@
"^@docmost/db/(.*)$": "<rootDir>/database/$1",
"^@docmost/transactional/(.*)$": "<rootDir>/integrations/transactional/$1",
"^@docmost/ee/(.*)$": "<rootDir>/ee/$1",
"^@docmost/token-estimate$": "<rootDir>/../../../packages/token-estimate/src/index.ts",
"^src/(.*)$": "<rootDir>/$1",
"^@tiptap/react$": "<rootDir>/../test/stubs/tiptap-react.js"
}
@@ -43,6 +43,9 @@ function makeRepo(overrides: Record<string, jest.Mock> = {}) {
workspaceId: v.workspaceId,
})),
update: jest.fn(async () => ({ id: 'run-1' })),
// #487: terminal finalize now goes through the CONDITIONAL write. Default
// returns a truthy row (the run WAS active -> this call wrote it).
finalizeIfActive: jest.fn(async () => ({ id: 'run-1', status: 'succeeded' })),
markStopRequested: jest.fn(async () => ({ id: 'run-1' })),
findActiveByChat: jest.fn(async () => undefined),
findLatestByChat: jest.fn(async () => undefined),
@@ -336,14 +339,12 @@ describe('AiChatRunService run lifecycle', () => {
await svc.finalizeRun('run-1', 'ws-1', 'error', 'provider blew up');
expect(svc.isLocallyActive('run-1')).toBe(false);
expect(repo.update).toHaveBeenCalledWith(
// #487: the terminal write is CONDITIONAL (finalizeIfActive); finishedAt is
// stamped inside the repo method, so the service passes just status + error.
expect(repo.finalizeIfActive).toHaveBeenCalledWith(
'run-1',
'ws-1',
expect.objectContaining({
status: 'failed',
error: 'provider blew up',
finishedAt: expect.any(Date),
}),
expect.objectContaining({ status: 'failed', error: 'provider blew up' }),
);
});
@@ -366,8 +367,8 @@ describe('AiChatRunService run lifecycle', () => {
// A second settle (e.g. a streamText callback firing after the catch) no-ops.
await svc.finalizeRun('run-1', 'ws-1', 'completed', undefined);
expect(repo.update).toHaveBeenCalledTimes(1);
expect(repo.update).toHaveBeenCalledWith(
expect(repo.finalizeIfActive).toHaveBeenCalledTimes(1);
expect(repo.finalizeIfActive).toHaveBeenCalledWith(
'run-1',
'ws-1',
expect.objectContaining({ status: 'failed', error: 'first' }),
@@ -389,8 +390,8 @@ describe('AiChatRunService run lifecycle', () => {
const updateGate = new Promise((res) => {
resolveUpdate = res;
});
const update = jest.fn(() => updateGate);
const repo = makeRepo({ update });
const finalizeIfActive = jest.fn(() => updateGate);
const repo = makeRepo({ finalizeIfActive });
const svc = new AiChatRunService(repo as never, makeEnv() as never);
await svc.beginRun({
chatId: 'chat-1',
@@ -399,23 +400,23 @@ describe('AiChatRunService run lifecycle', () => {
});
// Fire both before the (pending) update resolves. The first synchronously
// claims the entry (active.delete) and awaits update; the second, started in
// the same macrotask, finds the entry already gone and returns at the claim
// WITHOUT ever calling update.
// claims the entry (active.delete) and awaits the write; the second, started
// in the same macrotask, finds the entry already gone and returns at the claim
// WITHOUT ever writing.
const p1 = svc.finalizeRun('run-1', 'ws-1', 'completed');
const p2 = svc.finalizeRun('run-1', 'ws-1', 'error', 'safety-net');
// The decisive assertion: exactly one caller reached the terminal UPDATE.
expect(update).toHaveBeenCalledTimes(1);
expect(finalizeIfActive).toHaveBeenCalledTimes(1);
// Let the single in-flight update land; both calls resolve cleanly.
resolveUpdate({ id: 'run-1' });
resolveUpdate({ id: 'run-1', status: 'succeeded' });
await Promise.all([p1, p2]);
expect(update).toHaveBeenCalledTimes(1);
expect(finalizeIfActive).toHaveBeenCalledTimes(1);
// The winner is the FIRST caller ('completed' -> 'succeeded'); the late
// 'error' settle never wrote, so it could not clobber the real status.
expect(update).toHaveBeenCalledWith(
expect(finalizeIfActive).toHaveBeenCalledWith(
'run-1',
'ws-1',
expect.objectContaining({ status: 'succeeded' }),
@@ -431,10 +432,10 @@ describe('AiChatRunService run lifecycle', () => {
// 409s until a restart. The fix updates FIRST and retries.
let calls = 0;
const repo = makeRepo({
update: jest.fn(async () => {
finalizeIfActive: jest.fn(async () => {
calls += 1;
if (calls === 1) throw new Error('deadlock detected');
return { id: 'run-1' };
return { id: 'run-1', status: 'succeeded' };
}),
});
jest.spyOn(Logger.prototype, 'warn').mockImplementation(() => undefined);
@@ -447,26 +448,29 @@ describe('AiChatRunService run lifecycle', () => {
await svc.finalizeRun('run-1', 'ws-1', 'completed');
// The retry landed the terminal write: the entry is dropped (slot freed) and
// the row carries the real terminal status — NOT stranded at 'running'.
// The retry landed the terminal write: the entry is dropped (slot freed), no
// zombie left, and the row carries the real terminal status.
expect(svc.isLocallyActive('run-1')).toBe(false);
expect(repo.update).toHaveBeenCalledTimes(2);
expect(repo.update).toHaveBeenLastCalledWith(
expect(svc.hasZombie('run-1')).toBe(false);
expect(repo.finalizeIfActive).toHaveBeenCalledTimes(2);
expect(repo.finalizeIfActive).toHaveBeenLastCalledWith(
'run-1',
'ws-1',
expect.objectContaining({ status: 'succeeded' }),
);
});
it('F6: if the terminal write keeps failing, the entry is RETAINED and a LATER settle completes it (chat not permanently 409d)', async () => {
it('#487 give-up: if the terminal write keeps failing, finalizeRun leaves a ZOMBIE (does NOT restore the entry) and settleZombie re-drives it', async () => {
// Worst case: the DB is down for the whole first finalize (all attempts fail).
// The run must NOT be silently lost — the entry stays so a subsequent settle
// (a streamText callback, requestStop -> onAbort, or a future sweep) can retry.
// #487 changes the give-up behaviour: the entry is NOT restored (a restored
// entry is indistinguishable from a live run). Instead a ZOMBIE record holds
// the intended terminal status, and a re-drive (settleZombie — called by the
// reconcile / supersede / opportunistic paths) applies it later.
let healthy = false;
const repo = makeRepo({
update: jest.fn(async () => {
finalizeIfActive: jest.fn(async () => {
if (!healthy) throw new Error('pool exhausted');
return { id: 'run-1' };
return { id: 'run-1', status: 'succeeded' };
}),
});
jest.spyOn(Logger.prototype, 'warn').mockImplementation(() => undefined);
@@ -480,35 +484,83 @@ describe('AiChatRunService run lifecycle', () => {
userId: 'user-1',
});
// First settle: every bounded attempt fails -> entry retained, NOT settled.
// First settle: every bounded attempt fails -> ZOMBIE, entry NOT restored.
await svc.finalizeRun('run-1', 'ws-1', 'completed');
expect(svc.isLocallyActive('run-1')).toBe(true);
// F12: the give-up emits ONE explicit, greppable ERROR (run + chat context)
// so an operator can tell "gave up, run held in memory" from a per-attempt
// blip — distinct from the per-attempt warns.
expect(svc.isLocallyActive('run-1')).toBe(false); // NOT a live entry
expect(svc.hasZombie('run-1')).toBe(true);
expect(svc.zombieRunIds()).toContain('run-1');
// The give-up emits ONE explicit, greppable ERROR mentioning the zombie.
const gaveUp = errorSpy.mock.calls.some(
(c) =>
/NON-TERMINAL/.test(String(c[0])) &&
/ZOMBIE/.test(String(c[0])) &&
/run-1/.test(String(c[0])) &&
/chat-1/.test(String(c[0])),
);
expect(gaveUp).toBe(true);
// The settle notifier resolved as terminalWriteFailed (a subscriber learns the
// slot still needs the intended status applied).
const outcome = await svc.peekSettled('run-1');
expect(outcome).toEqual({
status: 'succeeded',
error: null,
terminalWriteFailed: true,
});
// The DB recovers; a later settle now succeeds and frees the slot.
// The DB recovers; a re-drive settles the zombie via the conditional UPDATE.
healthy = true;
await svc.finalizeRun('run-1', 'ws-1', 'completed');
expect(svc.isLocallyActive('run-1')).toBe(false);
expect(repo.update).toHaveBeenLastCalledWith(
const redriven = await svc.settleZombie('run-1');
expect(redriven).toBe(true);
expect(svc.hasZombie('run-1')).toBe(false);
expect(repo.finalizeIfActive).toHaveBeenLastCalledWith(
'run-1',
'ws-1',
expect.objectContaining({ status: 'succeeded' }),
);
// And it is now idempotent: a further settle no-ops (terminal row already
// written), so a double-settle can never clobber the real status.
const callsBefore = repo.update.mock.calls.length;
// A later finalizeRun is idempotent (row already terminal): it no-ops at the
// once-gate, never re-writing.
const callsBefore = repo.finalizeIfActive.mock.calls.length;
await svc.finalizeRun('run-1', 'ws-1', 'error', 'late');
expect(repo.update).toHaveBeenCalledTimes(callsBefore);
expect(repo.finalizeIfActive).toHaveBeenCalledTimes(callsBefore);
});
it('#487 double-settle collapses to a benign no-op (conditional write; notifier resolves once)', async () => {
// A second concurrent settle is stopped at the synchronous active.delete
// claim, so the terminal write runs exactly once and the notifier resolves
// exactly once with the FIRST settler's outcome.
const repo = makeRepo();
const svc = new AiChatRunService(repo as never, makeEnv() as never);
await svc.beginRun({ chatId: 'chat-1', workspaceId: 'ws-1', userId: 'u1' });
await svc.finalizeRun('run-1', 'ws-1', 'aborted');
await svc.finalizeRun('run-1', 'ws-1', 'error', 'late'); // no-op
expect(repo.finalizeIfActive).toHaveBeenCalledTimes(1);
const outcome = await svc.peekSettled('run-1');
// peekSettled after resolve+delete falls through (notifier dropped, no zombie)
// -> undefined; the FIRST settler already resolved any earlier subscriber.
expect(outcome).toBeUndefined();
});
it('#487 late settledPromise subscriber gets the resolved outcome', async () => {
const repo = makeRepo();
const svc = new AiChatRunService(repo as never, makeEnv() as never);
await svc.beginRun({ chatId: 'chat-1', workspaceId: 'ws-1', userId: 'u1' });
// Subscribe BEFORE settle: hold the promise reference (as supersede does).
const early = svc.peekSettled('run-1');
expect(early).toBeDefined();
await svc.finalizeRun('run-1', 'ws-1', 'completed');
// The reference grabbed before settle resolves with the written outcome, even
// though the notifier was dropped from the map on resolve (bounded).
await expect(early).resolves.toEqual({
status: 'succeeded',
error: null,
terminalWriteFailed: false,
});
});
it('recordStep / linkAssistantMessage are best-effort: a repo failure is swallowed', async () => {
@@ -525,3 +577,197 @@ describe('AiChatRunService run lifecycle', () => {
).resolves.toBeUndefined();
});
});
describe('#487 AiChatRunService.supersede (CAS)', () => {
const chat = 'chat-1';
const ws = 'ws-1';
it('degrade: no active run on the chat -> caller sends a normal turn', async () => {
const repo = makeRepo({
findById: jest.fn(async () => undefined),
findActiveByChat: jest.fn(async () => undefined),
});
const svc = new AiChatRunService(repo as never, makeEnv() as never);
expect(await svc.supersede(chat, 'run-x', ws)).toEqual({ kind: 'degrade' });
});
it('invalid: the target run belongs to a DIFFERENT chat -> 400', async () => {
const repo = makeRepo({
findById: jest.fn(async () => ({
id: 'run-x',
chatId: 'other-chat',
workspaceId: ws,
})),
});
const svc = new AiChatRunService(repo as never, makeEnv() as never);
expect(await svc.supersede(chat, 'run-x', ws)).toEqual({ kind: 'invalid' });
});
it('mismatch: a DIFFERENT run is active than the one targeted -> current runId', async () => {
const repo = makeRepo({
findById: jest.fn(async () => ({ id: 'run-x', chatId: chat, workspaceId: ws })),
findActiveByChat: jest.fn(async () => ({
id: 'run-live',
chatId: chat,
workspaceId: ws,
status: 'running',
})),
});
const svc = new AiChatRunService(repo as never, makeEnv() as never);
expect(await svc.supersede(chat, 'run-x', ws)).toEqual({
kind: 'mismatch',
activeRunId: 'run-live',
});
});
it('ready: the target IS active -> stop it, await its (fast) settle, free the slot', async () => {
// Simulate a live long TOOL (NOT a slow UPDATE): the run stays active until an
// explicit Stop unwinds it; commit-1's race makes that settle land quickly.
// The abort listener stands in for streamText's onAbort -> finalizeRun.
const repo = makeRepo({
findById: jest.fn(async () => ({
id: 'run-1',
chatId: chat,
workspaceId: ws,
status: 'aborted',
error: null,
})),
findActiveByChat: jest.fn(async () => ({
id: 'run-1',
chatId: chat,
workspaceId: ws,
status: 'running',
})),
});
const svc = new AiChatRunService(repo as never, makeEnv() as never);
const handle = await svc.beginRun({ chatId: chat, workspaceId: ws, userId: 'u1' });
handle.signal.addEventListener('abort', () => {
void svc.finalizeRun('run-1', ws, 'aborted');
});
// supersede: getRun -> getActiveByChat(==target) -> requestStop -> the abort
// listener settles the run -> awaitSettled resolves -> ready.
expect(await svc.supersede(chat, 'run-1', ws, 10_000)).toEqual({
kind: 'ready',
});
expect(handle.signal.aborted).toBe(true); // Stop reached the run
});
it('timeout: the target never settles within W -> 409 SUPERSEDE_TIMEOUT (nothing persisted)', async () => {
const repo = makeRepo({
findById: jest.fn(async () => ({ id: 'run-1', chatId: chat, workspaceId: ws })),
findActiveByChat: jest.fn(async () => ({
id: 'run-1',
chatId: chat,
workspaceId: ws,
status: 'running',
})),
});
const svc = new AiChatRunService(repo as never, makeEnv() as never);
await svc.beginRun({ chatId: chat, workspaceId: ws, userId: 'u1' });
// Do NOT settle the run: a tiny W elapses -> timeout.
const result = await svc.supersede(chat, 'run-1', ws, 30);
expect(result).toEqual({ kind: 'timeout' });
});
it('ready then a DUPLICATE supersede POST degrades (the run is already gone)', async () => {
let active: unknown = {
id: 'run-1',
chatId: chat,
workspaceId: ws,
status: 'running',
};
const repo = makeRepo({
findById: jest.fn(async () => ({
id: 'run-1',
chatId: chat,
workspaceId: ws,
status: 'aborted',
error: null,
})),
findActiveByChat: jest.fn(async () => active),
finalizeIfActive: jest.fn(async () => {
active = undefined; // settling frees the active slot
return { id: 'run-1', status: 'aborted' };
}),
});
const svc = new AiChatRunService(repo as never, makeEnv() as never);
const handle = await svc.beginRun({ chatId: chat, workspaceId: ws, userId: 'u1' });
handle.signal.addEventListener('abort', () => {
void svc.finalizeRun('run-1', ws, 'aborted');
});
expect(await svc.supersede(chat, 'run-1', ws, 10_000)).toEqual({
kind: 'ready',
});
// The duplicate POST for the same target now finds no active run -> degrade.
expect(await svc.supersede(chat, 'run-1', ws)).toEqual({ kind: 'degrade' });
});
it('reconcileStaleRuns: aborts a stale run with NO entry/zombie; NEVER touches a live entry', async () => {
const finalizeIfActive = jest.fn(async () => ({ id: 'x', status: 'aborted' }));
const repo = makeRepo({
insert: jest.fn(async (v: any) => ({
id: 'live-1',
status: 'running',
chatId: v.chatId,
workspaceId: v.workspaceId,
})),
finalizeIfActive,
findStaleActive: jest.fn(async () => [
{ id: 'orphan-1', workspaceId: ws, chatId: 'c-orphan' },
{ id: 'live-1', workspaceId: ws, chatId: 'c-live' },
]),
});
const svc = new AiChatRunService(repo as never, makeEnv() as never);
// A LIVE run this replica owns (in the `active` map).
await svc.beginRun({ chatId: 'c-live', workspaceId: ws, userId: 'u1' });
expect(svc.isLocallyActive('live-1')).toBe(true);
const aborted = await svc.reconcileStaleRuns(15 * 60 * 1000);
expect(aborted).toBe(1);
// The orphan (no entry) was aborted; the live entry was NEVER passed to the DB.
expect(finalizeIfActive).toHaveBeenCalledTimes(1);
expect(finalizeIfActive).toHaveBeenCalledWith(
'orphan-1',
ws,
expect.objectContaining({ status: 'aborted' }),
);
expect(svc.isLocallyActive('live-1')).toBe(true);
});
it('gave-up zombie: supersede applies the intended status (settleZombie) then is ready', async () => {
let healthy = false;
let active: unknown = {
id: 'run-1',
chatId: chat,
workspaceId: ws,
status: 'running',
};
const repo = makeRepo({
findById: jest.fn(async () => ({ id: 'run-1', chatId: chat, workspaceId: ws })),
findActiveByChat: jest.fn(async () => active),
finalizeIfActive: jest.fn(async () => {
if (!healthy) throw new Error('db down');
active = undefined;
return { id: 'run-1', status: 'aborted' };
}),
});
jest.spyOn(Logger.prototype, 'warn').mockImplementation(() => undefined);
jest.spyOn(Logger.prototype, 'error').mockImplementation(() => undefined);
const svc = new AiChatRunService(repo as never, makeEnv() as never);
await svc.beginRun({ chatId: chat, workspaceId: ws, userId: 'u1' });
// The run's terminal write gives up -> zombie (row still 'running').
await svc.finalizeRun('run-1', ws, 'aborted');
expect(svc.hasZombie('run-1')).toBe(true);
// The DB recovers; supersede awaits the (already-resolved, terminalWriteFailed)
// settle, then settleZombie applies the intended status -> ready.
healthy = true;
expect(await svc.supersede(chat, 'run-1', ws, 10_000)).toEqual({
kind: 'ready',
});
expect(svc.hasZombie('run-1')).toBe(false);
});
});
@@ -34,6 +34,88 @@ export class RunAlreadyActiveError extends Error {
export type TurnTerminalStatus = 'completed' | 'error' | 'aborted';
export type RunTerminalStatus = 'succeeded' | 'failed' | 'aborted';
/** The terminal run statuses — the row is done once it reads one of these. */
export const RUN_TERMINAL_STATUSES: readonly RunTerminalStatus[] = [
'succeeded',
'failed',
'aborted',
];
/** Whether a persisted run status is terminal (settled). */
export function isRunTerminal(status: string | null | undefined): boolean {
return (
status === 'succeeded' || status === 'failed' || status === 'aborted'
);
}
/**
* #487: the outcome a run's {@link AiChatRunService.finalizeRun} settled with.
* `terminalWriteFailed` = the terminal write GAVE UP after the bounded retry, so
* the row is still non-terminal ('running') and a ZOMBIE record holds the
* `intended` status for a later re-drive (reconcile / supersede / boot sweep). A
* subscriber (supersede, #487 commit 3) uses this to decide whether the slot is
* genuinely free or must first have the intended status applied.
*/
export interface RunSettleOutcome {
status: RunTerminalStatus;
error: string | null;
terminalWriteFailed: boolean;
}
/**
* #487: how long a supersede waits for the target run to settle after Stop before
* it degrades to `SUPERSEDE_TIMEOUT`. W=10s is generous under a HEALTHY DB: commit
* 1's race-on-abort makes an in-app tool abort->settle in ms/hundreds of ms, so a
* live run releases its slot well within the window. Under a DB brownout the
* timeout is normal (the write cannot land); W must NOT be raised to paper
* over a slow DB — a SUPERSEDE_TIMEOUT is the honest signal (nothing persisted,
* the composer keeps the user's text). Env-tunable for ops, default 10s.
*/
export const SUPERSEDE_SETTLE_TIMEOUT_MS = (() => {
const raw = Number(process.env.AI_CHAT_SUPERSEDE_TIMEOUT_MS);
return Number.isFinite(raw) && raw > 0 ? raw : 10_000;
})();
/**
* #487: the result of the supersede CAS ({@link AiChatRunService.supersede}).
* - `degrade` : no active run on the chat (it ended between click and POST) —
* the caller sends a NORMAL turn (NOT a mismatch);
* - `invalid` : the target runId belongs to a DIFFERENT chat (malformed CAS 400);
* - `mismatch` : a DIFFERENT run is active than the one the client targeted —
* 409 SUPERSEDE_TARGET_MISMATCH carrying the current `activeRunId`
* (the client does NOT auto-retry);
* - `timeout` : the target did not settle within W — 409 SUPERSEDE_TIMEOUT,
* nothing persisted;
* - `ready` : the target was stopped AND settled (or its zombie's intended was
* applied) — the slot is free; the caller may beginRun the new run.
*/
export type SupersedeResult =
| { kind: 'degrade' }
| { kind: 'invalid' }
| { kind: 'mismatch'; activeRunId: string }
| { kind: 'timeout' }
| { kind: 'ready' };
/** A one-shot settle notifier (#487): `resolve` is called EXACTLY ONCE. */
interface Deferred<T> {
promise: Promise<T>;
resolve: (value: T) => void;
}
/**
* #487: a run whose terminal write GAVE UP (every bounded attempt failed). The
* row is stranded non-terminal ('running'); this record is the ONLY thing that
* distinguishes it from a live run, and carries the `intended` terminal status so
* a re-drive can apply it via the conditional UPDATE. Process-local (phase-1
* single-process assumption): a restart drops it, and the boot sweep then writes
* 'aborted' over the intended — a documented loss (see finalizeRun).
*/
interface ZombieRun {
workspaceId: string;
chatId: string;
intended: { status: RunTerminalStatus; error: string | null };
}
export function mapTurnStatusToRun(
status: TurnTerminalStatus,
): RunTerminalStatus {
@@ -101,6 +183,22 @@ export class AiChatRunService implements OnModuleInit {
// uptime — negligible in phase 1's single process.
private readonly settled = new Set<string>();
// #487 runId -> one-shot settle notifier. Kept in a SEPARATE map from `active`
// ON PURPOSE: it must OUTLIVE the `active.delete` claim inside finalizeRun (the
// claim frees the slot the instant finalize starts), so a subscriber can still
// await the outcome after the entry is gone. Created in beginRun, resolved
// EXACTLY ONCE in finalizeRun, then removed (bounded). Absence => this replica
// has no live notifier: a subscriber falls back to the zombie map, then to the
// row (see peekSettled). Process-local (phase-1 single-process assumption).
private readonly settledPromises = new Map<string, Deferred<RunSettleOutcome>>();
// #487 runId -> ZOMBIE record: a run whose terminal write gave up (row stranded
// non-terminal). BOUNDED — an entry is added only on give-up and removed on a
// successful re-drive (settleZombie) or when the row is found already terminal;
// a process restart clears it (and the boot sweep settles the stranded row).
// Process-local (phase-1 single-process assumption).
private readonly zombies = new Map<string, ZombieRun>();
// Bounded retry for the terminal write (F6): a single PK UPDATE can fail
// transiently under many fire-and-forget writes (pool exhaustion, deadlock, a
// brief connection blip). Riding out that blip in-place matters because the
@@ -224,6 +322,10 @@ export class AiChatRunService implements OnModuleInit {
chatId: args.chatId,
workspaceId: args.workspaceId,
});
// #487: arm the one-shot settle notifier BEFORE returning, so a subscriber
// that races in immediately after begin always finds a promise to await. It
// is resolved exactly once when the run settles (or gives up).
this.settledPromises.set(run.id, this.makeDeferred<RunSettleOutcome>());
return { runId: run.id, signal: controller.signal };
}
@@ -263,47 +365,43 @@ export class AiChatRunService implements OnModuleInit {
}
/**
* Finalize a run to its terminal status (succeeded / failed / aborted),
* stamping finishedAt + any error. Best-effort, but ROBUST against a transient
* terminal-write failure (F6) AND atomically safe against a concurrent settle.
* Finalize a run to its terminal status (succeeded / failed / aborted) via a
* CONDITIONAL UPDATE, stamping finishedAt + any error. Atomically safe against a
* concurrent settle AND robust against a transient terminal-write failure.
*
* ATOMIC ONCE-CLAIM (the gate must close in ONE synchronous tick): two
* finalizeRun calls for the SAME run can race — the documented real path is
* AiChatService.stream's safety-net catch settling the turn to 'error' while a
* streamText terminal callback (onFinish/onAbort/onError) ALSO settles it. The
* `settled.has` check alone is NOT a gate: it is read BEFORE the awaited UPDATE,
* so two callers can both see `false` and both write the row (last-write-wins
* clobbers the real terminal status, and the bounded retry only widens that
* window). The claim therefore happens via `active.delete`, a SYNCHRONOUS
* check-and-clear with NO await between the gate and the entry removal: the
* second concurrent caller finds the entry already gone and returns in the same
* tick, before any UPDATE. The transition "nobody is finalizing" -> "I am
* finalizing" is thus a single atomic step.
* claim happens via `active.delete`, a SYNCHRONOUS check-and-clear with NO await
* between the gate and the entry removal: the second concurrent caller finds the
* entry already gone and returns in the same tick, before any UPDATE.
*
* ORDER MATTERS (F6): once we own the claim, the terminal UPDATE happens FIRST;
* only once it SUCCEEDS do we record the run as settled. If the UPDATE fails on
* every bounded attempt we RESTORE the in-memory entry, leave the run UNsettled,
* and emit an ERROR signal that the row is left non-terminal 'running' (which
* would 409 every future turn in the chat until recovery). An in-process retry
* by a LATER settle is only POSSIBLE, never guaranteed: it needs (a) the entry
* to have been restored at the give-up path AND (b) a fresh settler to arrive
* AFTER that restore. A concurrent settler that arrives DURING the retry window
* — while the entry is deleted for backoff and not yet restored — is consumed at
* the synchronous `active.delete` claim (it finds nothing to delete and returns
* a no-op), so it does NOT become an in-process retrier. The NO-streamText path
* (the turn threw before streamText was wired, so ONLY the safety-net ever
* settles) likewise has no second in-process settler at all. The UNCONDITIONAL
* backstop in every case is the boot sweep on the next restart (phase 1 has no
* periodic in-process sweep); the retained entry is bounded (cleared on restart)
* and harmless meanwhile.
* ALL TERMINAL WRITES ARE CONDITIONAL (#487): `finalizeIfActive` only flips a
* row still in pending|running (mirror of the assistant message's
* `onlyIfStreaming`). So even a settle that DID reach the UPDATE (e.g. a
* reconcile stamp racing an owner finalize) can never clobber a terminal status
* — the loser matches nothing and is a benign no-op. `active.delete` is the
* fast, in-process gate; the conditional WHERE is the authoritative one.
*
* IDEMPOTENT on SUCCESS (#184 review): the terminal write happens AT MOST ONCE
* per run. After a successful write the once-gate keys off {@link settled} (the
* terminal row already written) so a settle arriving AFTER the entry was already
* dropped-and-settled returns early; a settle racing the in-flight write is
* stopped earlier still, by the `active.delete` claim. Either way a genuine
* double-settle collapses to a single write and a late settle can never clobber
* the real terminal status or double-write the row.
* ZOMBIE ON GIVE-UP (#487): if every bounded attempt THROWS (the DB is down for
* the whole finalize), we do NOT restore the entry. The row is stranded
* non-terminal ('running'); we record a ZOMBIE `{ terminalWriteFailed, intended
* }` (the ONLY thing distinguishing this dead run from a live one) and resolve
* the settle notifier with `terminalWriteFailed: true`. A restore would make the
* zombie indistinguishable from a live run to every reader; instead a re-drive
* (settleZombie, called by the periodic reconcile / supersede / opportunistic
* paths) applies the intended status later via the same conditional UPDATE.
*
* DOCUMENTED LOSS (#487, single-process phase 1): if the process RESTARTS before
* a zombie is re-driven, the in-memory zombie map is gone and the boot sweep
* (unconditional) writes 'aborted' over the ACTUAL intended status. This is
* unavoidable while the run lifecycle is single-process — there is no durable
* record of `intended`; a cross-process durable intent is deferred to phase 2.
*
* IDEMPOTENT: the settle notifier resolves EXACTLY ONCE; a second settle is
* stopped at `settled.has` or the `active.delete` claim, so a double-settle
* collapses to a single write and can never double-resolve or clobber the row.
*/
async finalizeRun(
runId: string,
@@ -314,13 +412,17 @@ export class AiChatRunService implements OnModuleInit {
// ---- Atomic once-claim (synchronous; NO await before the gate closes) ----
// Already terminally written -> idempotent no-op.
if (this.settled.has(runId)) return;
// Capture the entry BEFORE the delete so a total-failure path can restore it.
// Capture the entry BEFORE the delete for the give-up log context.
const entry = this.active.get(runId);
// SYNCHRONOUS check-and-clear: the FIRST caller deletes (claims) the entry;
// any concurrent SECOND caller finds nothing to delete and returns HERE, in
// the same tick, before any await — so it can never reach the UPDATE.
if (!this.active.delete(runId)) return;
const status = mapTurnStatusToRun(turnStatus);
const err = error ?? null;
const chatId = entry?.chatId ?? 'unknown';
let lastError: unknown;
for (
let attempt = 1;
@@ -328,47 +430,294 @@ export class AiChatRunService implements OnModuleInit {
attempt++
) {
try {
await this.runRepo.update(runId, workspaceId, {
status: mapTurnStatusToRun(turnStatus),
finishedAt: new Date(),
error: error ?? null,
const row = await this.runRepo.finalizeIfActive(runId, workspaceId, {
status,
error: err,
});
// Terminal write landed: arm the once-gate. The entry is already gone
// (claimed above); we do NOT restore it. The slot is now free.
// No throw => the row is now terminal (we wrote it, or it was ALREADY
// terminal — another writer won the conditional UPDATE, a benign no-op).
this.settled.add(runId);
this.zombies.delete(runId);
// Resolve with the persisted outcome: our status when WE wrote it, else
// the row's real terminal status (re-read on the already-terminal path so
// a subscriber never sees a status we did not actually persist).
const outcome: RunSettleOutcome = row
? { status, error: err, terminalWriteFailed: false }
: await this.readTerminalOutcome(runId, workspaceId, status, err);
this.resolveSettled(runId, outcome);
return;
} catch (err) {
lastError = err;
} catch (err2) {
lastError = err2;
this.logger.warn(
`Failed to finalize run ${runId} (attempt ${attempt}/${
AiChatRunService.FINALIZE_MAX_ATTEMPTS
}): ${err instanceof Error ? err.message : 'unknown error'}`,
}): ${err2 instanceof Error ? err2.message : 'unknown error'}`,
);
if (attempt < AiChatRunService.FINALIZE_MAX_ATTEMPTS) {
await this.delay(AiChatRunService.FINALIZE_RETRY_BASE_MS * attempt);
}
}
}
// Every attempt failed: this is a give-up, materially worse than a per-attempt
// blip — the row is left NON-TERMINAL ('running'), so emit ONE explicit,
// greppable ERROR so an operator can tell "survived a blip" from "gave up, run
// held in memory until recovery" (the last warn alone says only "attempt 3/3").
// Every attempt threw: GIVE UP. The row is stranded non-terminal ('running').
// Do NOT restore the entry (a restored entry is indistinguishable from a live
// run); leave a ZOMBIE record instead, and resolve the notifier as
// terminalWriteFailed so a subscriber knows the slot still needs the intended
// status applied. One explicit, greppable ERROR so an operator can tell a
// give-up from a per-attempt blip.
this.logger.error(
`Run ${runId} (chat ${entry?.chatId ?? 'unknown'}) left NON-TERMINAL ` +
`('running'): terminal write failed after ${
AiChatRunService.FINALIZE_MAX_ATTEMPTS
} attempts; entry retained in memory, recovery deferred to next settle / ` +
`boot sweep`,
`Run ${runId} (chat ${chatId}) left NON-TERMINAL ('running'): terminal ` +
`write failed after ${AiChatRunService.FINALIZE_MAX_ATTEMPTS} attempts; ` +
`ZOMBIE recorded (intended '${status}'), recovery deferred to reconcile / ` +
`supersede / boot sweep`,
lastError,
);
// RESTORE the claimed entry (and leave the run UNsettled) so a LATER settle
// that arrives AFTER this restore MAY retry the terminal write — but that
// in-process retry is NOT guaranteed (a concurrent settler caught in the retry
// window above is consumed at the `active.delete` claim, and the no-streamText
// path has no second settler at all). The UNCONDITIONAL backstop in every case
// is the boot sweep on the next restart; the restored entry is bounded and
// cleared on restart.
if (entry) this.active.set(runId, entry);
this.zombies.set(runId, {
workspaceId,
chatId,
intended: { status, error: err },
});
this.resolveSettled(runId, { status, error: err, terminalWriteFailed: true });
}
/**
* #487: re-drive a zombie run's intended terminal write (the conditional
* UPDATE). Called by the periodic reconcile (commit 4), an opportunistic
* single-chat reconcile, and supersede (commit 3). On success — the row is now
* terminal (written OR found already terminal) — the zombie is cleared and the
* once-gate armed; on another failure the zombie is kept for a later retry.
* Returns true when the row is now terminal. Best-effort; never throws.
*/
async settleZombie(runId: string): Promise<boolean> {
const z = this.zombies.get(runId);
if (!z) return false;
try {
await this.runRepo.finalizeIfActive(runId, z.workspaceId, {
status: z.intended.status,
error: z.intended.error,
});
this.zombies.delete(runId);
this.settled.add(runId);
return true;
} catch (err) {
this.logger.warn(
`Re-drive of zombie run ${runId} (chat ${z.chatId}) failed; will retry ` +
`later: ${err instanceof Error ? err.message : 'unknown error'}`,
);
return false;
}
}
/**
* #487 reconcile clause (c): abort runs the DB still shows active (pending|
* running) but that this replica does NOT own — NO live entry AND NO zombie —
* and that have been UNTOUCHED past `staleMs` (from last-progress `updated_at`,
* NOT startedAt, so a legit long marathon is never a candidate). "No entry" is
* the PRIMARY gate: a live entry (an actively-executing run on this replica) is
* NEVER aborted, whatever its age. Returns the number aborted. Best-effort —
* never throws (a periodic-job failure must not crash the process).
*/
async reconcileStaleRuns(staleMs: number): Promise<number> {
let candidates: Array<{ id: string; workspaceId: string; chatId: string }>;
try {
candidates = await this.runRepo.findStaleActive(staleMs);
} catch (err) {
this.logger.warn(
`Reconcile (stale runs) query failed: ${
err instanceof Error ? err.message : 'unknown error'
}`,
);
return 0;
}
let aborted = 0;
for (const c of candidates) {
// PRIMARY gate: never touch a live entry, and never race a zombie we are
// already re-driving (settleZombie owns those).
if (this.active.has(c.id) || this.zombies.has(c.id)) continue;
try {
const row = await this.runRepo.finalizeIfActive(c.id, c.workspaceId, {
status: 'aborted',
error: 'Run aborted by reconcile: no live runner (stale).',
});
if (row) {
aborted += 1;
this.settled.add(c.id);
}
} catch (err) {
this.logger.warn(
`Reconcile abort of stale run ${c.id} failed: ${
err instanceof Error ? err.message : 'unknown error'
}`,
);
}
}
return aborted;
}
/**
* #487: the run's settle outcome as seen by THIS replica, or undefined when it
* has no record (the caller then reads the row — the DB is the source of truth).
* A LIVE deferred (still settling, or resolved-but-not-yet-consumed) wins; a
* ZOMBIE synthesizes the give-up outcome. A subscriber (supersede) races this
* against a timeout.
*/
peekSettled(runId: string): Promise<RunSettleOutcome> | undefined {
const d = this.settledPromises.get(runId);
if (d) return d.promise;
const z = this.zombies.get(runId);
if (z) {
return Promise.resolve({
status: z.intended.status,
error: z.intended.error,
terminalWriteFailed: true,
});
}
return undefined;
}
/**
* #487: await a run's settle outcome, bounded by `timeoutMs`. Returns the
* outcome on settle, or undefined on TIMEOUT (or when this replica has no record
* of the run and its row is not terminal). Uses the LIVE settle notifier / the
* zombie synth when present; else reads the row (the DB is the source of truth
* once the in-memory record is gone). The subscriber (supersede) grabs this
* right after Stop; commit 1's race makes the settle land in ms on a healthy DB.
*/
async awaitSettled(
runId: string,
workspaceId: string,
timeoutMs: number,
): Promise<RunSettleOutcome | undefined> {
const pending = this.peekSettled(runId);
if (pending) {
let timer: ReturnType<typeof setTimeout> | undefined;
const timeout = new Promise<undefined>((resolve) => {
timer = setTimeout(() => resolve(undefined), timeoutMs);
timer.unref?.();
});
try {
return await Promise.race([pending, timeout]);
} finally {
if (timer) clearTimeout(timer);
}
}
// No live notifier and no zombie: read the row (already settled-and-written,
// or unknown here). A terminal row is an outcome; anything else -> undefined.
const row = await this.runRepo.findById(runId, workspaceId);
if (row && isRunTerminal(row.status)) {
return {
status: row.status as RunTerminalStatus,
error: row.error ?? null,
terminalWriteFailed: false,
};
}
return undefined;
}
/**
* #487: the SERVER supersede CAS for `POST /stream { supersede: { runId: X } }`.
* Atomically transitions "X is the chat's active run" -> "X is stopped, settled,
* slot free" so the caller can start a replacement run. See {@link
* SupersedeResult} for the branch semantics.
*
* On a `ready` result the caller MUST still go through the normal beginRun gate
* (the partial unique index) — between the slot freeing here and beginRun a
* neighbouring tab's ordinary POST can win the slot (documented SLOT-THEFT: the
* loser then gets a MISMATCH carrying the NEW runId). There is also NO side-
* effect quiescence: an in-flight write of the stopped run may still land AFTER
* the new run starts (commit 1 stops the NEXT call, not one already committing),
* so the caller adds a prompt note to the new run.
*/
async supersede(
chatId: string,
targetRunId: string,
workspaceId: string,
timeoutMs: number = SUPERSEDE_SETTLE_TIMEOUT_MS,
): Promise<SupersedeResult> {
// Validate the target belongs to THIS chat (a CAS targeting another chat's run
// is malformed -> 400). A missing row is NOT invalid: the run may have ended
// and been pruned; the active-run check below decides degrade vs mismatch.
const target = await this.getRun(targetRunId, workspaceId);
if (target && target.chatId !== chatId) return { kind: 'invalid' };
const active = await this.getActiveForChat(chatId, workspaceId);
// No active run: it ended between the client's click and this POST — this is a
// DEGRADE to a normal send, NOT a mismatch (the user's intent still holds).
if (!active) return { kind: 'degrade' };
// A DIFFERENT run is active than the one the client saw -> mismatch. The
// client does not auto-retry; it surfaces the new runId.
if (active.id !== targetRunId) {
return { kind: 'mismatch', activeRunId: active.id };
}
// The target IS active: stop it, then await its settle within W.
await this.requestStop(targetRunId, workspaceId);
const outcome = await this.awaitSettled(targetRunId, workspaceId, timeoutMs);
if (!outcome) return { kind: 'timeout' };
// Gave up (terminal write failed): apply the intended status via the
// conditional UPDATE so the slot actually frees. If that ALSO fails, the row
// is still stranded -> treat as a timeout (nothing persisted for the new run).
if (outcome.terminalWriteFailed) {
const settled = await this.settleZombie(targetRunId);
if (!settled) return { kind: 'timeout' };
}
return { kind: 'ready' };
}
/** #487 test/diagnostic seam: whether a give-up zombie is held for this run. */
hasZombie(runId: string): boolean {
return this.zombies.has(runId);
}
/** #487: every zombie runId held on this replica (reconcile clause a, commit 4). */
zombieRunIds(): string[] {
return [...this.zombies.keys()];
}
/** #487: create a one-shot deferred (resolve captured for a later single call). */
private makeDeferred<T>(): Deferred<T> {
let resolve!: (value: T) => void;
const promise = new Promise<T>((r) => {
resolve = r;
});
return { promise, resolve };
}
/** #487: resolve a run's settle notifier EXACTLY ONCE, then drop it (bounded).
* A subscriber that already grabbed the promise still resolves; a later one
* falls back to the zombie map / the row (see peekSettled). */
private resolveSettled(runId: string, outcome: RunSettleOutcome): void {
const d = this.settledPromises.get(runId);
if (!d) return;
this.settledPromises.delete(runId);
d.resolve(outcome);
}
/** #487: read the persisted terminal outcome when the conditional finalize was a
* no-op (the row was already terminal). Falls back to the intended status when
* the read fails or the row is unexpectedly missing/non-terminal. */
private async readTerminalOutcome(
runId: string,
workspaceId: string,
fallbackStatus: RunTerminalStatus,
fallbackError: string | null,
): Promise<RunSettleOutcome> {
try {
const row = await this.runRepo.findById(runId, workspaceId);
if (row && isRunTerminal(row.status)) {
return {
status: row.status as RunTerminalStatus,
error: row.error ?? null,
terminalWriteFailed: false,
};
}
} catch {
// Fall through to the intended status — best-effort only.
}
return {
status: fallbackStatus,
error: fallbackError,
terminalWriteFailed: false,
};
}
/** Small async backoff between terminal-write retries (F6). Isolated so it is
@@ -115,7 +115,7 @@ describe('finalizeAssistant dispatch (planFinalizeAssistant + applyFinalize)', (
// Drive the SAME applyFinalize the service calls (no duplicated logic).
async function dispatchFinalize(
repo: { insert: jest.Mock; update: jest.Mock },
repo: { insert: jest.Mock; finalizeOwner: jest.Mock },
assistantId: string | undefined,
flushed: AssistantFlush,
): Promise<void> {
@@ -135,21 +135,22 @@ describe('finalizeAssistant dispatch (planFinalizeAssistant + applyFinalize)', (
expect(planFinalizeAssistant(undefined)).toEqual({ kind: 'insert' });
});
it('(a) upfront insert succeeded -> finalize UPDATEs the row by id', async () => {
const repo = { insert: jest.fn(), update: jest.fn() };
it('(a) upfront insert succeeded -> finalize CONDITIONALLY updates the row by id (#487 owner-write)', async () => {
const repo = { insert: jest.fn(), finalizeOwner: jest.fn() };
const flushed = flushAssistant([], 'final answer', 'completed', {
finishReason: 'stop',
});
await dispatchFinalize(repo, 'a1', flushed);
expect(repo.update).toHaveBeenCalledWith('a1', workspaceId, flushed);
// #487: the owner write is the CONDITIONAL finalizeOwner, not a raw update.
expect(repo.finalizeOwner).toHaveBeenCalledWith('a1', workspaceId, flushed);
expect(repo.insert).not.toHaveBeenCalled();
});
it('(b) upfront insert failed -> finalize INSERTs the terminal payload', async () => {
const repo = { insert: jest.fn(), update: jest.fn() };
const repo = { insert: jest.fn(), finalizeOwner: jest.fn() };
const flushed = flushAssistant([], 'partial', 'error', { error: 'boom' });
await dispatchFinalize(repo, undefined, flushed);
expect(repo.update).not.toHaveBeenCalled();
expect(repo.finalizeOwner).not.toHaveBeenCalled();
expect(repo.insert).toHaveBeenCalledTimes(1);
const arg = repo.insert.mock.calls[0][0];
// The fallback insert carries the terminal content/status/metadata.
@@ -0,0 +1,279 @@
import {
BadRequestException,
ConflictException,
ForbiddenException,
HttpException,
} from '@nestjs/common';
import { AiChatController } from './ai-chat.controller';
import type { User, Workspace } from '@docmost/db/types/entity.types';
/**
* #487 commit 3 — the single concurrency GATE (both modes) + the server supersede
* CAS, at the controller boundary. The gate + CAS run BEFORE res.hijack(), so a
* rejected concurrent start / a CAS branch returns clean JSON (an HttpException
* the controller's post-hijack catch re-serializes). These assert the OBSERVABLE
* HTTP contract against the real controller + a stubbed run service.
*/
describe('#487 AiChatController.stream — gate + supersede', () => {
const user = { id: 'u1' } as User;
function wsWith(autonomousRuns: boolean): Workspace {
return {
id: 'ws1',
settings: { ai: { chat: true, autonomousRuns } },
} as unknown as Workspace;
}
function makeReqRes(body: Record<string, unknown>) {
const req = {
raw: { sessionId: 'sess', once: jest.fn(), destroyed: false },
body,
};
const res = {
raw: {
writableEnded: false,
headersSent: false,
on: jest.fn(),
once: jest.fn(),
setHeader: jest.fn(),
end: jest.fn(),
statusCode: 200,
flushHeaders: jest.fn(),
},
hijack: jest.fn(),
status: jest.fn().mockReturnThis(),
send: jest.fn(),
};
return { req, res };
}
function makeController(
runServiceOverrides: Record<string, jest.Mock>,
// The chat assertOwnedChat resolves. Default: a chat OWNED by `user` (u1), so
// the ownership gate is transparent to the gate/CAS assertions below. Pass a
// foreign-owner (or undefined) chat to exercise the #487 owner rejection.
chat: { creatorId: string } | undefined = { creatorId: 'u1' },
) {
const aiChatService = {
resolveRoleForRequest: jest.fn().mockResolvedValue(null),
getChatModel: jest.fn().mockResolvedValue({}),
stream: jest.fn().mockResolvedValue(undefined),
};
const aiChatRunService = {
getActiveForChat: jest.fn().mockResolvedValue(undefined),
supersede: jest.fn(),
beginRun: jest.fn().mockResolvedValue({
runId: 'run-new',
signal: new AbortController().signal,
}),
linkAssistantMessage: jest.fn(),
recordStep: jest.fn(),
finalizeRun: jest.fn(),
requestStop: jest.fn(),
...runServiceOverrides,
};
const aiChatRepo = { findById: jest.fn().mockResolvedValue(chat) };
const controller = new AiChatController(
aiChatService as never,
aiChatRunService as never,
aiChatRepo as never, // aiChatRepo
{} as never, // aiChatMessageRepo
{} as never, // aiTranscription
{} as never, // pageRepo
);
return { controller, aiChatService, aiChatRunService, aiChatRepo };
}
const codeOf = (err: unknown) =>
(((err as HttpException).getResponse() as Record<string, unknown>) ?? {})
.code;
describe('single concurrency gate — BOTH modes reject the second tab with 409', () => {
for (const autonomousRuns of [true, false]) {
it(`rejects a concurrent start with 409 A_RUN_ALREADY_ACTIVE (autonomousRuns=${autonomousRuns})`, async () => {
const { controller, aiChatRunService } = makeController({
getActiveForChat: jest
.fn()
.mockResolvedValue({ id: 'run-live', chatId: 'c1' }),
});
const { req, res } = makeReqRes({ chatId: 'c1' });
let thrown: unknown;
try {
await controller.stream(
req as never,
res as never,
user,
wsWith(autonomousRuns),
);
} catch (e) {
thrown = e;
}
expect(thrown).toBeInstanceOf(ConflictException);
expect((thrown as HttpException).getStatus()).toBe(409);
expect(codeOf(thrown)).toBe('A_RUN_ALREADY_ACTIVE');
// Rejected BEFORE committing to the stream (no hijack, no service.stream).
expect(res.hijack).not.toHaveBeenCalled();
expect(aiChatRunService.getActiveForChat).toHaveBeenCalledWith(
'c1',
'ws1',
);
});
}
});
// #487 [security, F1]: stream() MUST owner-gate an existing chat exactly like its
// six sibling endpoints, BEFORE the supersede CAS. Otherwise a same-workspace
// non-owner could POST a supersede against another user's chat and (a) harvest
// that user's active runId from the 409 SUPERSEDE_TARGET_MISMATCH body, then (b)
// requestStop the foreign run. The gate must reject FIRST — no run lookup, no
// supersede, no stop, no runId leak.
describe('cross-user ownership gate (F1)', () => {
it('a non-owner streaming against someone else\'s chat is rejected (403) with NO runId leak and NO foreign requestStop', async () => {
// A live run exists on the victim's chat. Without the gate the supersede CAS
// would run and (faithful to the run service) return a MISMATCH carrying the
// victim's runId — the exact leak. With the gate it must never be reached.
const getActiveForChat = jest
.fn()
.mockResolvedValue({ id: 'run-victim', chatId: 'c-other' });
const supersede = jest
.fn()
.mockResolvedValue({ kind: 'mismatch', activeRunId: 'run-victim' });
const requestStop = jest.fn();
const { controller, aiChatService } = makeController(
{ getActiveForChat, supersede, requestStop },
{ creatorId: 'someone-else' }, // the chat is NOT owned by u1
);
const { req, res } = makeReqRes({
chatId: 'c-other',
supersede: { runId: 'guessed-uuid' },
});
let thrown: unknown;
try {
await controller.stream(req as never, res as never, user, wsWith(true));
} catch (e) {
thrown = e;
}
// Rejected by the ownership gate (403), the SAME shape the neighbors use.
expect(thrown).toBeInstanceOf(ForbiddenException);
expect((thrown as HttpException).getStatus()).toBe(403);
// Crucially NOT a 409 that would carry activeRunId — no runId is leaked.
const payload = JSON.stringify(
(thrown as HttpException).getResponse() ?? {},
);
expect(payload).not.toContain('run-victim');
expect(codeOf(thrown)).not.toBe('SUPERSEDE_TARGET_MISMATCH');
// The gate short-circuits BEFORE any run machinery runs.
expect(getActiveForChat).not.toHaveBeenCalled();
expect(supersede).not.toHaveBeenCalled();
expect(requestStop).not.toHaveBeenCalled();
expect(aiChatService.stream).not.toHaveBeenCalled();
expect(res.hijack).not.toHaveBeenCalled();
});
});
it('supersede MISMATCH -> 409 SUPERSEDE_TARGET_MISMATCH carrying the current runId', async () => {
const { controller } = makeController({
supersede: jest
.fn()
.mockResolvedValue({ kind: 'mismatch', activeRunId: 'run-other' }),
});
const { req, res } = makeReqRes({
chatId: 'c1',
supersede: { runId: 'run-x' },
});
let thrown: unknown;
try {
await controller.stream(req as never, res as never, user, wsWith(true));
} catch (e) {
thrown = e;
}
expect(thrown).toBeInstanceOf(ConflictException);
expect(codeOf(thrown)).toBe('SUPERSEDE_TARGET_MISMATCH');
expect(
((thrown as HttpException).getResponse() as Record<string, unknown>)
.activeRunId,
).toBe('run-other');
expect(res.hijack).not.toHaveBeenCalled();
});
it('supersede TIMEOUT -> 409 SUPERSEDE_TIMEOUT, nothing streamed', async () => {
const { controller } = makeController({
supersede: jest.fn().mockResolvedValue({ kind: 'timeout' }),
});
const { req, res } = makeReqRes({
chatId: 'c1',
supersede: { runId: 'run-x' },
});
let thrown: unknown;
try {
await controller.stream(req as never, res as never, user, wsWith(false));
} catch (e) {
thrown = e;
}
expect(thrown).toBeInstanceOf(ConflictException);
expect(codeOf(thrown)).toBe('SUPERSEDE_TIMEOUT');
expect(res.hijack).not.toHaveBeenCalled();
});
it('supersede INVALID (target on another chat) -> 400 SUPERSEDE_INVALID', async () => {
const { controller } = makeController({
supersede: jest.fn().mockResolvedValue({ kind: 'invalid' }),
});
const { req, res } = makeReqRes({
chatId: 'c1',
supersede: { runId: 'run-x' },
});
let thrown: unknown;
try {
await controller.stream(req as never, res as never, user, wsWith(true));
} catch (e) {
thrown = e;
}
expect(thrown).toBeInstanceOf(BadRequestException);
expect(codeOf(thrown)).toBe('SUPERSEDE_INVALID');
});
it('supersede without chatId -> 400 SUPERSEDE_INVALID', async () => {
const { controller, aiChatRunService } = makeController({});
const { req, res } = makeReqRes({ supersede: { runId: 'run-x' } });
let thrown: unknown;
try {
await controller.stream(req as never, res as never, user, wsWith(true));
} catch (e) {
thrown = e;
}
expect(thrown).toBeInstanceOf(BadRequestException);
expect(codeOf(thrown)).toBe('SUPERSEDE_INVALID');
expect(aiChatRunService.supersede).not.toHaveBeenCalled();
});
it('supersede READY -> proceeds to stream with superseded=true', async () => {
const { controller, aiChatService } = makeController({
supersede: jest.fn().mockResolvedValue({ kind: 'ready' }),
getActiveForChat: jest.fn().mockResolvedValue(undefined), // slot free after CAS
});
const { req, res } = makeReqRes({
chatId: 'c1',
supersede: { runId: 'run-x' },
});
await controller.stream(req as never, res as never, user, wsWith(true));
expect(res.hijack).toHaveBeenCalled();
expect(aiChatService.stream).toHaveBeenCalledTimes(1);
expect(aiChatService.stream.mock.calls[0][0].superseded).toBe(true);
// The run hooks are always present now (both modes).
expect(aiChatService.stream.mock.calls[0][0].runHooks).toBeDefined();
});
it('supersede DEGRADE -> proceeds to a normal send (superseded=false)', async () => {
const { controller, aiChatService } = makeController({
supersede: jest.fn().mockResolvedValue({ kind: 'degrade' }),
});
const { req, res } = makeReqRes({
chatId: 'c1',
supersede: { runId: 'run-x' },
});
await controller.stream(req as never, res as never, user, wsWith(false));
expect(aiChatService.stream).toHaveBeenCalledTimes(1);
expect(aiChatService.stream.mock.calls[0][0].superseded).toBe(false);
});
});
@@ -418,6 +418,19 @@ export class AiChatController {
const body = (req.body ?? {}) as AiChatStreamBody;
// #487 [security]: gate cross-user access to an EXISTING chat BEFORE anything
// reads its runs. Every sibling endpoint (getRun/stop/history/rename/delete/
// attachRunStream) owner-checks the chat via assertOwnedChat; stream() must too.
// Without this a same-workspace member who is NOT the chat owner could POST a
// supersede against another user's chat and (a) harvest that user's active runId
// out of the 409 SUPERSEDE_TARGET_MISMATCH body, then (b) requestStop the foreign
// run. Gate on the chatId the client sent, when present — a brand-new chat (no
// chatId) has no prior owner to check. Mirrors /stop's owner check (403 as the
// neighbors do), and runs pre-hijack so it returns clean JSON.
if (body.chatId) {
await this.assertOwnedChat(body.chatId, user, workspace);
}
// Resolve the agent role for this turn BEFORE hijack: existing chats read it
// from ai_chats.role_id (authoritative), a new chat from body.roleId. The
// role drives both the persona and the optional model override below.
@@ -432,12 +445,66 @@ export class AiChatController {
// HttpException) instead of breaking mid-stream.
const model = await this.aiChatService.getChatModel(workspace.id, role);
// #184: one active run per chat. For an EXISTING chat reject a concurrent
// start with a clean 409 BEFORE hijack (the common double-submit / second-tab
// case), so the user gets JSON, not a mid-stream error. A brand-new chat
// (no chatId) cannot have a prior run, and the DB partial unique index is the
// backstop against any race that slips past this check.
if (autonomousRuns && body.chatId) {
// #487: server-side supersede CAS ("interrupt and send now"). When the client
// asks to replace a live run, atomically STOP it and wait for it to settle
// before this turn claims the slot. Runs BEFORE hijack so every branch returns
// clean JSON (the client keeps the composer text on a 409). See
// AiChatRunService.supersede for the branch semantics.
let superseded = false;
const supersedeRunId = body.supersede?.runId;
if (supersedeRunId) {
if (!body.chatId) {
throw new BadRequestException({
message: 'supersede requires chatId',
code: 'SUPERSEDE_INVALID',
});
}
const result = await this.aiChatRunService.supersede(
body.chatId,
supersedeRunId,
workspace.id,
);
switch (result.kind) {
case 'invalid':
throw new BadRequestException({
message: 'The run to supersede does not belong to this chat',
code: 'SUPERSEDE_INVALID',
});
case 'mismatch':
// A DIFFERENT run is active than the one the client targeted. Surface
// the CURRENT runId; the client does NOT auto-retry (a stale CAS).
throw new ConflictException({
message: 'A different agent run is now active on this chat',
code: 'SUPERSEDE_TARGET_MISMATCH',
activeRunId: result.activeRunId,
});
case 'timeout':
// The target did not settle within W — nothing was persisted, the
// composer keeps the text. NOT a rollback: the stop is already issued.
throw new ConflictException({
message:
'The previous run did not stop in time; nothing was sent — please try again',
code: 'SUPERSEDE_TIMEOUT',
});
case 'ready':
// The target stopped and settled: the slot is free. Prompt the new run
// that the old run's last operations may still be applying.
superseded = true;
break;
case 'degrade':
// The run already ended between click and POST — send normally.
break;
}
}
// #487: one active run per chat — ENFORCED IN BOTH MODES now (legacy mode used
// to have NO gate, so two tabs streamed two parallel turns on one chat, which
// interleaved history and crashed convertToModelMessages). Reject a concurrent
// start with a clean pre-hijack 409 (double-submit / second-tab). A brand-new
// chat (no chatId) cannot have a prior run, and the DB partial unique index in
// beginRun is the authoritative backstop for any race that slips past here
// (including a slot stolen between a supersede release and beginRun).
if (body.chatId) {
const active = await this.aiChatRunService.getActiveForChat(
body.chatId,
workspace.id,
@@ -446,107 +513,94 @@ export class AiChatController {
throw new ConflictException({
message: 'An agent run is already in progress for this chat',
code: 'A_RUN_ALREADY_ACTIVE',
activeRunId: active.id,
});
}
}
// Run-lifecycle hooks (#184), only when the flag is on. They wrap the turn in
// a durable run whose abort is governed by the run (explicit stop), persist
// its progress, and settle its terminal status — see AiChatRunService.
const runHooks: AiChatRunHooks | undefined = autonomousRuns
? {
begin: async (chatId) => {
const handle = await this.aiChatRunService.beginRun({
chatId,
workspaceId: workspace.id,
userId: user.id,
trigger: 'user',
});
// #184 phase 1.5: register the run-stream entry at BEGIN (before any
// frame) so a tab that attaches in the begin->seed window finds an
// entry to wait on. Gated on AI_CHAT_RESUMABLE_STREAM: with the flag
// off nothing is registered and attach always 204s.
if (
handle?.runId &&
this.environment?.isAiChatResumableStreamEnabled?.()
) {
this.streamRegistry?.open(chatId, handle.runId);
}
return handle;
},
onAssistantSeeded: (runId, messageId) =>
this.aiChatRunService.linkAssistantMessage(
runId,
workspace.id,
messageId,
),
onStep: (runId, stepCount) =>
void this.aiChatRunService.recordStep(
runId,
workspace.id,
stepCount,
),
onSettled: (runId, status, error) =>
this.aiChatRunService.finalizeRun(
runId,
workspace.id,
status,
error,
),
// #487: the turn is ALWAYS a first-class RUN now (both modes). The mode
// difference is only the abort semantics on a browser disconnect (onClose
// below). currentRunId is captured at begin so a legacy disconnect can stop
// the run through its stop lever.
let currentRunId: string | undefined;
const runHooks: AiChatRunHooks = {
begin: async (chatId) => {
const handle = await this.aiChatRunService.beginRun({
chatId,
workspaceId: workspace.id,
userId: user.id,
trigger: 'user',
});
currentRunId = handle?.runId;
// #184 phase 1.5: register the run-stream entry at BEGIN (before any
// frame) so a tab that attaches in the begin->seed window finds an entry
// to wait on. Gated on AI_CHAT_RESUMABLE_STREAM.
if (
handle?.runId &&
this.environment?.isAiChatResumableStreamEnabled?.()
) {
this.streamRegistry?.open(chatId, handle.runId);
}
: undefined;
return handle;
},
onAssistantSeeded: (runId, messageId) =>
this.aiChatRunService.linkAssistantMessage(
runId,
workspace.id,
messageId,
),
onStep: (runId, stepCount) =>
void this.aiChatRunService.recordStep(runId, workspace.id, stepCount),
onSettled: (runId, status, error) =>
this.aiChatRunService.finalizeRun(runId, workspace.id, status, error),
};
// Abort the agent loop when the client disconnects. `close` also fires on
// normal completion, so only abort when the response has not finished
// writing (a genuine disconnect). `once` fires at most once and self-removes;
// we also drop it on response `finish` so it never lingers after the stream
// completes normally (the AI SDK pipes the response fire-and-forget, so we
// cannot simply remove it once `stream()` returns).
// Handle a client disconnect. `close` also fires on normal completion, so only
// act when the response has not finished writing (a genuine disconnect). `once`
// fires at most once and self-removes; we also drop it on response `finish`.
// DIAGNOSTIC (Safari stream-drop investigation) — temporary: wall-clock at
// which a Safari disconnect is observed, measured from request receipt.
const reqStartedAt = Date.now();
const controller = new AbortController();
const onClose = (): void => {
// A genuine disconnect leaves the response unfinished (unlike a normal
// completion, which also fires `close`). Such a drop — e.g. a reverse
// proxy cutting the SSE mid-answer — is otherwise invisible server-side,
// so log it here.
if (!res.raw.writableEnded) {
if (autonomousRuns) {
// #184: the turn is a DETACHED run. A disconnect must NOT abort it —
// the run keeps executing and persisting server-side; the client
// reconnects via /ai-chat/run (or re-stops via /ai-chat/stop). Log only.
// #184: a DETACHED run — a disconnect must NOT stop it. The run keeps
// executing and persisting server-side; the client reconnects via
// /ai-chat/run (or re-stops via /ai-chat/stop). Log only.
this.logger.log(
`AI chat stream: client disconnected; run continues server-side ` +
`(elapsed=${Date.now() - reqStartedAt}ms since request received)`,
);
} else {
// #487: legacy — a disconnect ENDS the turn, but the turn is now a RUN,
// so stop it through the run's stop lever (requestStop). streamText no
// longer consumes the socket signal (effectiveSignal is the run signal),
// so aborting `controller` would do nothing; requestStop aborts the run.
this.logger.warn(
`AI chat stream: client disconnected before completion; aborting turn ` +
`(elapsed=${Date.now() - reqStartedAt}ms since request received)`,
`AI chat stream: client disconnected before completion; stopping the ` +
`run (elapsed=${Date.now() - reqStartedAt}ms since request received)`,
);
controller.abort();
if (currentRunId) {
void this.aiChatRunService.requestStop(currentRunId, workspace.id);
}
}
}
};
req.raw.once('close', onClose);
res.raw.once('finish', () => req.raw.off('close', onClose));
// #184: in detached mode the turn is NOT aborted on disconnect, so the SDK's
// pipe keeps writing to a socket the client may have dropped — for the rest of
// the (continuing) run. A write to the dead socket can emit an 'error' on the
// raw response; without a listener that surfaces as an unhandled error event.
// Swallow it (the run continues server-side regardless). Legacy mode aborts on
// disconnect, so it does not need this and keeps its exact prior behavior.
if (autonomousRuns) {
res.raw.on('error', (err) => {
this.logger.debug(
`AI chat detached stream: post-disconnect socket error swallowed: ${
err instanceof Error ? err.message : String(err)
}`,
);
});
}
// #184/#487: the run/pipe can outlive the socket in BOTH modes now (autonomous
// keeps going; legacy keeps going until requestStop's abort unwinds the turn).
// The SDK's pipe may then write to a dropped socket and emit an 'error' on the
// raw response — swallow it so it never surfaces as an unhandled error event.
res.raw.on('error', (err) => {
this.logger.debug(
`AI chat stream: post-disconnect socket error swallowed: ${
err instanceof Error ? err.message : String(err)
}`,
);
});
// Commit to streaming: hijack so Fastify stops managing the response and
// the AI SDK can write the UI-message stream directly to the Node socket.
@@ -562,8 +616,10 @@ export class AiChatController {
signal: controller.signal,
model,
role,
// #184: present only when the flag is on; wraps the turn in a durable run.
// #487: the turn is always run-wrapped now (both modes).
runHooks,
// #487: warn the new run that a superseded run's last ops may still apply.
superseded,
});
} catch (err) {
// Any failure AFTER hijack can no longer go through Nest's exception
@@ -0,0 +1,142 @@
// #489 — client-parts validation + resilient history conversion.
//
// These unit tests exercise the two exported helpers against the REAL
// `convertToModelMessages` from `ai` (NOT a mock): a genuinely malformed part
// (a `null` element inside a parts array) makes the real converter throw
// ("Cannot read properties of null"), which is the actual production
// "bricked chat" mechanism this fix defends against. Asserting against the real
// converter (rather than a mock-shaped error) is the whole point — a mock would
// hide a version change in the converter's throw behaviour.
import { convertToModelMessages, type UIMessage } from 'ai';
import {
sanitizeUserParts,
convertHistoryResilient,
TOOL_CONTEXT_OMITTED_MARKER,
} from './ai-chat.service';
type Row = Omit<UIMessage, 'id'> & { id: string };
describe('sanitizeUserParts (#489, branch: validation on receipt)', () => {
it('keeps whitelisted text parts unchanged', () => {
const drops: string[] = [];
const out = sanitizeUserParts(
[
{ type: 'text', text: 'a' },
{ type: 'text', text: 'b' },
] as UIMessage['parts'],
(t) => drops.push(t),
);
expect(out).toEqual([
{ type: 'text', text: 'a' },
{ type: 'text', text: 'b' },
]);
expect(drops).toEqual([]);
});
it('drops a non-text part (a tool-part in input-available) and reports its type', () => {
const drops: string[] = [];
const out = sanitizeUserParts(
[
{ type: 'text', text: 'hi' },
{
type: 'tool-getPage',
toolCallId: 't1',
state: 'input-available',
input: { pageId: 'p' },
},
] as unknown as UIMessage['parts'],
(t) => drops.push(t),
);
expect(out).toEqual([{ type: 'text', text: 'hi' }]);
expect(drops).toEqual(['tool-getPage']);
});
it('drops a null part (the shape that would poison convertToModelMessages)', () => {
const drops: string[] = [];
const out = sanitizeUserParts(
[{ type: 'text', text: 'hi' }, null] as unknown as UIMessage['parts'],
(t) => drops.push(t),
);
expect(out).toEqual([{ type: 'text', text: 'hi' }]);
expect(drops).toEqual(['(unknown)']);
});
it('returns undefined when nothing survives (so a null metadata is persisted)', () => {
const out = sanitizeUserParts(
[
{ type: 'tool-x', toolCallId: 't', state: 'input-available' },
] as unknown as UIMessage['parts'],
() => undefined,
);
expect(out).toBeUndefined();
});
it('returns undefined for a non-array input', () => {
expect(
sanitizeUserParts(undefined as unknown as UIMessage['parts'], () => undefined),
).toBeUndefined();
});
});
describe('convertHistoryResilient (#489, branches: happy + per-row degradation)', () => {
it('happy path: healthy history converts identically to convertToModelMessages, no degrade', async () => {
const history: Row[] = [
{ id: 'u1', role: 'user', parts: [{ type: 'text', text: 'hi' }] },
{ id: 'a1', role: 'assistant', parts: [{ type: 'text', text: 'hello' }] },
];
const degrades: number[] = [];
const out = await convertHistoryResilient(history, (i) => degrades.push(i));
const expected = await convertToModelMessages(history as UIMessage[]);
expect(out).toEqual(expected);
expect(degrades).toEqual([]);
});
it('REAL poison: a null part throws in the batch converter but is isolated and degraded to a marker', async () => {
// Sanity: the real converter genuinely throws on this shape.
const poisoned: Row = {
id: 'a1',
role: 'assistant',
parts: [
{ type: 'text', text: 'earlier answer' },
null,
] as unknown as UIMessage['parts'],
};
await expect(
convertToModelMessages([poisoned as UIMessage]),
).rejects.toThrow();
const history: Row[] = [
{ id: 'u1', role: 'user', parts: [{ type: 'text', text: 'first' }] },
poisoned,
{ id: 'u2', role: 'user', parts: [{ type: 'text', text: 'second' }] },
];
const degrades: number[] = [];
const out = await convertHistoryResilient(history, (i) => degrades.push(i));
// Only the poisoned row (index 1) is degraded.
expect(degrades).toEqual([1]);
// Healthy rows survive verbatim.
const flat = JSON.stringify(out);
expect(flat).toContain('first');
expect(flat).toContain('second');
// The degraded row carries its readable text AND the truncation marker so the
// model sees that tool context was omitted (never a silent loss).
expect(flat).toContain('earlier answer');
expect(flat).toContain(TOOL_CONTEXT_OMITTED_MARKER);
// The whole batch converted (3 model messages, none dropped).
expect(out).toHaveLength(3);
});
it('a fully-poisoned row (no readable text) still degrades to just the marker', async () => {
const history: Row[] = [
{
id: 'a1',
role: 'assistant',
parts: [null] as unknown as UIMessage['parts'],
},
];
const out = await convertHistoryResilient(history, () => undefined);
expect(out).toHaveLength(1);
expect(JSON.stringify(out)).toContain(TOOL_CONTEXT_OMITTED_MARKER);
});
});
@@ -101,6 +101,22 @@ const INTERRUPT_NOTE =
'assume your previous response was complete, and do not silently restart the ' +
'partial work — build on it or follow the new instruction.';
/**
* #487: injected on a turn started by SUPERSEDING a previous run (the user hit
* "interrupt and send now" while a run was live). The previous run was Stopped,
* but there is NO side-effect quiescence — a write it had already committed, or
* one committing at the moment of Stop, may land with a small delay AFTER this new
* run starts. So the model is told its picture of the page/state may be a beat
* stale and to re-read before assuming an edit did or did not apply.
*/
const SUPERSEDE_NOTE =
'NOTE: A previous agent run in this conversation was just interrupted so this ' +
'new turn could start. That run was stopped, but any operation it had already ' +
'begun (e.g. a page edit) may still be applied with a short delay. Do not ' +
'assume the document/state is exactly as the interrupted run left it — if you ' +
'need to rely on the current content, RE-READ it with the page tools before ' +
'acting rather than trusting a cached view.';
/**
* Injected on a turn where the open page was hand-edited by the user (or anyone
* else) AFTER the agent's previous response ended (#274). The server takes a
@@ -203,6 +219,14 @@ export interface BuildSystemPromptInput {
* (partial) answer was cut off by the user's new message.
*/
interrupted?: boolean;
/**
* #487: true when THIS turn was started by superseding a still-live previous run
* ("interrupt and send now"). Adds SUPERSEDE_NOTE so the model knows the previous
* run's last operations may still be applying and to re-read state it depends on.
* Distinct from `interrupted` (which is about a PARTIAL prior answer in history);
* both can be set together. Self-clears — set only for the superseding turn.
*/
superseded?: boolean;
/**
* Set only when the open page was edited by the user AFTER the agent's previous
* turn ended (#274), confirmed server-side by diffing the current page against
@@ -311,6 +335,7 @@ export function buildSystemPrompt({
openedPage,
mcpInstructions,
interrupted,
superseded,
pageChanged,
deferredToolsEnabled,
toolCatalog,
@@ -360,6 +385,13 @@ export function buildSystemPrompt({
context += `\n${INTERRUPT_NOTE}`;
}
// Supersede note (#487): present only for a turn that stopped and replaced a
// still-live previous run — warns the model the previous run's last operations
// may still be applying (no side-effect quiescence).
if (superseded) {
context += `\n${SUPERSEDE_NOTE}`;
}
// Per-turn page-change note (#274). Added to the context section (inside the
// safety sandwich), present only when the server detected that the open page
// was edited by the user since the agent's last turn ended. The diff content is
@@ -89,11 +89,22 @@ describe('AiChatService.stream run-lifecycle safety net (#184)', () => {
const runRepo = {
insert: jest.fn().mockResolvedValue({ id: 'run-1', status: 'running' }),
update: jest.fn().mockResolvedValue({ id: 'run-1' }),
// #487: the terminal settle now goes through the CONDITIONAL write.
finalizeIfActive: jest
.fn()
.mockResolvedValue({ id: 'run-1', status: 'failed' }),
findById: jest.fn().mockResolvedValue(undefined),
};
const runService = new AiChatRunService(runRepo as never, { isCloud: () => false } as never);
// The user-message insert (the first bare await after beginRun) throws.
// The user-message insert throws. #489 runs the history load + convert BEFORE
// the insert (convert-before-insert, so a retry cannot duplicate the user row),
// so `findAllByChat` (a real repo method) is now called first — stub it to an
// empty history so the flow reaches the insert. Both awaits are AFTER beginRun,
// so the "exception after beginRun -> settled to error" invariant is unchanged;
// the throw point simply moved from insert to a later insert after a no-op load.
const aiChatMessageRepo = {
findAllByChat: jest.fn().mockResolvedValue([]),
insert: jest.fn().mockRejectedValue(new Error('insert boom')),
};
const aiChatRepo = {
@@ -148,9 +159,10 @@ describe('AiChatService.stream run-lifecycle safety net (#184)', () => {
// The run was begun...
expect(runRepo.insert).toHaveBeenCalledTimes(1);
// ...then settled to a terminal FAILED status by the safety net...
expect(runRepo.update).toHaveBeenCalledTimes(1);
expect(runRepo.update).toHaveBeenCalledWith(
// ...then settled to a terminal FAILED status by the safety net (via the
// #487 conditional write)...
expect(runRepo.finalizeIfActive).toHaveBeenCalledTimes(1);
expect(runRepo.finalizeIfActive).toHaveBeenCalledWith(
'run-1',
'ws1',
expect.objectContaining({ status: 'failed' }),
@@ -155,6 +155,8 @@ describe('AiChatService.stream — abortSignal wiring (#184 F3)', () => {
insert: jest.fn(async () => ({ id: 'msg-1' })),
findAllByChat: jest.fn(async () => []),
update: jest.fn(async () => ({ id: 'msg-1' })),
finalizeOwner: jest.fn(async () => ({ id: 'msg-1' })),
findStreamingWithTerminalRun: jest.fn(async () => []),
};
const aiSettings = { resolve: jest.fn(async () => ({})) };
const tools = { forUser: jest.fn(async () => ({})) };
@@ -179,7 +181,7 @@ describe('AiChatService.stream — abortSignal wiring (#184 F3)', () => {
{} as never, // pageAccess
{ isAiChatDeferredToolsEnabled: () => false, isAiChatFinalStepLockdownEnabled: () => false } as never, // environment
);
return { svc };
return { svc, aiChatMessageRepo };
}
const body = {
@@ -285,7 +287,7 @@ describe('AiChatService.stream — abortSignal wiring (#184 F3)', () => {
// Drive stream() to the point streamText is called, capturing the options object
// (which carries onStepFinish/onFinish/onError/onAbort) and the run hooks.
async function captureStreamCallbacks() {
const { svc } = makeService();
const { svc, aiChatMessageRepo } = makeService();
let capturedOpts: any;
streamTextMock.mockImplementation((opts: any) => {
capturedOpts = opts;
@@ -312,7 +314,7 @@ describe('AiChatService.stream — abortSignal wiring (#184 F3)', () => {
runHooks: runHooks as never,
});
expect(capturedOpts).toBeDefined();
return { capturedOpts, runHooks };
return { capturedOpts, runHooks, aiChatMessageRepo };
}
it('F9: onStepFinish bumps the run step count, onFinish settles the run "completed" (the dominant autonomous-run path)', async () => {
@@ -332,7 +334,13 @@ describe('AiChatService.stream — abortSignal wiring (#184 F3)', () => {
usage: {},
steps: [],
});
expect(runHooks.onSettled).toHaveBeenCalledWith('run-1', 'completed');
// #487: onFinish passes the (undefined) error slot so a message-finalize
// failure could error-mark the run; on the success path it is undefined.
expect(runHooks.onSettled).toHaveBeenCalledWith(
'run-1',
'completed',
undefined,
);
});
it('F9: onAbort settles the run "aborted"', async () => {
@@ -361,6 +369,51 @@ describe('AiChatService.stream — abortSignal wiring (#184 F3)', () => {
expect.stringContaining('provider exploded'),
);
});
// #490 reactive branch: a provider CONTEXT-OVERFLOW 400 in onError is classified,
// records a distinguishable cause, and stamps metadata.replayOverflow so the NEXT
// turn's budgeter trims aggressively (the recovery that un-bricks the chat).
it('#490: a context-overflow 400 stamps replayOverflow on the finalized row', async () => {
jest
.spyOn(Logger.prototype, 'error')
.mockImplementation(() => undefined as never);
jest
.spyOn(Logger.prototype, 'warn')
.mockImplementation(() => undefined as never);
const { capturedOpts, aiChatMessageRepo } = await captureStreamCallbacks();
const overflow = Object.assign(new Error('too large'), {
statusCode: 400,
message:
"This model's maximum context length is 128000 tokens. However, your messages resulted in 214000 tokens. Please reduce the length.",
});
await capturedOpts.onError({ error: overflow });
// The seed row exists (finalizeOwner is the owner-write path).
expect(aiChatMessageRepo.finalizeOwner).toHaveBeenCalled();
const calls = aiChatMessageRepo.finalizeOwner.mock.calls as any[][];
const patch = calls[calls.length - 1][2] as {
status: string;
metadata: Record<string, unknown>;
};
expect(patch.status).toBe('error');
expect(patch.metadata.replayOverflow).toBe(true);
expect(patch.metadata.error).toContain('контекстное окно');
});
it('#490: a non-overflow error does NOT stamp replayOverflow', async () => {
jest
.spyOn(Logger.prototype, 'error')
.mockImplementation(() => undefined as never);
const { capturedOpts, aiChatMessageRepo } = await captureStreamCallbacks();
await capturedOpts.onError({ error: new Error('network reset') });
const calls = aiChatMessageRepo.finalizeOwner.mock.calls as any[][];
const patch = calls[calls.length - 1][2] as {
status: string;
metadata: Record<string, unknown>;
};
expect('replayOverflow' in patch.metadata).toBe(false);
});
});
/**
@@ -415,6 +468,8 @@ describe('AiChatService.stream — begin-failure fails the turn (#184 F14 / #486
insert: jest.fn(async () => ({ id: 'msg-1' })),
findAllByChat: jest.fn(async () => []),
update: jest.fn(async () => ({ id: 'msg-1' })),
finalizeOwner: jest.fn(async () => ({ id: 'msg-1' })),
findStreamingWithTerminalRun: jest.fn(async () => []),
};
const aiSettings = { resolve: jest.fn(async () => ({})) };
const tools = { forUser: jest.fn(async () => ({})) };
@@ -13,6 +13,7 @@ import {
compactToolOutput,
assistantParts,
serializeSteps,
type StepPartsCache,
rowToUiMessage,
prepareAgentStep,
stepBudgetWarning,
@@ -28,6 +29,9 @@ import {
FINAL_STEP_NUDGE,
STEP_LIMIT_NO_ANSWER_MARKER,
OUTPUT_DEGENERATION_ERROR,
lastAssistantContextTokens,
lastAssistantReplayOverflow,
seedActivatedTools,
} from './ai-chat.service';
import type { AiChatMessage, Workspace } from '@docmost/db/types/entity.types';
import { buildSystemPrompt } from './ai-chat.prompt';
@@ -114,6 +118,54 @@ describe('compactToolOutput', () => {
describe('assistantParts', () => {
type AnyPart = Record<string, unknown>;
// #490 memoization: assistantParts builds each step's parts once and caches
// them by the step OBJECT's identity, so a mid-stream flush does not
// re-stringify every prior step's (large) output. Observable property: with a
// shared cache, the second call over the SAME step object returns the cached
// (identical) part array even if the step's underlying output was swapped —
// proving the work was memoized, not redone.
it('memoizes a step by identity (shared cache => one build per step)', () => {
const cache: StepPartsCache = new WeakMap();
const step = {
text: 'x',
toolCalls: [{ toolCallId: 'c1', toolName: 'getPage', input: {} }],
toolResults: [{ toolCallId: 'c1', toolName: 'getPage', output: { v: 1 } }],
};
const first = assistantParts([step], '', cache) as AnyPart[];
expect((first.find((p) => p.type === 'tool-getPage')!.output as any).v).toBe(
1,
);
// Swap the output for a NEW value; a re-build would pick it up, a cache hit
// keeps the first result.
step.toolResults[0] = {
toolCallId: 'c1',
toolName: 'getPage',
output: { v: 2 },
};
const second = assistantParts([step], '', cache) as AnyPart[];
expect((second.find((p) => p.type === 'tool-getPage')!.output as any).v).toBe(
1,
);
// Same cached part objects are reused.
expect(second.find((p) => p.type === 'tool-getPage')).toBe(
first.find((p) => p.type === 'tool-getPage'),
);
});
it('without a cache, each call rebuilds (no stale memo)', () => {
const step = {
text: 'x',
toolCalls: [{ toolCallId: 'c1', toolName: 'getPage', input: {} }],
toolResults: [{ toolCallId: 'c1', toolName: 'getPage', output: { v: 1 } }],
};
const first = assistantParts([step], '') as AnyPart[];
step.toolResults[0].output = { v: 2 };
const second = assistantParts([step], '') as AnyPart[];
expect((second.find((p) => p.type === 'tool-getPage')!.output as any).v).toBe(
2,
);
});
it('emits output-available for a tool-call WITH a paired result', () => {
const steps = [
{
@@ -231,61 +283,299 @@ describe('assistantParts', () => {
});
});
describe('serializeSteps', () => {
// #490 trace format v2: per call the trace stores { input } for the call and an
// OUTCOME element — { ok: true } on success, { error, kind: 'thrown' } on a
// thrown tool-error, { error, kind: 'interrupted' } on a mid-step abort. The tool
// OUTPUT is no longer duplicated here (it lives once in metadata.parts).
describe('serializeSteps (trace v2)', () => {
it('returns null when there are no calls or results', () => {
expect(serializeSteps([])).toBeNull();
});
it('flattens calls and results into a compact trace', () => {
it('pairs a successful call with an { ok: true } outcome and NO output', () => {
const trace = serializeSteps([
{
toolCalls: [{ toolName: 'getPage', input: { id: 'p1' } }],
toolResults: [{ toolName: 'getPage', output: { title: 'T' } }],
toolCalls: [{ toolCallId: 'c1', toolName: 'getPage', input: { id: 'p1' } }],
toolResults: [{ toolCallId: 'c1', toolName: 'getPage' }],
},
]) as Array<Record<string, unknown>>;
expect(trace).toHaveLength(2);
expect(trace[0]).toEqual({ toolName: 'getPage', input: { id: 'p1' } });
expect(trace[1]).toEqual({ toolName: 'getPage', output: { title: 'T' } });
expect(trace[1]).toEqual({ toolName: 'getPage', ok: true });
// The output is NOT stored in the trace any more (dedup: it lives in parts).
expect(trace.some((e) => 'output' in e)).toBe(false);
});
it('records a THROWN tool failure (tool-error part) with its error message', () => {
it('records a THROWN failure with { error, kind: "thrown" }', () => {
const trace = serializeSteps([
{
toolCalls: [{ toolName: 'editPageText', input: { id: 'p1' } }],
toolCalls: [
{ toolCallId: 'c1', toolName: 'editPageText', input: { id: 'p1' } },
],
toolResults: [],
content: [
{
type: 'tool-error',
toolCallId: 'c1',
toolName: 'editPageText',
error: new Error('page is locked'),
},
],
},
]) as Array<Record<string, unknown>>;
// The call element is followed by a paired error element (mirroring how a
// successful result is appended), so the failure survives in the trace.
expect(trace).toHaveLength(2);
expect(trace[0]).toEqual({ toolName: 'editPageText', input: { id: 'p1' } });
expect(trace[1]).toEqual({
toolName: 'editPageText',
error: 'page is locked',
kind: 'thrown',
});
});
it('truncates a very long tool-error message to the tool-output limit', () => {
it('marks an interrupted call (no result, no throw) with kind "interrupted"', () => {
const trace = serializeSteps([
{
toolCalls: [
{ toolCallId: 'c1', toolName: 'createComment', input: { x: 1 } },
],
toolResults: [],
content: [],
},
]) as Array<Record<string, unknown>>;
expect(trace).toHaveLength(2);
expect(trace[1]).toEqual({
toolName: 'createComment',
error: 'Tool call did not complete.',
kind: 'interrupted',
});
// Structurally distinct from a thrown hard-fail so it never inflates an
// error-rate scan.
expect((trace[1] as { kind: string }).kind).not.toBe('thrown');
});
it('truncates a very long thrown-error message to the tool-output limit', () => {
const long = 'x'.repeat(5000);
const trace = serializeSteps([
{
toolCalls: [{ toolName: 'editPageText', input: {} }],
toolCalls: [{ toolCallId: 'c1', toolName: 'editPageText', input: {} }],
toolResults: [],
content: [{ type: 'tool-error', toolName: 'editPageText', error: long }],
content: [
{
type: 'tool-error',
toolCallId: 'c1',
toolName: 'editPageText',
error: long,
},
],
},
]) as Array<Record<string, unknown>>;
const errorText = trace[1].error as string;
// Truncated (not the full 5000 chars) and carries the omission marker.
expect(errorText.length).toBeLessThan(long.length);
expect(errorText).toContain('chars omitted');
});
it('pairs parallel calls in one step with their outcomes by id', () => {
const trace = serializeSteps([
{
toolCalls: [
{ toolCallId: 'a', toolName: 'getPage', input: {} },
{ toolCallId: 'b', toolName: 'searchPages', input: {} },
],
toolResults: [{ toolCallId: 'b', toolName: 'searchPages' }],
content: [
{ type: 'tool-error', toolCallId: 'a', toolName: 'getPage', error: 'nope' },
],
},
]) as Array<Record<string, unknown>>;
// call a, outcome a (thrown), call b, outcome b (ok)
expect(trace).toHaveLength(4);
expect(trace[1]).toEqual({ toolName: 'getPage', error: 'nope', kind: 'thrown' });
expect(trace[3]).toEqual({ toolName: 'searchPages', ok: true });
});
});
// #490: every assistant row flushAssistant writes carries the v2 era marker so a
// dual-shape diagnostic query can branch on the trace shape without inspecting it.
describe('toolTraceVersion era marker (#490)', () => {
it('stamps metadata.toolTraceVersion = 2 on every flushed row', () => {
const seed = flushAssistant([], '', 'streaming');
expect(seed.metadata.toolTraceVersion).toBe(2);
const done = flushAssistant(
[
{
text: 'ok',
toolCalls: [{ toolCallId: 'c1', toolName: 'getPage', input: {} }],
toolResults: [{ toolCallId: 'c1', toolName: 'getPage' }],
},
],
'',
'completed',
{ finishReason: 'stop' },
);
expect(done.metadata.toolTraceVersion).toBe(2);
});
});
// #490 replay-budget signal helpers over persisted history.
describe('lastAssistantContextTokens', () => {
const row = (
role: string,
metadata: Record<string, unknown> | null,
): AiChatMessage => ({ role, metadata }) as unknown as AiChatMessage;
it('reads the most recent assistant turn contextTokens (provider fact)', () => {
const hist = [
row('user', null),
row('assistant', { contextTokens: 12000 }),
row('user', null),
row('assistant', { contextTokens: 41000 }),
];
expect(lastAssistantContextTokens(hist)).toBe(41000);
});
it('returns undefined when the last assistant turn recorded no usage', () => {
const hist = [row('assistant', { error: 'boom' }), row('user', null)];
expect(lastAssistantContextTokens(hist)).toBeUndefined();
expect(lastAssistantContextTokens([])).toBeUndefined();
});
});
// #490 snapshotOpenPage fast-path: skip the full Markdown export + upsert when a
// snapshot already exists at the page's CURRENT version (same updated_at instant).
describe('snapshotOpenPage fast-path (#490)', () => {
function makeSvc(existingSnapshot: unknown, pageUpdatedAt: Date) {
const exportPageMarkdown = jest.fn(async () => '# md');
const upsert = jest.fn(async () => undefined);
const findByChatPage = jest.fn(async () => existingSnapshot);
const pageRepo = {
findById: jest.fn(async () => ({
id: 'p1',
workspaceId: 'ws1',
updatedAt: pageUpdatedAt,
})),
};
const svc = new AiChatService(
{} as never, // ai
{} as never, // aiChatRepo
{} as never, // aiChatMessageRepo
{ findByChatPage, upsert } as never, // aiChatPageSnapshotRepo
{} as never, // aiSettings
{ exportPageMarkdown } as never, // tools
{} as never, // mcpClients
{} as never, // aiAgentRoleRepo
pageRepo as never, // pageRepo
{} as never, // pageAccess
{} as never, // environment
);
return { svc, exportPageMarkdown, upsert, findByChatPage };
}
const args = () =>
[
'chat1',
'p1',
{ id: 'ws1' } as never,
{ id: 'u1' } as never,
'sess',
] as const;
it('skips export + upsert when the snapshot is already at this page version', async () => {
const t = new Date('2026-07-07T10:00:00Z');
const { svc, exportPageMarkdown, upsert } = makeSvc(
{ pageUpdatedAt: t, contentMd: '# md' },
t,
);
await (svc as unknown as { snapshotOpenPage: (...a: unknown[]) => Promise<void> })
.snapshotOpenPage(...args());
expect(exportPageMarkdown).not.toHaveBeenCalled();
expect(upsert).not.toHaveBeenCalled();
});
it('exports + upserts when the page advanced since the snapshot', async () => {
const { svc, exportPageMarkdown, upsert } = makeSvc(
{ pageUpdatedAt: new Date('2026-07-07T10:00:00Z'), contentMd: 'old' },
new Date('2026-07-07T11:00:00Z'),
);
await (svc as unknown as { snapshotOpenPage: (...a: unknown[]) => Promise<void> })
.snapshotOpenPage(...args());
expect(exportPageMarkdown).toHaveBeenCalledTimes(1);
expect(upsert).toHaveBeenCalledTimes(1);
});
it('seeds (exports + upserts) on the first turn (no snapshot yet)', async () => {
const { svc, exportPageMarkdown, upsert } = makeSvc(
undefined,
new Date('2026-07-07T10:00:00Z'),
);
await (svc as unknown as { snapshotOpenPage: (...a: unknown[]) => Promise<void> })
.snapshotOpenPage(...args());
expect(exportPageMarkdown).toHaveBeenCalledTimes(1);
expect(upsert).toHaveBeenCalledTimes(1);
});
});
// #490 deferred-tool activation persisted across turns.
describe('seedActivatedTools', () => {
const valid = new Set(['Search_web', 'getPageJson', 'diffPageVersions']);
it('seeds from persisted metadata, intersected with current valid names', () => {
expect(
seedActivatedTools(
{ activatedTools: ['Search_web', 'getPageJson'] },
valid,
),
).toEqual(['Search_web', 'getPageJson']);
});
it('drops a stored tool that is no longer valid (allowlist/role changed)', () => {
// 'Habr_publish' was activated before but is not in the current allowlist.
expect(
seedActivatedTools({ activatedTools: ['Search_web', 'Habr_publish'] }, valid),
).toEqual(['Search_web']);
});
it('is empty/robust for missing, non-array, or unknown-shaped metadata', () => {
expect(seedActivatedTools(undefined, valid)).toEqual([]);
expect(seedActivatedTools({}, valid)).toEqual([]);
expect(seedActivatedTools({ activatedTools: 'nope' }, valid)).toEqual([]);
expect(
seedActivatedTools({ activatedTools: [1, 'getPageJson', null] }, valid),
).toEqual(['getPageJson']);
});
it('de-duplicates stored names', () => {
expect(
seedActivatedTools(
{ activatedTools: ['getPageJson', 'getPageJson'] },
valid,
),
).toEqual(['getPageJson']);
});
});
describe('lastAssistantReplayOverflow', () => {
const row = (
role: string,
metadata: Record<string, unknown> | null,
): AiChatMessage => ({ role, metadata }) as unknown as AiChatMessage;
it('is true only when the LAST assistant turn overflowed', () => {
expect(
lastAssistantReplayOverflow([
row('assistant', { replayOverflow: true }),
row('user', null),
]),
).toBe(true);
// A recovered (later, non-overflow) assistant turn clears it.
expect(
lastAssistantReplayOverflow([
row('assistant', { replayOverflow: true }),
row('user', null),
row('assistant', { contextTokens: 5 }),
]),
).toBe(false);
expect(lastAssistantReplayOverflow([])).toBe(false);
});
});
describe('rowToUiMessage', () => {
@@ -618,6 +908,23 @@ describe('flushAssistant', () => {
expect(flushed.metadata.error).toBe('boom');
});
// #490 observability: the replay budgeter's decision is stamped on the turn.
it('records replayTrimmedToTokens + replayOverflow when provided', () => {
const f = flushAssistant([], '', 'error', {
error: 'ctx',
replayTrimmedToTokens: 42_000,
replayOverflow: true,
});
expect(f.metadata.replayTrimmedToTokens).toBe(42_000);
expect(f.metadata.replayOverflow).toBe(true);
});
it('omits the replay metadata when not provided', () => {
const f = flushAssistant([], '', 'completed', { finishReason: 'stop' });
expect('replayTrimmedToTokens' in f.metadata).toBe(false);
expect('replayOverflow' in f.metadata).toBe(false);
});
// #274 observability: the page-change diff the agent saw this turn is persisted
// to metadata.pageChanged when a non-empty diff was injected, and omitted when
// the diff is empty/whitespace or the arg is not supplied.
@@ -1315,8 +1622,12 @@ describe('AiChatService page-change lifecycle (#274)', () => {
describe('isInterruptResume', () => {
// history tail is the just-inserted user row; [len-2] is the previous turn.
const withPrev = (
prev: { role: string; status?: string | null } | null,
): Array<{ role: string; status?: string | null }> =>
prev: {
role: string;
status?: string | null;
metadata?: unknown;
} | null,
): Array<{ role: string; status?: string | null; metadata?: unknown }> =>
prev
? [prev, { role: 'user', status: null }]
: [{ role: 'user', status: null }];
@@ -1357,6 +1668,33 @@ describe('isInterruptResume', () => {
it('false when there is no preceding turn (only the new user row)', () => {
expect(isInterruptResume(withPrev(null), true)).toBe(false);
});
it('#487 EXCLUDES a reconcile stamp (finalizeFailed) — not a genuine interruption', () => {
// A row a reconcile settled to 'aborted' carries metadata.finalizeFailed. It
// must NOT be treated as an interrupt-resume (that would inject a false
// "you were interrupted" note), even though its status is 'aborted'.
expect(
isInterruptResume(
withPrev({
role: 'assistant',
status: 'aborted',
metadata: { finalizeFailed: true },
}),
true,
),
).toBe(false);
// A genuine abort (no finalizeFailed) still counts.
expect(
isInterruptResume(
withPrev({
role: 'assistant',
status: 'aborted',
metadata: { parts: [] },
}),
true,
),
).toBe(true);
});
});
/**
@@ -1409,7 +1747,7 @@ describe('AiChatService.stream — resumable pipe options (#184 phase 1.5)', ()
}
// Wire only the deps reached on the way to the pipe call, plus a spy registry.
function makeService(opts: { resumable: boolean }) {
function makeService(opts: { resumable: boolean; history?: unknown[] }) {
const aiChatRepo = {
findById: jest.fn(async () => ({ id: 'chat-1', workspaceId: 'ws-1' })),
insert: jest.fn(),
@@ -1417,8 +1755,11 @@ describe('AiChatService.stream — resumable pipe options (#184 phase 1.5)', ()
const aiChatMessageRepo = {
// Both the user insert and the assistant seed return the same row id.
insert: jest.fn(async () => ({ id: 'msg-1' })),
findAllByChat: jest.fn(async () => []),
findAllByChat: jest.fn(async () => opts.history ?? []),
update: jest.fn(async () => ({ id: 'msg-1' })),
// #487: the terminal owner-write + the opportunistic reconcile query.
finalizeOwner: jest.fn(async () => ({ id: 'msg-1' })),
findStreamingWithTerminalRun: jest.fn(async () => []),
};
const aiSettings = { resolve: jest.fn(async () => ({})) };
const tools = { forUser: jest.fn(async () => ({})) };
@@ -1453,7 +1794,7 @@ describe('AiChatService.stream — resumable pipe options (#184 phase 1.5)', ()
} as never,
streamRegistry as never,
);
return { svc, streamRegistry };
return { svc, streamRegistry, aiChatMessageRepo };
}
const body = {
@@ -1536,6 +1877,86 @@ describe('AiChatService.stream — resumable pipe options (#184 phase 1.5)', ()
await expect(drive(svc, makeRunHooks())).rejects.toThrow('boom');
expect(streamRegistry.abortEntry).toHaveBeenCalledWith('chat-1', 'run-1');
});
// #489 REGRESSION (against the REAL convertToModelMessages — not mocked here):
// a persisted history row whose parts contain a `null` element makes the real
// convertToModelMessages THROW ("Cannot read properties of null"). Pre-fix that
// 500-ed every turn forever and each retry appended a duplicate user row. The
// fix converts BEFORE the insert and isolates the poisoned row per-row, degrading
// it to text with a "[tool context omitted]" marker. Assert the turn still runs,
// the marker reaches the model, and exactly ONE user row is inserted.
it('#489: a poisoned OLD-history row keeps the chat working; the marker reaches the model; one user insert', async () => {
const { svc, aiChatMessageRepo } = makeService({
resumable: false,
history: [
{
id: 'old-1',
role: 'assistant',
content: 'earlier answer',
// A null part is the poison: rowToUiMessage keeps it (the array is
// non-empty) and the real convertToModelMessages throws on it.
metadata: { parts: [{ type: 'text', text: 'earlier answer' }, null] },
status: 'completed',
},
],
});
// Must NOT throw — the poisoned row is degraded, not fatal.
await drive(svc, makeRunHooks());
expect(streamTextMock).toHaveBeenCalledTimes(1);
const passedMessages = streamTextMock.mock.calls[0][0].messages;
const serialized = JSON.stringify(passedMessages);
// The model sees the truncation marker (silent tool-context loss is not ok)
// AND the row's readable text is preserved alongside it.
expect(serialized).toContain('[tool context omitted]');
expect(serialized).toContain('earlier answer');
// Exactly ONE user row inserted (no duplicate), inserted AFTER conversion.
const userInserts = aiChatMessageRepo.insert.mock.calls
.map((c: unknown[]) => c[0] as { role?: string })
.filter((r) => r.role === 'user');
expect(userInserts).toHaveLength(1);
});
// #489: client-supplied non-text parts (a tool-part in `input-available`, the
// exact "bricking" payload) are dropped ON RECEIPT — never persisted — so they
// can never poison future turns. Only the text survives into metadata.parts.
it('#489: a non-text client part is stripped before persist (only text survives)', async () => {
const { svc, aiChatMessageRepo } = makeService({ resumable: false });
await svc.stream({
user: { id: 'u1' } as never,
workspace: { id: 'ws-1' } as never,
sessionId: 's1',
body: {
chatId: 'chat-1',
messages: [
{
id: 'm1',
role: 'user',
parts: [
{ type: 'text', text: 'hello' },
// untrusted tool-part — must be dropped, never persisted
{
type: 'tool-getPage',
toolCallId: 't1',
state: 'input-available',
input: { pageId: 'p' },
},
],
},
],
} as never,
res: makeRes() as never,
signal: new AbortController().signal,
model: {} as never,
role: null,
runHooks: makeRunHooks() as never,
});
const userInsert = aiChatMessageRepo.insert.mock.calls
.map((c: unknown[]) => c[0] as { role?: string; metadata?: unknown })
.find((r) => r.role === 'user');
const parts = (userInsert?.metadata as { parts?: Array<{ type: string }> })
?.parts;
expect(parts).toEqual([{ type: 'text', text: 'hello' }]);
});
});
/**
@@ -1623,6 +2044,19 @@ describe('AiChatService.stream — token-degeneration reaction (#444)', () => {
return { id };
},
),
// #487: the terminal owner-write records into the SAME `updated` recorder so
// assertions on the terminal 'completed'/'error'/'aborted' write still hold.
finalizeOwner: jest.fn(
async (
id: string,
workspaceId: string,
patch: Record<string, unknown>,
) => {
updated.push({ id, workspaceId, patch });
return { id };
},
),
findStreamingWithTerminalRun: jest.fn(async () => []),
};
const aiSettings = { resolve: jest.fn(async () => ({})) };
const tools = { forUser: jest.fn(async () => ({})) };
@@ -1882,3 +2316,148 @@ describe('AiChatService.stream — token-degeneration reaction (#444)', () => {
expect(patch.content).not.toContain(STEP_LIMIT_NO_ANSWER_MARKER);
});
});
// #487 F3 — the reconcile() / reconcileChat() ORCHESTRATORS. The individual
// clauses are exercised elsewhere; these pin the production orchestration the
// per-clause specs do not: the clause ORDER, the per-clause try/catch ISOLATION
// (one clause throwing must NOT abort the others), and reconcileChat() (which runs
// at the start of every turn and was entirely uncovered).
describe('AiChatService.reconcile / reconcileChat orchestrators (#487 F3)', () => {
let warnSpy: jest.SpyInstance;
beforeEach(() => {
// Silence the intentional clause-failure warnings (kept out of test output).
warnSpy = jest
.spyOn(Logger.prototype, 'warn')
.mockImplementation(() => undefined);
});
afterEach(() => {
warnSpy.mockRestore();
});
function makeService(opts: {
messageRepo?: Record<string, jest.Mock>;
runService?: Record<string, jest.Mock>;
}) {
const aiChatMessageRepo = {
findStreamingWithTerminalRun: jest.fn(async () => []),
stampTerminalIfStreaming: jest.fn(async () => undefined),
sweepStreamingWithoutActiveRun: jest.fn(async () => 0),
...(opts.messageRepo ?? {}),
};
const aiChatRunService = opts.runService
? {
zombieRunIds: jest.fn(() => []),
settleZombie: jest.fn(async () => true),
reconcileStaleRuns: jest.fn(async () => 0),
...opts.runService,
}
: undefined;
const svc = new AiChatService(
{} as never, // ai
{} as never, // aiChatRepo
aiChatMessageRepo as never,
{} as never, // aiChatPageSnapshotRepo
{} as never, // aiSettings
{} as never, // tools
{} as never, // mcpClients
{} as never, // aiAgentRoleRepo
{} as never, // pageRepo
{} as never, // pageAccess
{} as never, // environment
{} as never, // streamRegistry
aiChatRunService as never, // aiChatRunService (#487)
);
return { svc, aiChatMessageRepo, aiChatRunService };
}
it('reconcile() fires all four clauses IN ORDER (a -> b -> c -> d)', async () => {
const order: string[] = [];
const { svc } = makeService({
messageRepo: {
findStreamingWithTerminalRun: jest.fn(async () => {
order.push('b:find');
return [
{ messageId: 'm1', workspaceId: 'ws1', runStatus: 'succeeded' },
];
}),
stampTerminalIfStreaming: jest.fn(async () => {
order.push('b:stamp');
}),
sweepStreamingWithoutActiveRun: jest.fn(async () => {
order.push('d');
return 0;
}),
},
runService: {
zombieRunIds: jest.fn(() => ['z1']),
settleZombie: jest.fn(async () => {
order.push('a');
return true;
}),
reconcileStaleRuns: jest.fn(async () => {
order.push('c');
return 0;
}),
},
});
await svc.reconcile();
expect(order).toEqual(['a', 'b:find', 'b:stamp', 'c', 'd']);
});
it('a clause that THROWS does not abort the remaining clauses (per-clause try/catch isolation)', async () => {
const { svc, aiChatMessageRepo, aiChatRunService } = makeService({
messageRepo: {
// Clause (b) blows up mid-reconcile.
findStreamingWithTerminalRun: jest.fn(async () => {
throw new Error('clause b DB blip');
}),
},
runService: {
zombieRunIds: jest.fn(() => ['z1']),
},
});
// reconcile() must SETTLE (the clause-b failure is swallowed), not reject.
await expect(svc.reconcile()).resolves.toBeUndefined();
// (a) ran before (b); crucially (c) and (d) STILL ran despite (b) throwing —
// the property a missing try/catch would break. MUTATION-VERIFY: drop clause
// (b)'s try/catch and this reddens (the throw propagates, skipping c + d).
expect(aiChatRunService!.settleZombie).toHaveBeenCalled(); // (a)
expect(aiChatRunService!.reconcileStaleRuns).toHaveBeenCalled(); // (c)
expect(
aiChatMessageRepo.sweepStreamingWithoutActiveRun,
).toHaveBeenCalled(); // (d)
});
it('reconcileChat() settles THIS chat\'s stuck streaming rows by their run status', async () => {
const { svc, aiChatMessageRepo } = makeService({
messageRepo: {
findStreamingWithTerminalRun: jest.fn(async () => [
{ messageId: 'm1', workspaceId: 'ws1', runStatus: 'failed' },
{ messageId: 'm2', workspaceId: 'ws1', runStatus: 'succeeded' },
]),
},
});
await svc.reconcileChat('chat-1', 'ws1');
// Scoped to THIS chat and bounded at 50 (the user-facing opportunistic path).
expect(
aiChatMessageRepo.findStreamingWithTerminalRun,
).toHaveBeenCalledWith(50, { chatId: 'chat-1', workspaceId: 'ws1' });
// failed-run -> 'error'; every other terminal status -> 'aborted'.
expect(aiChatMessageRepo.stampTerminalIfStreaming).toHaveBeenCalledWith(
'm1',
'ws1',
'error',
);
expect(aiChatMessageRepo.stampTerminalIfStreaming).toHaveBeenCalledWith(
'm2',
'ws1',
'aborted',
);
});
});
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,209 @@
import { randomBytes } from 'crypto';
import { Client } from 'pg';
import { flushAssistant, serializeSteps } from './ai-chat.service';
/**
* #490 write-volume regression — an OBSERVABLE-PROPERTY test on a LIVE Postgres,
* not "bytes through a mock repo" (a mock measures exactly the thing that does not
* hurt). It drives a realistic 50-step run where each step returns a ~100 KB tool
* output and, at every `onStepFinish`, UPDATEs the assistant row the way the
* service does — then reads the REAL write volume via the `pg_current_wal_lsn()`
* delta around the run.
*
* The property proven: v2 stores each tool OUTPUT only in `metadata.parts`, no
* longer ALSO in the `tool_calls` trace. So:
* 1. the trace (`tool_calls`) column's write volume is now O(Σ steps) — tiny,
* linear outcome flags — vs the pre-#490 O(N²) that re-persisted every prior
* output on every step; and
* 2. the FULL-row write volume drops sharply (the duplicated output copy is gone).
*
* Connects to the local gitmost test Postgres (docker `gitmost-test-pg` on :5432);
* SKIPS cleanly when that DB is not reachable so it never breaks a DB-less CI.
*/
const CONN =
process.env.WAL_TEST_DATABASE_URL ??
'postgresql://docmost:docmost_dev_pw@localhost:5432/docmost';
// A step whose tool output is ~100 KB (a page read), in the SDK StepLike shape.
// The body is INCOMPRESSIBLE random text — a `'x'.repeat()` filler would TOAST-
// compress to nothing and hide the real write volume (a page body does not).
function makeStep(i: number, outputBytes = 100_000) {
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 },
},
],
};
}
// The pre-#490 (v1) trace: outputs stored a SECOND time in `tool_calls`
// (the duplication #490 removed). Mirrors the OLD serializeSteps shape.
function v1Trace(steps: ReturnType<typeof makeStep>[]): unknown {
const calls: unknown[] = [];
for (const s of steps) {
for (const c of s.toolCalls) calls.push({ toolName: c.toolName, input: c.input });
for (const r of s.toolResults)
calls.push({ toolName: r.toolName, output: r.output });
}
return calls;
}
async function walDelta(
client: Client,
fn: () => Promise<void>,
): Promise<number> {
const before = (await client.query('SELECT pg_current_wal_lsn() AS l')).rows[0]
.l as string;
await fn();
// NOTE: do NOT pg_switch_wal() here — a segment switch pads the LSN to the next
// 16 MB boundary and would swamp the actual write delta. The raw LSN advances by
// the bytes of WAL emitted, which is exactly what we want to measure.
const after = (await client.query('SELECT pg_current_wal_lsn() AS l')).rows[0]
.l as string;
return Number(
(await client.query('SELECT pg_wal_lsn_diff($1,$2) AS d', [after, before]))
.rows[0].d,
);
}
describe('#490 write-volume on a live Postgres (pg_current_wal_lsn delta)', () => {
let client: Client | undefined;
let available = false;
beforeAll(async () => {
try {
client = new Client(CONN);
await client.connect();
await client.query('SELECT pg_current_wal_lsn()');
available = true;
} catch {
available = false;
client = undefined;
}
});
afterAll(async () => {
await client?.end().catch(() => undefined);
});
const STEPS = 50;
it('v2 trace write volume is O(Σ steps) — a tiny fraction of the v1 duplicate', async () => {
if (!available || !client) {
console.warn('SKIP: gitmost-test-pg not reachable; skipping WAL test.');
return;
}
const c = client;
// Isolated table so we measure only the tool_calls (trace) column's writes.
await c.query('DROP TABLE IF EXISTS _wal_trace');
await c.query('CREATE TABLE _wal_trace(id int primary key, tool_calls jsonb)');
await c.query("INSERT INTO _wal_trace VALUES (1, '[]'::jsonb)");
const steps: ReturnType<typeof makeStep>[] = [];
// v1: each step re-persists ALL prior outputs into the trace (the O(N²) churn).
const v1 = await walDelta(c, async () => {
const acc: ReturnType<typeof makeStep>[] = [];
for (let i = 0; i < STEPS; i++) {
acc.push(makeStep(i));
await c.query('UPDATE _wal_trace SET tool_calls=$1 WHERE id=1', [
JSON.stringify(v1Trace(acc)),
]);
}
steps.push(...acc);
});
await c.query("UPDATE _wal_trace SET tool_calls='[]'::jsonb WHERE id=1");
// v2: the REAL serializeSteps — outcome flags only, NO outputs.
const v2 = await walDelta(c, async () => {
const acc: ReturnType<typeof makeStep>[] = [];
for (let i = 0; i < STEPS; i++) {
acc.push(makeStep(i));
await c.query('UPDATE _wal_trace SET tool_calls=$1 WHERE id=1', [
JSON.stringify(serializeSteps(acc)),
]);
}
});
await c.query('DROP TABLE IF EXISTS _wal_trace');
// eslint-disable-next-line no-console
console.log(
`[#490 WAL] trace column over ${STEPS} steps: v1=${(v1 / 1e6).toFixed(1)}MB ` +
`v2=${(v2 / 1e6).toFixed(2)}MB (${(v1 / v2).toFixed(0)}x smaller)`,
);
// The trace no longer carries outputs: v2 is a tiny fraction of v1's WAL.
expect(v2).toBeLessThan(v1 * 0.1);
// And v2's trace WAL is small in absolute terms — O(Σ steps) of flags, not
// O(N² × output). 50 steps of ~40-byte flags is well under a few MB of WAL.
expect(v2).toBeLessThan(5_000_000);
// v1's duplicate alone is huge (≈ the 100 KB output re-written N² times).
expect(v1).toBeGreaterThan(50_000_000);
}, 120_000);
it('the full assistant row write drops sharply once the duplicate is gone', async () => {
if (!available || !client) return;
const c = client;
await c.query('DROP TABLE IF EXISTS _wal_full');
await c.query(
'CREATE TABLE _wal_full(id int primary key, content text, tool_calls jsonb, metadata jsonb, status text)',
);
await c.query("INSERT INTO _wal_full VALUES (1, '', '[]'::jsonb, '{}'::jsonb, 'streaming')");
const writeRow = async (patch: {
content: string;
toolCalls: unknown;
metadata: unknown;
status: string;
}) =>
c.query(
'UPDATE _wal_full SET content=$1, tool_calls=$2, metadata=$3, status=$4 WHERE id=1',
[
patch.content,
JSON.stringify(patch.toolCalls ?? null),
JSON.stringify(patch.metadata),
patch.status,
],
);
// v2 (real flushAssistant): outputs live once, in metadata.parts.
const v2 = await walDelta(c, async () => {
const acc: ReturnType<typeof makeStep>[] = [];
for (let i = 0; i < STEPS; i++) {
acc.push(makeStep(i));
await writeRow(flushAssistant(acc as never, '', 'streaming'));
}
});
await c.query("UPDATE _wal_full SET content='', tool_calls='[]'::jsonb, metadata='{}'::jsonb WHERE id=1");
// v1: same row PLUS the duplicated outputs in the trace column.
const v1 = await walDelta(c, async () => {
const acc: ReturnType<typeof makeStep>[] = [];
for (let i = 0; i < STEPS; i++) {
acc.push(makeStep(i));
const f = flushAssistant(acc as never, '', 'streaming');
await writeRow({ ...f, toolCalls: v1Trace(acc) });
}
});
await c.query('DROP TABLE IF EXISTS _wal_full');
// eslint-disable-next-line no-console
console.log(
`[#490 WAL] full row over ${STEPS} steps: v1=${(v1 / 1e6).toFixed(1)}MB ` +
`v2=${(v2 / 1e6).toFixed(1)}MB (saved ${((1 - v2 / v1) * 100).toFixed(0)}%)`,
);
// Removing the duplicated trace copy is a large, real write-volume reduction.
expect(v2).toBeLessThan(v1 * 0.75);
}, 120_000);
});
@@ -0,0 +1,261 @@
import { errors } from 'undici';
import {
McpClientsService,
isRetryableConnectError,
} from './mcp-clients.service';
/**
* #489 — external-MCP in-run transport recovery.
*
* The transport-error classification + retry gate are exercised against the REAL
* undici error CLASSES prod throws (`errors.SocketError` / `errors.BodyTimeoutError`,
* carrying the true `UND_ERR_*` codes and class names), wrapped EXACTLY as undici's
* `fetch` wraps them — a `TypeError('fetch failed'|'terminated')` whose `.cause` is
* the undici error. These are the real classes, not hand-rolled `{code:'...'}`
* mocks: constructing the genuine class is what makes this a faithful test of the
* prod predicate (epic root-cause #4 — a mock-shaped predicate would leave the
* evict/retry path silently dead in production while CI stays green). We construct
* rather than drive a live fetch because Jest's environment degrades the live-fetch
* error to a generic `Error` cause (no undici code), which would NOT be the prod
* shape.
*/
/** A REAL undici socket reset, wrapped as fetch wraps it. */
function realSocketResetError(): unknown {
const err = new TypeError('fetch failed');
(err as { cause?: unknown }).cause = new errors.SocketError('other side closed');
return err;
}
/** A REAL undici body timeout, wrapped as fetch wraps it. */
function realBodyTimeoutError(): unknown {
const err = new TypeError('terminated');
(err as { cause?: unknown }).cause = new errors.BodyTimeoutError();
return err;
}
type FakeServer = {
id: string;
name: string;
transport: 'http' | 'sse';
url: string;
headersEnc: string | null;
toolAllowlist: string[] | null;
instructions: string | null;
};
const server = (over: Partial<FakeServer> = {}): FakeServer => ({
id: 's1',
name: 'srv',
transport: 'http',
url: 'http://example.test/mcp',
headersEnc: null,
toolAllowlist: null,
instructions: null,
...over,
});
function buildService(servers: FakeServer[], trusted = false) {
const repo = { listEnabled: jest.fn().mockResolvedValue(servers) };
const service = new McpClientsService(repo as never, {} as never);
// Seed a DETERMINISTIC write-class map so the retry gate is controlled here
// (the production map loads from @docmost/mcp via a dynamic ESM import). getPage
// is a read, patchNode is a write — the real classifications.
(
service as unknown as { writeClassMapPromise: Promise<unknown> }
).writeClassMapPromise = Promise.resolve({
getPage: 'readOnly',
patchNode: 'write',
});
// The service only APPLIES that map to a TRUSTED internal Docmost server
// (isInternalDocmostServer, really false for every third-party row). A retry
// test needs a trusted server to exercise the readOnly-retry path at all, so it
// passes trusted=true to model a Docmost-origin server; the third-party
// double-apply test leaves it at the real value (false).
if (trusted) {
jest
.spyOn(
service as unknown as {
isInternalDocmostServer: (s: FakeServer) => boolean;
},
'isInternalDocmostServer',
)
.mockReturnValue(true);
}
return { service, repo };
}
/** Spy the private `connect` so each call yields a controlled fake client whose
* single tool's execute is the supplied function. Returns the connect spy. */
function stubConnect(
service: McpClientsService,
toolName: string,
execs: Array<(...a: unknown[]) => Promise<unknown>>,
) {
let n = 0;
return jest
.spyOn(
service as unknown as { connect: (s: FakeServer) => Promise<unknown> },
'connect',
)
.mockImplementation(async () => {
const exec = execs[Math.min(n, execs.length - 1)];
n += 1;
return {
tools: async () => ({ [toolName]: { description: 'x', execute: exec } }),
close: jest.fn().mockResolvedValue(undefined),
};
});
}
const opts = (abortSignal?: AbortSignal) =>
({ toolCallId: 't', messages: [], abortSignal }) as never;
describe('isRetryableConnectError (#489, REAL error shapes)', () => {
it('classifies a real undici socket reset and body timeout as retryable', async () => {
const socketErr = await realSocketResetError();
const bodyErr = await realBodyTimeoutError();
expect(isRetryableConnectError(socketErr)).toBe(true);
expect(isRetryableConnectError(bodyErr)).toBe(true);
// Unwraps a wrapped cause chain (e.g. an MCPClientError around the socket err).
const wrapped = new Error('mcp call failed');
(wrapped as { cause?: unknown }).cause = socketErr;
expect(isRetryableConnectError(wrapped)).toBe(true);
});
it('does NOT classify an application-level error as a transport break', () => {
expect(isRetryableConnectError(new Error('validation failed'))).toBe(false);
expect(isRetryableConnectError({ name: 'HttpError', status: 400 })).toBe(false);
expect(isRetryableConnectError(undefined)).toBe(false);
expect(isRetryableConnectError('boom')).toBe(false);
});
});
describe('McpClientsService in-run transport recovery (#489)', () => {
afterEach(() => jest.restoreAllMocks());
it('a readOnly tool whose transport breaks reconnects and retries WITHIN the same run', async () => {
const realErr = await realSocketResetError();
const { service } = buildService([server()], true);
const first = jest.fn().mockRejectedValue(realErr);
const second = jest.fn().mockResolvedValue({ ok: true });
const connectSpy = stubConnect(service, 'getPage', [first, second]);
const toolset = await service.toolsFor('ws-1');
const tool = toolset.tools['srv_getPage'];
const result = await (tool.execute as (a: unknown, o: unknown) => Promise<unknown>)(
{ pageId: 'p' },
opts(),
);
// The repeat call within the run got a LIVE client and succeeded.
expect(result).toEqual({ ok: true });
expect(first).toHaveBeenCalledTimes(1);
expect(second).toHaveBeenCalledTimes(1);
// Exactly one reconnect was minted (initial build connect + one recovery).
expect(connectSpy).toHaveBeenCalledTimes(2);
// The run accumulated BOTH leases (old + reconnected) — released together at end.
expect(toolset.clients).toHaveLength(2);
await Promise.all(toolset.clients.map((c) => c.close()));
});
it('a WRITE tool does NOT auto-retry on a transport error (indeterminate)', async () => {
const realErr = await realSocketResetError();
const { service } = buildService([server()], true);
const exec = jest.fn().mockRejectedValue(realErr);
const connectSpy = stubConnect(service, 'patchNode', [exec]);
const toolset = await service.toolsFor('ws-2');
const tool = toolset.tools['srv_patchNode'];
await expect(
(tool.execute as (a: unknown, o: unknown) => Promise<unknown>)(
{ pageId: 'p' },
opts(),
),
).rejects.toThrow(/MAY have already applied/);
// Called exactly once — NO blind retry (avoids double-apply, the #435 class).
expect(exec).toHaveBeenCalledTimes(1);
// No fresh connection was minted for a write.
expect(connectSpy).toHaveBeenCalledTimes(1);
await Promise.all(toolset.clients.map((c) => c.close()));
});
it('does NOT retry (or reconnect) after the run is aborted (Stop)', async () => {
const realErr = await realSocketResetError();
const { service } = buildService([server()], true);
const controller = new AbortController();
// The transport error arrives, but the run was Stopped in the same tick.
const first = jest.fn().mockImplementation(async () => {
controller.abort();
throw realErr;
});
const second = jest.fn().mockResolvedValue({ ok: true });
const connectSpy = stubConnect(service, 'getPage', [first, second]);
const toolset = await service.toolsFor('ws-3');
const tool = toolset.tools['srv_getPage'];
await expect(
(tool.execute as (a: unknown, o: unknown) => Promise<unknown>)(
{ pageId: 'p' },
opts(controller.signal),
),
).rejects.toBeDefined();
// getPage IS readOnly, but the Stop blocks the retry — no second call, no mint.
expect(second).not.toHaveBeenCalled();
expect(connectSpy).toHaveBeenCalledTimes(1);
await Promise.all(toolset.clients.map((c) => c.close()));
});
it('an app-level (non-transport) tool error is surfaced verbatim, never retried', async () => {
const { service } = buildService([server()], true);
const appErr = new Error('tool says: bad input');
const exec = jest.fn().mockRejectedValue(appErr);
const connectSpy = stubConnect(service, 'getPage', [exec]);
const toolset = await service.toolsFor('ws-4');
const tool = toolset.tools['srv_getPage'];
await expect(
(tool.execute as (a: unknown, o: unknown) => Promise<unknown>)(
{ pageId: 'p' },
opts(),
),
).rejects.toThrow('tool says: bad input');
expect(exec).toHaveBeenCalledTimes(1);
expect(connectSpy).toHaveBeenCalledTimes(1); // no reconnect for an app error
await Promise.all(toolset.clients.map((c) => c.close()));
});
// #489 (review, MEDIUM) — the Docmost write-class map keys by DOCMOST tool
// names; a THIRD-PARTY server may name a WRITE tool `getPage` (a Docmost read
// name). It must NOT inherit readOnly and must NOT auto-retry on a transport
// error — a blind retry of that write is a double-apply (the #435 class). Here
// the server is UNTRUSTED (buildService default, isInternalDocmostServer=false),
// so the map is not applied and `getPage` classifies as a write.
//
// MUTATION-VERIFY: forcing the server "trusted" (buildService(..., true)) makes
// `getPage` inherit readOnly -> it WOULD reconnect+retry (connect twice) and the
// assertions below fail — i.e. removing the trust scope re-opens the bug.
it('a THIRD-PARTY WRITE tool named like a Docmost read does NOT auto-retry (no double-apply)', async () => {
const realErr = await realSocketResetError();
// Untrusted: default trusted=false — a real third-party server.
const { service } = buildService([server()]);
const exec = jest.fn().mockRejectedValue(realErr);
const connectSpy = stubConnect(service, 'getPage', [exec, exec]);
const toolset = await service.toolsFor('ws-5');
const tool = toolset.tools['srv_getPage'];
await expect(
(tool.execute as (a: unknown, o: unknown) => Promise<unknown>)(
{ pageId: 'p' },
opts(),
),
).rejects.toThrow(/MAY have already applied/);
// Exactly one call, NO reconnect — the name collision granted no readOnly-retry.
expect(exec).toHaveBeenCalledTimes(1);
expect(connectSpy).toHaveBeenCalledTimes(1);
await Promise.all(toolset.clients.map((c) => c.close()));
});
});
@@ -106,8 +106,11 @@ describe('McpClientsService.decryptHeaders', () => {
describe('McpClientsService.guardedFetch (SSRF per-request guard)', () => {
// The bound guardedFetch closure lives on the instance as a private field.
// #489 split it into per-transport HTTP/SSE bindings (they differ only in the
// dispatcher's bodyTimeout); the SSRF guard is identical, so testing the HTTP
// one is sufficient.
const guardedFetchOf = (service: McpClientsService) =>
(service as unknown as { guardedFetch: typeof fetch }).guardedFetch;
(service as unknown as { guardedFetchHttp: typeof fetch }).guardedFetchHttp;
let fetchSpy: jest.SpiedFunction<typeof fetch>;
@@ -1,5 +1,6 @@
import { isIP } from 'node:net';
import { lookup as dnsLookup, type LookupAddress } from 'node:dns';
import { pathToFileURL } from 'node:url';
import { Injectable, Logger } from '@nestjs/common';
import { type Tool, type ToolCallOptions } from 'ai';
import { createMCPClient } from '@ai-sdk/mcp';
@@ -10,9 +11,29 @@ import {
streamingDispatcherOptions,
mcpStreamTimeoutMs,
mcpCallTimeoutMs,
mcpSseBodyTimeoutMs,
} from '../../../integrations/ai/ai-streaming-fetch';
import { SecretBoxService } from '../../../integrations/crypto/secret-box';
import { isUrlAllowed, isIpAllowed } from './ssrf-guard';
// TYPE-ONLY (erased at compile): @docmost/mcp is ESM-only and cannot be a runtime
// `require()` from this commonjs module (same constraint as docmost-client.loader).
// The write-class MAP is loaded lazily via the dynamic-import trick below.
import type { ToolWriteClass } from '@docmost/mcp';
// TS(commonjs) downlevels a literal `import()` to `require()`, which cannot load
// the ESM-only @docmost/mcp. Indirect through Function so the real dynamic
// `import()` survives compilation (same trick as docmost-client.loader.ts).
const esmImport = new Function(
'specifier',
'return import(specifier)',
) as (specifier: string) => Promise<unknown>;
/** Local read-only predicate — avoids a value import of the ESM-only package.
* Only a pure read is retry-safe after a transport break (a write is
* indeterminate). Kept in lockstep with @docmost/mcp's isRetryableWriteClass. */
function isReadOnlyWriteClass(writeClass: ToolWriteClass | undefined): boolean {
return writeClass === 'readOnly';
}
/** A closable external MCP client handle. */
export interface Closable {
@@ -81,12 +102,52 @@ const MAX_TOOL_NAME_LENGTH = 64;
* close until the turn releases it, so a TTL expiry mid-turn never closes a
* client a stream is still executing against.
*/
/**
* Where a merged (namespaced) tool came from, so the per-run recovery wrapper
* (#489) can, on a transport error, reconnect THAT server and re-resolve the SAME
* underlying tool by its raw name. `writeClass` gates the single auto-retry (a
* read is retry-safe; a write is indeterminate). `serverIndex` indexes the
* entry's `servers` array (which server config to reconnect).
*/
interface ToolProvenance {
serverIndex: number;
rawName: string;
writeClass: ToolWriteClass | undefined;
}
/** A live reconnected server (its fresh client + raw call-timeout-wrapped tools). */
interface RecoveredServerState {
client: McpClient;
tools: Record<string, Tool>;
}
/**
* Per-run, per-server recovery binding (#489). `current` is the server's LIVE
* target for this run: `null` means "use the ORIGINAL cached client/template";
* a non-null value is a reconnected throwaway client all this server's tools now
* call. `reconnecting` dedupes concurrent reconnects so only ONE fresh client is
* minted per death (a losing concurrent call awaits it and retries on the SAME
* new client — the CAS-by-identity rule).
*/
interface ServerBinding {
current: RecoveredServerState | null;
reconnecting?: Promise<RecoveredServerState>;
}
interface CacheEntry {
tools: Record<string, Tool>;
clients: McpClient[];
outcomes: ServerOutcome[];
/** Prompt guidance for qualifying servers (see McpServerInstruction). */
instructions: McpServerInstruction[];
/**
* The enabled server configs used to build this entry (#489), so the per-run
* recovery wrapper can reconnect a specific server by index. Parallel to the
* indices referenced by {@link toolMeta}.
*/
servers: AiMcpServer[];
/** merged-tool-key -> provenance (#489), for the per-run recovery wrapper. */
toolMeta: Record<string, ToolProvenance>;
expiresAt: number;
/** Active leases (turns currently using these clients). */
refCount: number;
@@ -120,20 +181,82 @@ export class McpClientsService {
*/
private readonly cache = new Map<string, Promise<CacheEntry>>();
/**
* A single shared SSRF-pinned dispatcher for ALL outbound external-MCP fetches.
* Its custom connect.lookup runs per connection, so one instance safely guards
* every server's connections (we never connect to an unvalidated IP).
* SSRF-pinned dispatchers for outbound external-MCP fetches. Both use the SAME
* custom connect.lookup (so every connection is IP-validated), but carry a
* DIFFERENT `bodyTimeout` (#489): the HTTP (streamable) transport opens a fresh
* request per call, so it keeps the tight silence timeout; the SSE transport
* holds ONE long-lived body open across many calls, so a >1-min idle BETWEEN
* calls is LEGITIMATE and must not break the socket — it gets a much larger
* bodyTimeout. (headersTimeout stays tight on both.)
*/
private readonly dispatcher: Dispatcher = buildPinnedDispatcher();
/** guardedFetch bound to the pinned dispatcher; reused by every transport. */
private readonly guardedFetch: typeof fetch = (input, init) =>
guardedFetch(this.dispatcher, input, init);
private readonly dispatcherHttp: Dispatcher = buildPinnedDispatcher(
mcpStreamTimeoutMs(),
);
private readonly dispatcherSse: Dispatcher = buildPinnedDispatcher(
mcpSseBodyTimeoutMs(),
);
/** guardedFetch bound to each dispatcher; picked by transport type in connect(). */
private readonly guardedFetchHttp: typeof fetch = (input, init) =>
guardedFetch(this.dispatcherHttp, input, init);
private readonly guardedFetchSse: typeof fetch = (input, init) =>
guardedFetch(this.dispatcherSse, input, init);
/**
* Memoized write-class map (#489), loaded lazily from @docmost/mcp via the
* dynamic-import trick. Keyed by tool name (=== mcpName). A tool NOT in the map
* (any third-party external MCP tool) classifies as `undefined` -> treated as a
* write by the retry gate (the safe default: never blind-retry an unknown tool).
* On any load failure the map is `{}` (every tool -> no auto-retry), so a
* missing/older @docmost/mcp build only DISABLES retries, never mis-retries.
*/
private writeClassMapPromise: Promise<Record<string, ToolWriteClass>> | null =
null;
constructor(
private readonly repo: AiMcpServerRepo,
private readonly secretBox: SecretBoxService,
) {}
/**
* Whether an external MCP server is the TRUSTED internal Docmost MCP server —
* the only server whose tools may be classified by the Docmost write-class map
* (#489 review). Today this is ALWAYS false: every `ai_mcp_servers` row is an
* admin-configured THIRD-PARTY endpoint (there is no builtin/self flag, sentinel
* URL, or synthetic server in this path — Docmost's OWN tools are exposed via the
* separate in-app tools path, never through this external-MCP client). So no
* third-party tool can inherit `readOnly` by a name collision with a Docmost read
* tool, and none is ever auto-retried on a transport error (which would risk a
* double-apply — the #435 class). Flip this (an explicit `kind`/`isBuiltin`
* column, or a configured self-MCP URL) if a trusted internal server is ever
* introduced. A method (not a free function) so it is a single, mockable seam.
*/
private isInternalDocmostServer(_server: AiMcpServer): boolean {
return false;
}
/** Lazily load + memoize the shared write-class map (see the field doc). */
private getWriteClassMap(): Promise<Record<string, ToolWriteClass>> {
if (!this.writeClassMapPromise) {
this.writeClassMapPromise = (async () => {
try {
const entry = require.resolve('@docmost/mcp');
const mod = (await esmImport(pathToFileURL(entry).href)) as {
SHARED_TOOL_WRITE_CLASS?: Record<string, ToolWriteClass>;
};
return mod.SHARED_TOOL_WRITE_CLASS ?? {};
} catch (err) {
this.logger.warn(
`Could not load MCP write-class map (auto-retry disabled): ${shortError(
err,
)}`,
);
return {};
}
})();
}
return this.writeClassMapPromise;
}
/**
* Build (or reuse a cached) external toolset for a workspace. Returns the
* merged tools, the open client handles to release, and per-server outcomes.
@@ -162,11 +285,37 @@ export class McpClientsService {
}
},
};
// One release handle drives the whole leased entry; closing it releases all
// underlying clients together (they share the same lease lifecycle).
// #489: the run accumulates a SET of leases — the primary cache lease PLUS any
// throwaway client minted by an in-run transport-recovery reconnect. They are
// NEVER released mid-run (releasing a swapped-out client while a concurrent
// in-flight call still holds it would INDUCE a second failure); the caller
// releases the WHOLE set together at turn-end. A recovery reconnect pushes its
// lease onto this live array, which the consumer closes over.
const leaseSet: Closable[] = [release];
// #489: per-RUN transport-recovery binding, one per server, SHARED by all of
// that server's tools so a swap by one call is seen by the next (CAS by
// identity). Kept per-run (here, not in the cached entry) because the binding
// + lease-set state is per-run.
const bindings = new Map<number, ServerBinding>();
const capMs = mcpCallTimeoutMs();
// Wrap each cached tool with the recovery layer. On a transport error a
// declared readOnly tool reconnects its server and retries ONCE; a write is
// never blind-retried (indeterminate — may have applied before the reset). A
// tool without provenance (a minimal stub entry in a test) passes through raw.
const tools: Record<string, Tool> = {};
for (const [key, tool] of Object.entries(entry.tools)) {
const meta = entry.toolMeta?.[key];
tools[key] = meta
? this.wrapWithTransportRecovery(entry, meta, tool, leaseSet, bindings, capMs)
: tool;
}
return {
tools: entry.tools,
clients: [release],
tools,
clients: leaseSet,
outcomes: entry.outcomes,
instructions: entry.instructions,
};
@@ -254,6 +403,16 @@ export class McpClientsService {
// Per-call total wall-clock cap, read once for this build (env-overridable).
const callTimeoutMs = mcpCallTimeoutMs();
const instructions: McpServerInstruction[] = [];
// merged-key -> provenance for the per-run recovery wrapper (#489).
const toolMeta: Record<string, ToolProvenance> = {};
// Shared Docmost write-class map (#489) — classifies a tool by its raw name.
// Loaded ONLY when at least one server is a TRUSTED internal Docmost server
// (see isInternalDocmostServer): for third-party servers the map is never
// applied (a name collision must not grant readOnly-retry), so we skip the
// dynamic ESM load entirely in that (currently universal) case.
const writeClassMap = servers.some((s) => this.isInternalDocmostServer(s))
? await this.getWriteClassMap()
: null;
// Per-server connect+tools result, still tagged with its server so the merge
// below can be applied in the SAME order as `servers` (see the parallel note).
@@ -327,11 +486,23 @@ export class McpClientsService {
// against names already merged from earlier servers, so no external
// tool is silently overwritten on collision. The returned count drives
// whether this server's prompt guidance is included (≥1 tool merged).
// #489 (review): the Docmost write-class map keys by DOCMOST tool names and
// may ONLY be trusted for a server KNOWN to be the internal Docmost MCP
// server. Every row here is an admin-configured THIRD-PARTY endpoint, so a
// third-party WRITE tool that happens to be named like a Docmost read
// (getPage, listPages, ...) must NOT inherit readOnly — that would auto-retry
// a mutation on a transport error (double-apply, the #435 class). Gate the
// map on the trust check; untrusted servers get writeClass=undefined -> the
// recovery wrapper treats them as writes and never auto-retries.
const trustWriteClass = this.isInternalDocmostServer(server);
const merged = this.mergeNamespaced(
tools,
result.guarded,
server.name,
server.id,
toolMeta,
i,
trustWriteClass ? writeClassMap : null,
);
outcomes.push({ name: server.name, ok: true });
// Include this server's guidance ONLY when it actually contributed at
@@ -353,6 +524,8 @@ export class McpClientsService {
clients,
outcomes,
instructions,
servers,
toolMeta,
expiresAt: Date.now() + CACHE_TTL_MS,
refCount: 0,
evicted: false,
@@ -379,18 +552,33 @@ export class McpClientsService {
picked: Record<string, Tool>,
serverName: string,
serverId: string,
toolMeta: Record<string, ToolProvenance>,
serverIndex: number,
// The Docmost write-class map, or `null` for an UNTRUSTED (third-party)
// server whose tools must all default to write (never auto-retried).
writeClassMap: Record<string, ToolWriteClass> | null,
): { count: number; prefix: string } {
let count = 0;
for (const [name, tool] of Object.entries(namespace(picked, serverName))) {
let key = name;
for (const { full, raw, tool } of namespace(picked, serverName)) {
let key = full;
if (key in target) {
const original = key;
key = disambiguate(name, serverId, (candidate) => candidate in target);
key = disambiguate(full, serverId, (candidate) => candidate in target);
this.logger.debug(
`External MCP tool name "${original}" collided; renamed to "${key}"`,
);
}
target[key] = tool;
// Record provenance so the per-run recovery wrapper (#489) can reconnect
// this tool's server and re-resolve it by its raw name. writeClass is set
// ONLY from a TRUSTED (internal-Docmost) map; for a third-party server the
// map is null -> writeClass stays undefined -> the wrapper treats the tool
// as a write and never auto-retries it (no double-apply on name collision).
toolMeta[key] = {
serverIndex,
rawName: raw,
writeClass: writeClassMap ? writeClassMap[raw] : undefined,
};
count += 1;
}
return { count, prefix: namespacePrefix(serverName) };
@@ -424,7 +612,10 @@ export class McpClientsService {
// Defense in depth: re-validate the actual request host on EVERY fetch
// AND pin the socket to a validated IP via the dispatcher's connect
// lookup, closing the DNS-rebinding TOCTOU between check and connect.
fetch: this.guardedFetch,
// #489: the SSE transport uses the raised-bodyTimeout dispatcher (idle
// between calls is legit); HTTP uses the tight one.
fetch:
transportType === 'sse' ? this.guardedFetchSse : this.guardedFetchHttp,
},
})) as unknown as McpClient;
return client;
@@ -505,6 +696,176 @@ export class McpClientsService {
}
}
/**
* Wrap one merged external tool with the per-run transport-recovery layer (#489).
*
* attempt 1 runs on the server's CURRENT binding (the cached client, or a client
* a sibling tool already reconnected this run). On a REAL transport error
* (undici/@ai-sdk socket/body-timeout shapes — {@link isRetryableConnectError},
* NOT a mock) and ONLY for a declared readOnly tool, it reconnects the server
* and retries EXACTLY ONCE on the fresh client; a write is surfaced as an
* indeterminate error (it may have applied before the reset — never
* blind-retried). A single per-call cap bounds BOTH attempts + the reconnect,
* and the run's abort signal is checked before the retry AND before minting a
* fresh connection (no connection is opened for a stopped run).
*/
private wrapWithTransportRecovery(
entry: CacheEntry,
meta: ToolProvenance,
template: Tool,
leaseSet: Closable[],
bindings: Map<number, ServerBinding>,
capMs: number,
): Tool {
const original = template.execute;
if (typeof original !== 'function') return template;
const service = this;
const { serverIndex, rawName, writeClass } = meta;
let binding = bindings.get(serverIndex);
if (!binding) {
binding = { current: null };
bindings.set(serverIndex, binding);
}
const boundBinding = binding;
const execute = async (args: unknown, options: ToolCallOptions) => {
// The per-call cap governs the WHOLE sequence (attempt1 + reconnect +
// attempt2). Compose it with the run's abort signal so a Stop or the cap
// ends any awaited call — @ai-sdk/mcp does not settle on abort, so we RACE.
const capController = new AbortController();
const capTimer = setTimeout(() => {
capController.abort(new Error(`MCP tool call timed out after ${capMs}ms`));
}, capMs);
capTimer.unref?.();
const runSignal = options?.abortSignal;
const composed = runSignal
? AbortSignal.any([runSignal, capController.signal])
: capController.signal;
const stopped = () => runSignal?.aborted === true || capController.signal.aborted;
const callOn = async (
exec: NonNullable<Tool['execute']>,
): Promise<unknown> => {
const aborted = new Promise<never>((_, reject) => {
const fail = () => reject(abortReason(composed));
if (composed.aborted) fail();
else composed.addEventListener('abort', fail, { once: true });
});
return Promise.race([exec(args, { ...options, abortSignal: composed }), aborted]);
};
const execFor = (
state: RecoveredServerState | null,
): NonNullable<Tool['execute']> | undefined =>
state ? (state.tools[rawName]?.execute as NonNullable<Tool['execute']>) : original;
try {
// Snapshot the target BEFORE the call so a swap by a concurrent call is
// detected by identity in the catch.
const attemptState = boundBinding.current;
const attemptExec = execFor(attemptState);
if (typeof attemptExec !== 'function') {
throw new Error(`external MCP tool "${rawName}" is not callable`);
}
try {
return await callOn(attemptExec);
} catch (err) {
// Never retry on a Stop or an exhausted cap.
if (stopped()) throw err;
// Only a genuine transport break is a recovery candidate.
if (!isRetryableConnectError(err)) throw err;
// A write tool is INDETERMINATE on a transport error (may have applied
// before the reset) — surface that; do NOT auto-retry (double-apply is
// the #435 incident class).
if (!isReadOnlyWriteClass(writeClass)) {
throw new Error(
`external MCP tool "${rawName}" hit a transport error and MAY have already ` +
`applied on the server — not retried automatically; verify state before ` +
`retrying. (${shortError(err)})`,
);
}
// Abort check BEFORE minting a fresh connection (no socket for a
// stopped run). LIMITATION (#489, LOW): the reconnect's own connect is
// bounded by CONNECT_TIMEOUT_MS but does NOT itself observe `composed`,
// so a Stop that lands DURING the handshake is only honored at the next
// `stopped()` gate (before the retry) — a bounded ≤5s late-abort window;
// the throwaway client is closed at turn-end regardless. Threading
// `composed` into the SHARED (CAS-deduped) reconnect is deliberately
// avoided: it would let the first caller's abort tear down a reconnect a
// concurrent still-live caller depends on.
if (stopped()) throw err;
// CAS-swap by IDENTITY: mint+swap only if nobody swapped since this
// call's snapshot; a losing concurrent call awaits the same reconnect
// and retries on the SAME fresh client.
let target: RecoveredServerState;
if (boundBinding.current === attemptState) {
if (!boundBinding.reconnecting) {
boundBinding.reconnecting = (async () => {
const server = entry.servers[serverIndex];
const fresh = await service.reconnectServer(server, capMs);
leaseSet.push(fresh.lease); // accumulate; released at turn-end
boundBinding.current = fresh.state;
return fresh.state;
})();
// Clear the in-flight marker once it settles (success or failure) so
// a LATER death of the new client can reconnect again.
void boundBinding.reconnecting.then(
() => (boundBinding.reconnecting = undefined),
() => (boundBinding.reconnecting = undefined),
);
}
target = await boundBinding.reconnecting;
} else {
target = boundBinding.current as RecoveredServerState;
}
// Abort check BEFORE the retry.
if (stopped()) throw err;
const retryExec = execFor(target);
if (typeof retryExec !== 'function') throw err;
return await callOn(retryExec);
}
} finally {
clearTimeout(capTimer);
}
};
return { ...template, execute } as unknown as Tool;
}
/**
* Reconnect ONE server for an in-run recovery (#489): open a fresh client and
* list+wrap its tools. The throwaway client is NOT cached — it is owned by the
* RUN via the returned lease (closed at turn-end), independent of the shared
* cache entry (whose TTL rebuild heals future turns). On a failure the fresh
* client is closed so its socket never leaks.
*/
private async reconnectServer(
server: AiMcpServer,
capMs: number,
): Promise<{ state: RecoveredServerState; lease: Closable }> {
const client = await this.connectWithTimeout(server, CONNECT_TIMEOUT_MS);
let tools: Record<string, Tool>;
try {
const raw = await withTimeout(client.tools(), CONNECT_TIMEOUT_MS);
const allow = server.toolAllowlist;
const picked =
Array.isArray(allow) && allow.length > 0 ? pick(raw, allow) : raw;
tools = wrapToolsWithCallTimeout(picked, capMs);
} catch (err) {
void client.close().catch(() => undefined);
throw err;
}
let released = false;
const lease: Closable = {
close: async () => {
if (released) return;
released = true;
await client.close().catch(() => undefined);
},
};
return { state: { client, tools }, lease };
}
/** Mark an entry evicted; close its clients now if nothing is leasing them. */
private evict(entry: CacheEntry): void {
clearTimeout(entry.timer);
@@ -554,22 +915,21 @@ export function validateResolvedAddresses(addrs: readonly LookupAddress[]): {
* certificate validation still uses the real hostname (we never rewrite the URL
* to an IP literal).
*/
function buildPinnedDispatcher(): Agent {
// External-MCP traffic uses a DEDICATED, shorter silence timeout
function buildPinnedDispatcher(bodyTimeoutMs: number): Agent {
// External-MCP traffic uses a DEDICATED, shorter HEADERS silence timeout
// (`AI_MCP_STREAM_TIMEOUT_MS`, default 1 min) — deliberately tighter than the
// chat provider's 15-min `streamTimeoutMs()` — so a byte-silent/hung MCP
// upstream is broken in ~1 min instead of 15. We keep the keep-alive options
// from `streamingDispatcherOptions()` but OVERRIDE headers/body timeouts.
// Accepted trade-off: a legitimately long but byte-silent single tool call,
// and an SSE transport idling >1 min BETWEEN tool calls, are also cut here; the
// per-call total cap (wrapToolsWithCallTimeout, `AI_MCP_CALL_TIMEOUT_MS`) is the
// complementary guard for chatty-but-stuck calls that keep the socket warm yet
// never return.
const mcpSilenceMs = mcpStreamTimeoutMs();
// from `streamingDispatcherOptions()` but OVERRIDE the timeouts. `bodyTimeout`
// is passed in per-transport (#489): tight for HTTP (fresh request per call),
// raised for SSE (one long-lived body across calls — idle BETWEEN calls is
// legit). The per-call total cap (`AI_MCP_CALL_TIMEOUT_MS`) is the complementary
// guard for chatty-but-stuck calls that keep the socket warm yet never return.
const headersMs = mcpStreamTimeoutMs();
return new Agent({
...streamingDispatcherOptions(),
headersTimeout: mcpSilenceMs,
bodyTimeout: mcpSilenceMs,
headersTimeout: headersMs,
bodyTimeout: bodyTimeoutMs,
connect: {
lookup: (hostname, _options, callback) => {
// Always resolve ALL addresses ourselves; do not trust the caller's
@@ -669,18 +1029,22 @@ function pick(
function namespace(
tools: Record<string, Tool>,
serverName: string,
): Record<string, Tool> {
): Array<{ full: string; raw: string; tool: Tool }> {
const prefix = namespacePrefix(serverName);
const out: Record<string, Tool> = {};
const out: Array<{ full: string; raw: string; tool: Tool }> = [];
const taken: Record<string, true> = {};
for (const [name, t] of Object.entries(tools)) {
const safe = sanitizeName(name);
let full = capName(`${prefix}_${safe}`);
// Duplicate names within ONE server can still collide after sanitize/
// truncate — suffix-disambiguate so the second tool is not overwritten.
if (full in out) {
full = disambiguate(full, '', (candidate) => candidate in out);
if (full in taken) {
full = disambiguate(full, '', (candidate) => candidate in taken);
}
out[full] = t;
taken[full] = true;
// Keep the RAW (un-namespaced) name alongside the merged key so the per-run
// recovery wrapper (#489) can re-resolve the same tool on a fresh client.
out.push({ full, raw: name, tool: t });
}
return out;
}
@@ -804,6 +1168,69 @@ export function wrapToolWithCallTimeout(tool: Tool, ms: number): Tool {
return { ...tool, execute } as unknown as Tool;
}
/**
* undici / Node network error CODES that mean the connection broke (not an
* application-level error) — a transient transport failure a readOnly call may
* safely retry after reconnecting. Matched against the REAL error shapes (#489):
* a socket reset surfaces as `TypeError: fetch failed` whose `.cause` is an
* undici `SocketError { code:'UND_ERR_SOCKET' }`; a body-timeout as
* `TypeError: terminated` whose `.cause` is `BodyTimeoutError`. Classifying by
* these real codes/names (not by mock errors) is essential — a mock-shaped
* predicate would leave eviction silently dead in production while CI is green.
*/
const RETRYABLE_TRANSPORT_ERROR_CODES: ReadonlySet<string> = new Set([
'ECONNRESET',
'ECONNREFUSED',
'ECONNABORTED',
'EPIPE',
'ETIMEDOUT',
'EAI_AGAIN',
'ENETUNREACH',
'EHOSTUNREACH',
'UND_ERR_SOCKET',
'UND_ERR_CONNECT_TIMEOUT',
'UND_ERR_HEADERS_TIMEOUT',
'UND_ERR_BODY_TIMEOUT',
'UND_ERR_CLOSED',
'UND_ERR_DESTROYED',
]);
/** undici error CLASS names for the same transport-break conditions. */
const RETRYABLE_TRANSPORT_ERROR_NAMES: ReadonlySet<string> = new Set([
'SocketError',
'BodyTimeoutError',
'HeadersTimeoutError',
'ConnectTimeoutError',
'ClientClosedError',
'ClientDestroyedError',
]);
/**
* Whether `err` is a retryable TRANSPORT break (a broken socket / body timeout),
* classified by the REAL undici/@ai-sdk error shapes (#489). undici surfaces a
* reset as `TypeError('fetch failed'|'terminated')` with the real error in
* `.cause`, and @ai-sdk/mcp may wrap it again in an `MCPClientError` (cause
* chain), so we walk `.cause` (bounded depth) checking `.code` and `.name`. An
* app-level tool error (a 4xx, a validation failure) is NOT retryable and returns
* false — only a connection-level failure heals with a reconnect.
*/
export function isRetryableConnectError(err: unknown, depth = 0): boolean {
if (!err || typeof err !== 'object' || depth > 6) return false;
const e = err as {
code?: unknown;
name?: unknown;
cause?: unknown;
};
if (typeof e.code === 'string' && RETRYABLE_TRANSPORT_ERROR_CODES.has(e.code)) {
return true;
}
if (typeof e.name === 'string' && RETRYABLE_TRANSPORT_ERROR_NAMES.has(e.name)) {
return true;
}
if (e.cause != null) return isRetryableConnectError(e.cause, depth + 1);
return false;
}
/** The signal's reason as an Error (informative thrown value on abort/timeout). */
function abortReason(signal: AbortSignal): Error {
const r = signal.reason;
@@ -0,0 +1,266 @@
import type { ModelMessage } from 'ai';
import {
resolveReplayBudget,
isContextOverflowError,
estimateMessagesTokens,
trimHistoryForReplay,
REPLAY_BUDGET_DEFAULT_TOKENS,
REPLAY_TRUNCATION_MARKER,
REPLAY_TURN_COLLAPSED_MARKER,
} from './history-budget';
describe('resolveReplayBudget', () => {
it('uses min(default, 0.7 x window) for a configured window', () => {
// 0.7 x 60k = 42k < 100k
expect(resolveReplayBudget(60_000)).toEqual({
thresholdTokens: 42_000,
usedDefault: false,
});
// 0.7 x 1M = 700k, capped to the 100k default
expect(resolveReplayBudget(1_000_000)).toEqual({
thresholdTokens: REPLAY_BUDGET_DEFAULT_TOKENS,
usedDefault: false,
});
});
it('accepts the raw ::text stored form', () => {
expect(resolveReplayBudget('60000').thresholdTokens).toBe(42_000);
});
// The crux (#490): a chat with NO context window configured must STILL be
// budgeted — those are exactly the installs that hit terminal overflow.
it('applies the flat default when the window is unset/empty', () => {
expect(resolveReplayBudget(undefined)).toEqual({
thresholdTokens: REPLAY_BUDGET_DEFAULT_TOKENS,
usedDefault: true,
});
expect(resolveReplayBudget('')).toEqual({
thresholdTokens: REPLAY_BUDGET_DEFAULT_TOKENS,
usedDefault: true,
});
expect(resolveReplayBudget(' ')).toEqual({
thresholdTokens: REPLAY_BUDGET_DEFAULT_TOKENS,
usedDefault: true,
});
});
it('treats an explicit 0 as the off-switch (distinct from unset)', () => {
expect(resolveReplayBudget(0)).toEqual({
thresholdTokens: null,
usedDefault: false,
});
expect(resolveReplayBudget('0')).toEqual({
thresholdTokens: null,
usedDefault: false,
});
});
it('falls back to the default on a negative/garbage value', () => {
expect(resolveReplayBudget(-5).usedDefault).toBe(true);
expect(resolveReplayBudget('abc').usedDefault).toBe(true);
});
});
describe('isContextOverflowError', () => {
it('classifies a real provider 400 context-overflow shape', () => {
// OpenAI-compatible shape.
expect(
isContextOverflowError({
statusCode: 400,
message:
"This model's maximum context length is 128000 tokens. However, your messages resulted in 214000 tokens. Please reduce the length of the messages.",
}),
).toBe(true);
// Anthropic-style wording.
expect(
isContextOverflowError({
status: 400,
message: 'prompt is too long: 250000 tokens > 200000 maximum',
}),
).toBe(true);
// Nested body + string status.
expect(
isContextOverflowError({
response: { status: '400' },
message: 'input is too long for the requested model',
}),
).toBe(true);
// Error instance with the cause carrying the body.
const e = new Error('Bad request');
(e as any).statusCode = 400;
(e as any).cause = new Error('maximum context window exceeded');
expect(isContextOverflowError(e)).toBe(true);
});
it('does NOT classify unrelated 400s or auth/rate-limit errors', () => {
expect(
isContextOverflowError({ statusCode: 400, message: 'invalid tool schema' }),
).toBe(false);
expect(
isContextOverflowError({
statusCode: 429,
message: 'context length exceeded but rate limited',
}),
).toBe(false);
expect(isContextOverflowError({ statusCode: 500, message: 'server error' })).toBe(
false,
);
expect(isContextOverflowError(undefined)).toBe(false);
expect(isContextOverflowError('some random string')).toBe(false);
});
});
// Helpers to build ModelMessage fixtures in the ai@6 shape.
const userMsg = (text: string): ModelMessage =>
({ role: 'user', content: [{ type: 'text', text }] }) as ModelMessage;
const assistantMsg = (
text: string,
toolCallId?: string,
toolName?: string,
): ModelMessage =>
({
role: 'assistant',
content: [
{ type: 'text', text },
...(toolCallId
? [{ type: 'tool-call', toolCallId, toolName, input: {} }]
: []),
],
}) as ModelMessage;
const toolMsg = (
toolCallId: string,
toolName: string,
value: unknown,
): ModelMessage =>
({
role: 'tool',
content: [
{ type: 'tool-result', toolCallId, toolName, output: { type: 'json', value } },
],
}) as ModelMessage;
describe('trimHistoryForReplay', () => {
it('null budget disables trimming (returns the same reference)', () => {
const msgs = [userMsg('hi'), assistantMsg('yo')];
const r = trimHistoryForReplay(msgs, null);
expect(r.trimmed).toBe(false);
expect(r.messages).toBe(msgs);
});
it('leaves history under budget untouched (same reference)', () => {
const msgs = [userMsg('hi'), assistantMsg('a short answer')];
const r = trimHistoryForReplay(msgs, 100_000);
expect(r.trimmed).toBe(false);
expect(r.messages).toBe(msgs);
});
it('truncates OLD tool outputs but keeps recent turns full', () => {
const big = 'X'.repeat(40_000); // ~16k tokens on its own
const msgs: ModelMessage[] = [];
// 6 OLD turns (indices 0..5), each with a huge tool output.
for (let i = 0; i < 6; i++) {
msgs.push(userMsg(`old q${i}`));
msgs.push(assistantMsg('looking', `c${i}`, 'getPage'));
msgs.push(toolMsg(`c${i}`, 'getPage', { body: big }));
msgs.push(assistantMsg(`old a${i}`));
}
// 3 small recent turns, then the CURRENT turn with its own huge tool output.
// With REPLAY_KEEP_RECENT_TURNS=4 the last 4 user-turns stay full, so only
// these small recent turns + the current big one are kept full; the 6 old
// turns above fall in the trim region.
for (let i = 0; i < 3; i++) {
msgs.push(userMsg(`recent q${i}`));
msgs.push(assistantMsg(`recent a${i}`));
}
msgs.push(userMsg('current q'));
msgs.push(assistantMsg('looking', 'cR', 'getPage'));
msgs.push(toolMsg('cR', 'getPage', { body: big }));
msgs.push(assistantMsg('current a'));
// Budget large enough that phase-1 tool truncation alone brings it under.
const r = trimHistoryForReplay(msgs, 30_000);
expect(r.trimmed).toBe(true);
const flat = JSON.stringify(r.messages);
// The CURRENT turn's tool output survives in full.
expect(flat).toContain(big);
// Old outputs were truncated with the marker.
expect(flat).toContain(REPLAY_TRUNCATION_MARKER);
// Phase 1 sufficed: the oldest turns were NOT collapsed.
expect(flat).not.toContain(REPLAY_TURN_COLLAPSED_MARKER);
expect(estimateMessagesTokens(r.messages)).toBeLessThan(
estimateMessagesTokens(msgs),
);
});
it('collapses the oldest turns when tool truncation is not enough', () => {
// Many turns with LARGE assistant TEXT (not tool output) so phase 1 can't help.
const bigText = 'слово '.repeat(8_000); // large Cyrillic text per turn
const msgs: ModelMessage[] = [];
for (let i = 0; i < 12; i++) {
msgs.push(userMsg(`q${i}`));
msgs.push(assistantMsg(bigText));
}
const r = trimHistoryForReplay(msgs, 30_000);
expect(r.trimmed).toBe(true);
// Oldest turns collapsed; result fits (best-effort) and is much smaller.
expect(estimateMessagesTokens(r.messages)).toBeLessThan(
estimateMessagesTokens(msgs),
);
// The LAST turn's text is preserved in full (recent turns stay full).
expect(JSON.stringify(r.messages[r.messages.length - 1])).toContain(bigText);
});
it('is deterministic / byte-stable for identical inputs', () => {
const big = 'Y'.repeat(30_000);
const build = (): ModelMessage[] => {
const m: ModelMessage[] = [];
for (let i = 0; i < 10; i++) {
m.push(userMsg(`q${i}`));
m.push(assistantMsg('t', `c${i}`, 'getPage'));
m.push(toolMsg(`c${i}`, 'getPage', { body: big }));
}
return m;
};
const a = trimHistoryForReplay(build(), 15_000);
const b = trimHistoryForReplay(build(), 15_000);
expect(JSON.stringify(a.messages)).toBe(JSON.stringify(b.messages));
});
it('never leaves an unpaired tool-call after collapsing (balanced history)', () => {
const big = 'Z'.repeat(40_000);
const msgs: ModelMessage[] = [];
for (let i = 0; i < 10; i++) {
msgs.push(userMsg(`q${i}`));
msgs.push(assistantMsg('t', `c${i}`, 'getPage'));
msgs.push(toolMsg(`c${i}`, 'getPage', { body: big }));
}
const r = trimHistoryForReplay(msgs, 8_000);
// Count tool-call vs tool-result parts in the trimmed output.
let calls = 0;
let results = 0;
for (const m of r.messages) {
if (!Array.isArray(m.content)) continue;
for (const p of m.content as Array<{ type?: string }>) {
if (p.type === 'tool-call') calls++;
if (p.type === 'tool-result' || p.type === 'tool-error') results++;
}
}
// Every surviving tool-call has a surviving result (collapsing drops BOTH).
expect(calls).toBe(results);
// Collapsed turns carry the marker.
expect(JSON.stringify(r.messages)).toContain(REPLAY_TURN_COLLAPSED_MARKER);
});
it('respects the provider fact: under-budget contextTokens skips trimming', () => {
const big = 'W'.repeat(60_000);
const msgs = [
userMsg('q'),
assistantMsg('t', 'c1', 'getPage'),
toolMsg('c1', 'getPage', { body: big }),
];
// char-estimate is high, but the provider says we are well under budget.
const r = trimHistoryForReplay(msgs, 100_000, 5_000);
expect(r.trimmed).toBe(false);
expect(r.messages).toBe(msgs);
});
});
@@ -0,0 +1,356 @@
/**
* History-replay token budget (#490).
*
* The whole persisted conversation is replayed to the provider on EVERY turn, so
* a long chat eventually exceeds the model's context window and the provider 400s
* on every turn terminally (the chat "bricks"). This module bounds the replayed
* history at REPLAY TIME only: it never mutates what is persisted (the DB stays
* the full record), and its output is a deterministic, byte-stable function of its
* input so the trimmed prefix is identical turn to turn (provider prompt-cache
* friendliness real money on long chats).
*
* The PRIMARY signal is the provider's own fact: `metadata.contextTokens` from the
* last turn. The chars-based {@link estimateTokens} (shared with the client) is
* used only for the DELTA of not-yet-sent messages, to decide WHAT to trim, and as
* the fallback for chats with no usage yet.
*/
import type { ModelMessage } from 'ai';
import { estimateTokens } from '@docmost/token-estimate';
/** Flat default budget when no context window is configured (tokens). */
export const REPLAY_BUDGET_DEFAULT_TOKENS = 100_000;
/** Fraction of a configured context window used as the budget. */
export const REPLAY_BUDGET_WINDOW_FRACTION = 0.7;
/**
* Fraction of the normal budget used for the REACTIVE re-trim after a provider
* context-overflow 400 the preventive estimate under-counted, so cut harder.
*/
export const REPLAY_AGGRESSIVE_FRACTION = 0.5;
/**
* Turns (a user message + its assistant/tool replies) kept FULL at the tail,
* including the current one never trimmed. Older turns are compacted first.
*/
export const REPLAY_KEEP_RECENT_TURNS = 4;
/** Leading chars kept from a truncated old tool output. */
export const REPLAY_TOOL_OUTPUT_HEAD = 800;
/** Trailing chars kept from a truncated old tool output. */
export const REPLAY_TOOL_OUTPUT_TAIL = 300;
/** Marker inserted where an old tool output was truncated for replay. */
export const REPLAY_TRUNCATION_MARKER =
'[…truncated for replay; call the tool again to read the full output]';
/** Marker for a whole old turn collapsed to its text. */
export const REPLAY_TURN_COLLAPSED_MARKER =
'[earlier tool activity omitted for replay]';
export interface ReplayBudget {
/** Token threshold above which replay history is trimmed; `null` = OFF. */
thresholdTokens: number | null;
/** True when the flat default was used (no context window configured). */
usedDefault: boolean;
}
/**
* Resolve the replay budget from the RAW stored `chatContextWindow` (text/number).
* - a positive value -> `min(default, floor(fraction × window))`
* - explicit `0` -> OFF (admin opt-out; `null` threshold)
* - unset/empty/invalid-> the flat default (still protects the installations
* that hit terminal overflow are exactly the ones that never set a window)
*
* Note the raw value is needed because the parsed `chatContextWindow` collapses
* both `0` and unset to `undefined`, which would erase the explicit off-switch.
*/
export function resolveReplayBudget(rawContextWindow: unknown): ReplayBudget {
let n: number | undefined;
if (typeof rawContextWindow === 'number') {
n = rawContextWindow;
} else if (typeof rawContextWindow === 'string') {
const t = rawContextWindow.trim();
n = t === '' ? undefined : Number(t);
}
// Unset / empty / non-numeric / negative -> flat default (the protective case).
if (n === undefined || !Number.isFinite(n) || n < 0) {
return { thresholdTokens: REPLAY_BUDGET_DEFAULT_TOKENS, usedDefault: true };
}
// Explicit 0 -> off-switch.
if (n === 0) {
return { thresholdTokens: null, usedDefault: false };
}
return {
thresholdTokens: Math.min(
REPLAY_BUDGET_DEFAULT_TOKENS,
Math.floor(REPLAY_BUDGET_WINDOW_FRACTION * n),
),
usedDefault: false,
};
}
/**
* True when a provider error is a CONTEXT-OVERFLOW rejection (the prompt exceeds
* the model's window). Providers surface this as an HTTP 400 with a recognizable
* message; match both the status and the message patterns robustly across
* OpenAI-compatible / Anthropic / Gemini wordings, since the exact shape varies.
*/
export function isContextOverflowError(error: unknown): boolean {
const status = extractStatus(error);
const msg = extractMessage(error).toLowerCase();
// Message patterns seen across providers for "prompt too long".
const overflowPattern =
/context (?:length|window)|maximum context|too many tokens|too large for|reduce the length|prompt is too long|input (?:is )?too long|exceeds? the (?:maximum )?(?:context|token)|maximum.*tokens|string too long/;
if (!overflowPattern.test(msg)) return false;
// A 400/413 with an overflow-shaped message is an overflow. Some providers
// omit/rewrite the status, so accept the message match when the status is
// unknown, but reject it for auth/rate-limit statuses that never mean overflow.
if (status === 400 || status === 413) return true;
if (status === 401 || status === 403 || status === 429) return false;
return true;
}
function extractStatus(error: unknown): number | undefined {
if (!error || typeof error !== 'object') return undefined;
const e = error as Record<string, unknown>;
for (const k of ['statusCode', 'status']) {
const v = e[k];
if (typeof v === 'number') return v;
if (typeof v === 'string' && /^\d+$/.test(v)) return Number(v);
}
// Nested (e.g. { response: { status } } / { cause: { statusCode } }).
for (const k of ['response', 'cause', 'data']) {
const nested = e[k];
if (nested && typeof nested === 'object') {
const s = extractStatus(nested);
if (s !== undefined) return s;
}
}
return undefined;
}
function extractMessage(error: unknown): string {
if (error == null) return '';
if (typeof error === 'string') return error;
if (error instanceof Error) {
// Include nested causes (provider libs wrap the real body in `cause`).
const cause = (error as { cause?: unknown }).cause;
return `${error.message} ${cause ? extractMessage(cause) : ''}`;
}
if (typeof error === 'object') {
const e = error as Record<string, unknown>;
const parts: string[] = [];
for (const k of ['message', 'error', 'body', 'responseBody', 'data']) {
const v = e[k];
if (typeof v === 'string') parts.push(v);
else if (v && typeof v === 'object') parts.push(extractMessage(v));
}
return parts.join(' ');
}
return String(error);
}
/** Rough token size of a ModelMessage array via the shared chars estimator. */
export function estimateMessagesTokens(
messages: ReadonlyArray<ModelMessage>,
): number {
let total = 0;
for (const m of messages) {
total += estimateTokens(serializeContent(m.content));
}
return total;
}
function serializeContent(content: unknown): string {
if (typeof content === 'string') return content;
try {
return JSON.stringify(content) ?? '';
} catch {
return '';
}
}
/** Deep JSON string of an arbitrary value, bounded so estimation never throws. */
function stringifyValue(value: unknown): string {
if (typeof value === 'string') return value;
try {
return JSON.stringify(value) ?? String(value);
} catch {
return String(value);
}
}
export interface TrimResult {
messages: ModelMessage[];
/** Whether any trimming was applied. */
trimmed: boolean;
/** Estimated tokens of the returned messages (chars-based). */
estimatedTokens: number;
}
/**
* Bound the replayed history to `budgetTokens`, deterministically. Returns the
* SAME array reference (no copy) when nothing needs trimming, so the common case
* is free and byte-identical. Trimming order (spec #490):
* 1. truncate OLD turns' tool outputs (head+tail + marker) the bulk of the size
* 2. mechanically collapse the OLDEST turns to their text (concatenation, no LLM)
* 3. the current + last {@link REPLAY_KEEP_RECENT_TURNS} turns stay FULL
*
* `budgetTokens === null` disables trimming. `priorContextTokens` (the provider's
* fact from last turn) short-circuits the decision: when it is known and already
* under budget we skip trimming even if the char-estimate is higher (the provider
* count is authoritative). The char-estimate drives WHAT to cut.
*/
export function trimHistoryForReplay(
messages: ModelMessage[],
budgetTokens: number | null,
priorContextTokens?: number,
): TrimResult {
if (budgetTokens == null) {
return { messages, trimmed: false, estimatedTokens: 0 };
}
const estimated = estimateMessagesTokens(messages);
// Decision signal: prefer the provider's fact (last turn's contextTokens) plus
// the estimated delta of the messages appended since; fall back to the pure
// char-estimate for a chat with no usage yet.
const projected =
priorContextTokens != null
? Math.max(priorContextTokens, estimated)
: estimated;
if (projected <= budgetTokens) {
return { messages, trimmed: false, estimatedTokens: estimated };
}
// The tail we always keep full: from the Nth-from-last user message onward.
const boundary = recentBoundaryIndex(messages, REPLAY_KEEP_RECENT_TURNS);
const tail = messages.slice(boundary);
let head = messages.slice(0, boundary).map(cloneMessage);
// Phase 1: truncate old tool outputs.
for (const m of head) {
if (m.role === 'tool') truncateToolMessage(m);
}
let out = [...head, ...tail];
let est = estimateMessagesTokens(out);
if (est <= budgetTokens) {
return { messages: out, trimmed: true, estimatedTokens: est };
}
// Phase 2: collapse the oldest turns (in `head`) to their text, one at a time,
// from the oldest, until we fit or the whole head is collapsed.
const turns = splitTurns(head);
const collapsed: ModelMessage[] = [];
let i = 0;
for (; i < turns.length; i++) {
if (est <= budgetTokens) break;
collapsed.push(...collapseTurn(turns[i]));
// Re-estimate the whole prospective output.
const remaining = turns.slice(i + 1).flat();
out = [...collapsed, ...remaining, ...tail];
est = estimateMessagesTokens(out);
}
// Include any turns we didn't need to collapse.
const remaining = turns.slice(i).flat();
out = [...collapsed, ...remaining, ...tail];
est = estimateMessagesTokens(out);
return { messages: out, trimmed: true, estimatedTokens: est };
}
/** Index of the first message of the Nth-from-last user turn (0 if fewer). */
function recentBoundaryIndex(
messages: ReadonlyArray<ModelMessage>,
keepTurns: number,
): number {
const userIdx: number[] = [];
for (let i = 0; i < messages.length; i++) {
if (messages[i].role === 'user') userIdx.push(i);
}
if (userIdx.length <= keepTurns) return 0;
return userIdx[userIdx.length - keepTurns];
}
/** Split a message list into turns; each turn starts at a `user` message. */
function splitTurns(messages: ModelMessage[]): ModelMessage[][] {
const turns: ModelMessage[][] = [];
for (const m of messages) {
if (m.role === 'user' || turns.length === 0) turns.push([m]);
else turns[turns.length - 1].push(m);
}
return turns;
}
/**
* Collapse a whole turn to its plain text (mechanical concatenation, not an LLM
* summary). Keeps the user message; replaces the assistant/tool messages with a
* single assistant text message = the assistant's concatenated text + a marker
* when tool activity was dropped. Dropping BOTH the tool-call and tool-result
* parts together keeps the rebuilt history balanced (no unpaired calls).
*/
function collapseTurn(turn: ModelMessage[]): ModelMessage[] {
const out: ModelMessage[] = [];
let assistantText = '';
let hadTools = false;
for (const m of turn) {
if (m.role === 'user') {
out.push(m);
} else if (m.role === 'assistant') {
const { text, tools } = extractAssistantText(m.content);
assistantText += text;
hadTools = hadTools || tools;
} else if (m.role === 'tool') {
hadTools = true;
} else {
out.push(m);
}
}
const note =
(assistantText ? assistantText.trimEnd() : '') +
(hadTools
? `${assistantText ? '\n\n' : ''}${REPLAY_TURN_COLLAPSED_MARKER}`
: '');
if (note) out.push({ role: 'assistant', content: note } as ModelMessage);
return out;
}
function extractAssistantText(content: unknown): {
text: string;
tools: boolean;
} {
if (typeof content === 'string') return { text: content, tools: false };
if (!Array.isArray(content)) return { text: '', tools: false };
let text = '';
let tools = false;
for (const part of content) {
const type = (part as { type?: string })?.type;
if (type === 'text') text += (part as { text?: string }).text ?? '';
else if (type === 'tool-call') tools = true;
}
return { text, tools };
}
/** Truncate every tool-result output in a `tool` message to head+tail+marker. */
function truncateToolMessage(message: ModelMessage): void {
const content = message.content;
if (!Array.isArray(content)) return;
for (const part of content) {
const p = part as { type?: string; output?: { type?: string; value?: unknown } };
if (p.type !== 'tool-result' && p.type !== 'tool-error') continue;
if (!p.output) continue;
const raw = stringifyValue(p.output.value);
const budget = REPLAY_TOOL_OUTPUT_HEAD + REPLAY_TOOL_OUTPUT_TAIL;
if (raw.length <= budget + REPLAY_TRUNCATION_MARKER.length) continue;
const truncated =
raw.slice(0, REPLAY_TOOL_OUTPUT_HEAD) +
`\n${REPLAY_TRUNCATION_MARKER}\n` +
raw.slice(raw.length - REPLAY_TOOL_OUTPUT_TAIL);
// Represent the shrunk output as a text output (a valid tool-result output).
p.output = { type: 'text', value: truncated };
}
}
/** Shallow-ish clone so trimming never mutates the caller's (persisted-derived)
* message objects only the OLD region is cloned before it is edited. */
function cloneMessage(m: ModelMessage): ModelMessage {
if (typeof m.content === 'string') return { ...m };
return {
...m,
content: (m.content as unknown[]).map((p) =>
p && typeof p === 'object' ? { ...(p as object) } : p,
),
} as ModelMessage;
}
@@ -0,0 +1,160 @@
import {
wrapInAppToolWithCap,
inAppToolCallCapMs,
type ToolAbortSignalSink,
} from './ai-chat-tools.service';
import type { Tool, ToolCallOptions } from 'ai';
/**
* #487 commit 1 in-app tool race-on-abort + safe-points + per-call cap.
*
* Tests assert the HONEST observable property the spec names "after Stop, NO
* new HTTP/WS call STARTS; an already-started single call may take either
* outcome" against the REAL wrapper mechanism (the composite abort signal it
* publishes on the client + the RACE it runs), NOT a timing-dependent proxy like
* "the write didn't land".
*/
// A minimal stand-in for the client's `toolAbortSignal` field. In production the
// wrapper publishes the composite here and the client's paginateAll /
// mutatePageContent safe-points read it; the fake "tool" below reads it the same
// way, so this exercises the real contract without a live DB / collab socket.
class FakeClient implements ToolAbortSignalSink {
private signal: AbortSignal | null = null;
setToolAbortSignal(signal: AbortSignal | null): void {
this.signal = signal;
}
getToolAbortSignal(): AbortSignal | null {
return this.signal;
}
}
// A ToolCallOptions with just the field the wrapper reads (abortSignal). The AI
// SDK passes a fuller object; the wrapper only spreads it and reads abortSignal.
const opts = (abortSignal?: AbortSignal): ToolCallOptions =>
({ toolCallId: 't1', messages: [], abortSignal }) as unknown as ToolCallOptions;
const tick = (ms = 5) => new Promise((r) => setTimeout(r, ms));
describe('#487 wrapInAppToolWithCap — race-on-abort + safe-points', () => {
it('after Stop, no NEW simulated call starts (multi-call tool)', async () => {
const client = new FakeClient();
const started: number[] = [];
// A multi-call tool that mirrors paginateAll: it consults the client signal
// at a safe-point BEFORE starting each simulated network call.
const multiCall: Tool = {
execute: (async (_args: unknown) => {
for (let i = 0; i < 6; i++) {
// Safe-point: exactly what paginateAll / mutatePageContent do.
client.getToolAbortSignal()?.throwIfAborted();
started.push(i);
await tick(10);
}
return 'done';
}) as unknown as Tool['execute'],
} as Tool;
const wrapped = wrapInAppToolWithCap(multiCall, client, 10_000);
const ac = new AbortController();
const call = (
wrapped.execute as (a: unknown, o: ToolCallOptions) => Promise<unknown>
)({}, opts(ac.signal));
// Let one or two calls start, then Stop.
await tick(12);
ac.abort(new Error('user stop'));
await expect(call).rejects.toThrow(); // wrapper rejects promptly
const startedAtStop = started.length;
// Give the abandoned loser ample time; its next safe-point must throw because
// the (aborted) composite is still published on the client.
await tick(60);
expect(started.length).toBe(startedAtStop);
// It must NOT have run the whole sequence (that would mean Stop did nothing).
expect(started.length).toBeLessThan(6);
});
it('rejects immediately on Stop even if the call never settles (discard loser)', async () => {
const client = new FakeClient();
let settled = false;
const hang: Tool = {
execute: (async () => {
await new Promise(() => undefined); // never resolves
settled = true;
}) as unknown as Tool['execute'],
} as Tool;
const wrapped = wrapInAppToolWithCap(hang, client, 10_000);
const ac = new AbortController();
const call = (
wrapped.execute as (a: unknown, o: ToolCallOptions) => Promise<unknown>
)({}, opts(ac.signal));
await tick(5);
ac.abort();
await expect(call).rejects.toThrow();
expect(settled).toBe(false);
});
it('per-call cap rejects a hung call with no Stop signal', async () => {
const client = new FakeClient();
const hang: Tool = {
execute: (async () => {
await new Promise(() => undefined);
}) as unknown as Tool['execute'],
} as Tool;
// Tiny cap; no options.abortSignal at all (composite = cap only).
const wrapped = wrapInAppToolWithCap(hang, client, 20);
const start = Date.now();
await expect(
(wrapped.execute as (a: unknown, o: ToolCallOptions) => Promise<unknown>)(
{},
opts(undefined),
),
).rejects.toThrow(/per-call cap/);
expect(Date.now() - start).toBeLessThan(2000);
});
it('publishes a composite signal on the client for the duration of the call', async () => {
const client = new FakeClient();
let seenDuringCall: AbortSignal | null = null;
const probe: Tool = {
execute: (async () => {
seenDuringCall = client.getToolAbortSignal();
return 'ok';
}) as unknown as Tool['execute'],
} as Tool;
const wrapped = wrapInAppToolWithCap(probe, client, 10_000);
const ac = new AbortController();
await (
wrapped.execute as (a: unknown, o: ToolCallOptions) => Promise<unknown>
)({}, opts(ac.signal));
expect(seenDuringCall).not.toBeNull();
// The published composite must reflect the turn's Stop signal.
ac.abort();
expect((seenDuringCall as unknown as AbortSignal).aborted).toBe(true);
});
it('a completed call returns its raw result unchanged', async () => {
const client = new FakeClient();
const ok: Tool = {
execute: (async () => ({ items: [1, 2, 3] })) as unknown as Tool['execute'],
} as Tool;
const wrapped = wrapInAppToolWithCap(ok, client, 10_000);
const res = await (
wrapped.execute as (a: unknown, o: ToolCallOptions) => Promise<unknown>
)({}, opts(new AbortController().signal));
expect(res).toEqual({ items: [1, 2, 3] });
});
it('cap is env-tunable with a 2-minute default', () => {
const prev = process.env.AI_CHAT_INAPP_TOOL_CALL_CAP_MS;
delete process.env.AI_CHAT_INAPP_TOOL_CALL_CAP_MS;
expect(inAppToolCallCapMs()).toBe(120_000);
process.env.AI_CHAT_INAPP_TOOL_CALL_CAP_MS = '5000';
expect(inAppToolCallCapMs()).toBe(5000);
process.env.AI_CHAT_INAPP_TOOL_CALL_CAP_MS = 'not-a-number';
expect(inAppToolCallCapMs()).toBe(120_000);
if (prev === undefined) delete process.env.AI_CHAT_INAPP_TOOL_CALL_CAP_MS;
else process.env.AI_CHAT_INAPP_TOOL_CALL_CAP_MS = prev;
});
});
@@ -1,5 +1,5 @@
import { Injectable, Logger } from '@nestjs/common';
import { tool, type Tool } from 'ai';
import { tool, type Tool, type ToolCallOptions } from 'ai';
import { z } from 'zod';
import { User } from '@docmost/db/types/entity.types';
import { TokenService } from '../../auth/services/token.service';
@@ -159,6 +159,129 @@ function __assertClientCallContract(client: DocmostClientLike): void {
* existing service-account `/mcp` path already calls loopback successfully, so
* this works for single-workspace self-host.
*/
/**
* #487: wall-clock cap for a SINGLE in-app tool call, env-tunable via
* `AI_CHAT_INAPP_TOOL_CALL_CAP_MS`. Bounds a read tool that would otherwise
* paginate for minutes and a content write whose collab commit hangs, and is the
* per-call CAP half of the composite abort signal every in-app tool is wrapped
* with (the other half is the turn's Stop signal). Default 2 minutes: generous
* for a legitimate long read/write, tight enough that a stuck call cannot pin the
* turn. The reconcile staleness floor (#487 commit 4) is derived as
* max(2 x this cap, 15 min), so keep this well under that.
*/
export function inAppToolCallCapMs(): number {
const raw = Number(process.env.AI_CHAT_INAPP_TOOL_CALL_CAP_MS);
return Number.isFinite(raw) && raw > 0 ? raw : 120_000;
}
/** #487: the composite signal's reason as an Error (informative thrown value). */
function inAppAbortReason(signal: AbortSignal): Error {
const r = signal.reason;
return r instanceof Error
? r
: new Error(typeof r === 'string' ? r : 'In-app tool call aborted');
}
/**
* The client surface {@link wrapInAppToolWithCap} drives (#487). Both methods are
* OPTIONAL: the real loopback DocmostClient implements them (so a Stop/cap reaches
* its pagination / pre-commit safe-points), but a client that omits them still
* gets the OUTER guarantee the race rejects on abort regardless. This keeps the
* wrapper decoupled from the exact client shape (unit-test doubles need not stub
* the plumbing).
*/
export interface ToolAbortSignalSink {
setToolAbortSignal?(signal: AbortSignal | null): void;
getToolAbortSignal?(): AbortSignal | null;
}
/**
* #487: wrap an in-app tool so a Stop (the turn's `options.abortSignal`) OR the
* per-call wall-clock cap REJECTS the call immediately, and so that SAME
* composite signal reaches the client's pagination / pre-commit safe-points (via
* `client.setToolAbortSignal`) making a Stop stop the NEXT HTTP/WS call from
* starting.
*
* Reuses the RACE pattern of `wrapToolWithCallTimeout` (mcp-clients.service.ts):
* the call is raced against the composite signal, so on abort we reject in the
* SAME tick and DISCARD the loser promise. Its network / collab teardown latency
* therefore never blocks the turn the supersede timeout W=10s (#487 commit 3)
* relies on this abort->settle latency being milliseconds, not a socket teardown.
* Awaiting the client's own signal-into-write path alone would NOT satisfy this
* (the loser could still be tearing down a collab socket).
*
* The composite is SET on the client at entry and deliberately NOT restored on
* unwind: after this wrapper rejects on abort, the ABANDONED loser promise keeps
* running, and its safe-points read the client field leaving the (aborted)
* composite there is exactly what makes the loser's NEXT call throw and stop. The
* next in-app tool call overwrites the field with its own fresh composite before
* any of its safe-points run, so a stale settled signal never leaks forward.
* SINGLE-WRITER by phase-1 assumption see DocmostClientContext.toolAbortSignal
* for the parallel-call caveat (#487).
*
* KNOWN LIMITATION (#487): a write tool that issues SEVERAL sequential collab
* commits can be aborted BETWEEN commits, leaving a partially-applied operation.
* Cancel guarantees "no NEW call starts", NOT "the write didn't land".
*/
export function wrapInAppToolWithCap(
toolDef: Tool,
client: ToolAbortSignalSink,
capMs: number,
): Tool {
const original = toolDef.execute;
if (typeof original !== 'function') return toolDef;
const execute = async (args: unknown, options: ToolCallOptions) => {
const capController = new AbortController();
const timer = setTimeout(() => {
capController.abort(
new Error(`In-app tool call exceeded the ${capMs}ms per-call cap`),
);
}, capMs);
timer.unref?.();
const composite = options?.abortSignal
? AbortSignal.any([options.abortSignal, capController.signal])
: capController.signal;
// Reject the MOMENT the composite fires, independent of whether `original`
// ever settles (a hung collab write / read would otherwise pin the turn). The
// losing `original` is left pending; Promise.race attaches a rejection
// handler to both inputs so a late rejection is never unhandled.
const aborted = new Promise<never>((_, reject) => {
const fail = () => reject(inAppAbortReason(composite));
if (composite.aborted) fail();
else composite.addEventListener('abort', fail, { once: true });
});
// Publish the composite so the client's pagination / pre-commit safe-points
// observe it (see the "not restored on unwind" rationale above). Guarded: a
// client without the plumbing still gets the OUTER race guarantee below.
client.setToolAbortSignal?.(composite);
try {
return await Promise.race([
(original as (a: unknown, o: ToolCallOptions) => Promise<unknown>)(
args,
{ ...options, abortSignal: composite },
),
aborted,
]);
} finally {
clearTimeout(timer);
}
};
return { ...toolDef, execute } as unknown as Tool;
}
/** #487: apply {@link wrapInAppToolWithCap} to every tool in a set. */
export function wrapInAppToolsWithCap(
tools: Record<string, Tool>,
client: ToolAbortSignalSink,
capMs: number,
): Record<string, Tool> {
const out: Record<string, Tool> = {};
for (const [name, t] of Object.entries(tools)) {
out[name] = wrapInAppToolWithCap(t, client, capMs);
}
return out;
}
@Injectable()
export class AiChatToolsService {
private readonly logger = new Logger(AiChatToolsService.name);
@@ -186,7 +309,12 @@ export class AiChatToolsService {
sessionId: string,
workspaceId: string,
aiChatId: string,
): Promise<DocmostClientLike> {
// #487: the returned client also carries the tool-cancellation plumbing
// (setToolAbortSignal/getToolAbortSignal). These are host plumbing, NOT part
// of the tool-execute surface (DocmostClientMethod), so they are surfaced here
// as an intersection rather than by widening that Pick — keeping the
// positional-call drift-guard (#446) scoped to the actual tool methods.
): Promise<DocmostClientLike & ToolAbortSignalSink> {
const apiUrl =
process.env.MCP_DOCMOST_API_URL ||
`http://127.0.0.1:${process.env.PORT || 3000}/api`;
@@ -630,7 +758,15 @@ export class AiChatToolsService {
// dependency and reuses the CASL enforcement already on `client`. When the
// loaded package predates #417 (factory undefined) or the loader is mocked in
// a unit test, signalling is a pure no-op and results are byte-identical.
if (!createCommentSignalTracker) return tools;
// #487: wrap every in-app tool with the race-on-abort + per-call cap guard so
// a Stop / cap rejects immediately AND reaches the client's write/pagination
// safe-points. Applied as the OUTERMOST wrapper (over the comment-signal
// wrapper below) so the race governs the whole call. The client carries the
// per-call composite signal via setToolAbortSignal.
const capMs = inAppToolCallCapMs();
if (!createCommentSignalTracker) {
return wrapInAppToolsWithCap(tools, client, capMs);
}
const tracker = createCommentSignalTracker({
probe: async (pageId: string, sinceMs: number) => {
@@ -659,7 +795,11 @@ export class AiChatToolsService {
},
});
return wrapToolsWithCommentSignal(tools, tracker);
return wrapInAppToolsWithCap(
wrapToolsWithCommentSignal(tools, tracker),
client,
capMs,
);
}
}
@@ -0,0 +1,24 @@
import { type Kysely, sql } from 'kysely';
export async function up(db: Kysely<any>): Promise<void> {
// Chat-level metadata bag (#490). First use: the deferred-tool ACTIVATION set
// (`activatedTools`) is persisted here so it survives across turns — previously
// the set was reset every turn, forcing the model to re-run loadTools and pay a
// fresh round-trip to re-activate the same tools each turn. On load the stored
// set is intersected with the current valid deferred names, so an allowlist /
// role change can never inject a now-nonexistent tool.
//
// jsonb, defaulted to '{}' so every row (incl. pre-migration ones, backfilled
// by the default) is a readable object — the app never has to null-guard the
// bag itself, only individual keys.
await db.schema
.alterTable('ai_chats')
.addColumn('metadata', 'jsonb', (col) =>
col.notNull().defaultTo(sql`'{}'::jsonb`),
)
.execute();
}
export async function down(db: Kysely<any>): Promise<void> {
await db.schema.alterTable('ai_chats').dropColumn('metadata').execute();
}
@@ -1,5 +1,6 @@
import { Injectable, Logger } from '@nestjs/common';
import { InjectKysely } from 'nestjs-kysely';
import { sql } from 'kysely';
import { KyselyDB, KyselyTransaction } from '../../types/kysely.types';
import { dbOrTx } from '../../utils';
import {
@@ -188,6 +189,144 @@ export class AiChatMessageRepo {
return query.returning(this.baseFields).executeTakeFirst();
}
/**
* #487 OWNER terminal write the streamText terminal callback's finalize. Like
* `update` but CONDITIONAL on `status='streaming' OR metadata.finalizeFailed`:
* the owner writes its real content EITHER when the row is still streaming (the
* normal case) OR when a reconcile stamp already flipped it to a terminal status
* but marked `finalizeFailed:true` the owner's real content OVERWRITES that
* placeholder stamp (owner-write priority, #487). A row that is properly terminal
* (no finalizeFailed) is left untouched (undefined) idempotent. The `patch`
* carries the real metadata WITHOUT finalizeFailed, so a successful write CLEARS
* the flag. Returns the updated row, or undefined when nothing matched.
*/
async finalizeOwner(
id: string,
workspaceId: string,
patch: Partial<{
content: string | null;
toolCalls: unknown;
metadata: unknown;
status: string | null;
}>,
trx?: KyselyTransaction,
): Promise<AiChatMessage | undefined> {
const db = dbOrTx(this.db, trx);
return db
.updateTable('aiChatMessages')
.set({ ...(patch as Record<string, unknown>), updatedAt: new Date() })
.where('id', '=', id)
.where('workspaceId', '=', workspaceId)
.where((eb) =>
eb.or([
eb('status', '=', 'streaming'),
eb(sql<string>`(metadata->>'finalizeFailed')`, '=', 'true'),
]),
)
.returning(this.baseFields)
.executeTakeFirst();
}
/**
* #487 RECONCILE status-only stamp settle a stuck 'streaming' row to a
* terminal status WITHOUT the owner's real content (which lived only in the
* dead process's memory — a documented loss). CONDITIONAL on `status='streaming'`
* (never touches an already-terminal row) AND it MERGES `finalizeFailed:true`
* into metadata (preserving the partial `parts` already persisted) so a LATER
* owner-write (finalizeOwner) can still OVERWRITE this placeholder with real
* content, and so `isInterruptResume` can EXCLUDE this row (a reconcile stamp is
* not a genuine user interruption). Returns the updated row, or undefined.
*/
async stampTerminalIfStreaming(
id: string,
workspaceId: string,
status: 'aborted' | 'error' | 'completed',
trx?: KyselyTransaction,
): Promise<AiChatMessage | undefined> {
const db = dbOrTx(this.db, trx);
return db
.updateTable('aiChatMessages')
.set({
status,
metadata: sql`coalesce(metadata, '{}'::jsonb) || jsonb_build_object('finalizeFailed', true)`,
updatedAt: new Date(),
})
.where('id', '=', id)
.where('workspaceId', '=', workspaceId)
.where('status', '=', 'streaming')
.returning(this.baseFields)
.executeTakeFirst();
}
/**
* #487 reconcile clause (b): streaming assistant rows whose linked RUN has
* already reached a terminal status an asymmetry ("run settled / message
* streaming forever") the periodic reconcile heals by stamping the message.
* Returns the message id + its run's terminal status, bounded.
*/
async findStreamingWithTerminalRun(
limit = 200,
// #487: scope to ONE chat for the opportunistic per-turn reconcile (removes
// reconcile latency from the user-visible path); omit for the periodic sweep.
chat?: { chatId: string; workspaceId: string },
): Promise<
Array<{ messageId: string; workspaceId: string; runStatus: string }>
> {
let query = this.db
.selectFrom('aiChatMessages as m')
.innerJoin('aiChatRuns as r', 'r.assistantMessageId', 'm.id')
.select([
'm.id as messageId',
'm.workspaceId as workspaceId',
'r.status as runStatus',
])
.where('m.status', '=', 'streaming')
.where('r.status', 'in', ['succeeded', 'failed', 'aborted']);
if (chat) {
query = query
.where('m.chatId', '=', chat.chatId)
.where('m.workspaceId', '=', chat.workspaceId);
}
return query.limit(limit).execute();
}
/**
* #487 reconcile clause (d) historical-row safety: streaming rows older than
* `staleMs` whose chat has NO active run row (double-gated). Settle them to
* 'aborted' + finalizeFailed (so a late owner-write could still overwrite).
* Returns the count. Used ONLY by the periodic reconcile, never at boot.
*/
async sweepStreamingWithoutActiveRun(
staleMs: number,
trx?: KyselyTransaction,
): Promise<number> {
const db = dbOrTx(this.db, trx);
const staleBefore = new Date(Date.now() - staleMs);
const rows = await db
.updateTable('aiChatMessages as m')
.set({
status: 'aborted',
metadata: sql`coalesce(m.metadata, '{}'::jsonb) || jsonb_build_object('finalizeFailed', true)`,
updatedAt: new Date(),
})
.where('m.status', '=', 'streaming')
.where('m.updatedAt', '<', staleBefore)
.where((eb) =>
eb.not(
eb.exists(
eb
.selectFrom('aiChatRuns as r')
.select('r.id')
.whereRef('r.chatId', '=', 'm.chatId')
.where('r.status', 'in', ['pending', 'running']),
),
),
)
.returning('m.id')
.execute();
return rows.length;
}
/**
* Crash-recovery sweep (#183): flip every assistant row still left in the
* 'streaming' state (a turn that died mid-write before reaching a terminal
@@ -200,13 +339,20 @@ export class AiChatMessageRepo {
* step, so an actively-streaming row never matches; this prevents a fresh
* replica's boot-sweep from aborting a turn another replica is still streaming
* in a multi-instance deploy.
*
* #487: the sweep now ALSO marks `finalizeFailed:true` so a late owner-write can
* overwrite this placeholder with real content (owner-write priority).
*/
async sweepStreaming(trx?: KyselyTransaction): Promise<number> {
const db = dbOrTx(this.db, trx);
const staleBefore = new Date(Date.now() - SWEEP_STREAMING_STALE_MS);
const rows = await db
.updateTable('aiChatMessages')
.set({ status: 'aborted', updatedAt: new Date() })
.set({
status: 'aborted',
metadata: sql`coalesce(metadata, '{}'::jsonb) || jsonb_build_object('finalizeFailed', true)`,
updatedAt: new Date(),
})
.where('status', '=', 'streaming')
.where('updatedAt', '<', staleBefore)
.returning('id')
@@ -143,6 +143,41 @@ export class AiChatRunRepo {
.executeTakeFirst();
}
/**
* #487: CONDITIONAL terminal finalize flip a run to a terminal status and
* stamp `finished_at` ONLY while it is still active (pending|running), mirroring
* the assistant message's `onlyIfStreaming` guard. A double-settle (a late or
* second writer, a supersede applying a zombie's intended, a reconcile stamp)
* matches NOTHING once the row is terminal and is a benign no-op so a terminal
* status can never be clobbered by a later writer (last-writer-wins is gone).
*
* Returns the updated row when it WAS active (this call wrote it), else
* undefined (the row was already terminal another writer won). The caller
* distinguishes the two to resolve the correct settle outcome.
*/
async finalizeIfActive(
id: string,
workspaceId: string,
patch: { status: string; error: string | null },
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,
})
.where('id', '=', id)
.where('workspaceId', '=', workspaceId)
.where('status', 'in', ACTIVE_RUN_STATUSES as unknown as string[])
.returning(this.baseFields)
.executeTakeFirst();
}
/**
* Mark an EXPLICIT stop request on an active run (distinct from a browser
* disconnect, which never stops a run). Stamps `stop_requested_at` ONLY while
@@ -184,6 +219,31 @@ export class AiChatRunRepo {
* sweeps only runs UNTOUCHED past the window. Phase 1 is single-process, so the
* boot path supplies no window.
*/
/**
* #487 reconcile clause (c): active (pending|running) runs UNTOUCHED past
* `staleMs` candidates for "no live runner" abort. Staleness is measured from
* `updated_at` (the LAST-PROGRESS timestamp recordStep bumps it), NOT
* `started_at`, so a legitimate long-running marathon (1125 min of steady
* progress) is never a candidate. The caller filters these against its in-memory
* `active` / zombie maps ("no entry" is the PRIMARY gate a live entry is never
* aborted) before settling any of them. Bounded.
*/
async findStaleActive(
staleMs: number,
limit = 200,
trx?: KyselyTransaction,
): Promise<Array<{ id: string; workspaceId: string; chatId: string }>> {
const db = dbOrTx(this.db, trx);
const staleBefore = new Date(Date.now() - staleMs);
return db
.selectFrom('aiChatRuns')
.select(['id', 'workspaceId', 'chatId'])
.where('status', 'in', ACTIVE_RUN_STATUSES as unknown as string[])
.where('updatedAt', '<', staleBefore)
.limit(limit)
.execute();
}
async sweepRunning(
opts: { staleMs?: number } = {},
trx?: KyselyTransaction,
+3
View File
@@ -606,6 +606,9 @@ export interface AiChats {
// The document the chat was created in (open page at first message). NULL =>
// started outside any document. ON DELETE SET NULL on the page FK.
pageId: string | null;
// Chat-level metadata bag (#490). jsonb, defaulted to '{}'. First key:
// `activatedTools` — the deferred-tool activation set persisted across turns.
metadata: Generated<Json>;
createdAt: Generated<Timestamp>;
updatedAt: Generated<Timestamp>;
deletedAt: Timestamp | null;
@@ -245,6 +245,9 @@ export class AiSettingsService {
// Max context window for the chat header badge denominator. Stored as
// ::text; 0/unset/invalid = no limit (undefined).
chatContextWindow: parsePositiveInt(provider.chatContextWindow),
// RAW stored value (#490): the replay budgeter reads this to distinguish an
// explicit `0` (off-switch) from unset, which parsePositiveInt cannot.
chatContextWindowRaw: provider.chatContextWindow,
// Plain passthrough; getChatModel defaults unset to 'openai-compatible'.
chatApiStyle: provider.chatApiStyle,
// Cheap model id for the anonymous public-share assistant; reuses the chat
@@ -129,6 +129,12 @@ const DEFAULT_MCP_STREAM_TIMEOUT_MS = 60_000;
/** Default total wall-clock cap for ONE external MCP tool call (2 min). */
const DEFAULT_MCP_CALL_TIMEOUT_MS = 120_000;
/**
* Default `bodyTimeout` for the EXTERNAL-MCP SSE transport (10 min) #489.
* Deliberately much LARGER than {@link DEFAULT_MCP_STREAM_TIMEOUT_MS}.
*/
const DEFAULT_MCP_SSE_BODY_TIMEOUT_MS = 600_000;
/**
* SILENCE timeout (ms) for EXTERNAL-MCP transport ONLY. Override with
* `AI_MCP_STREAM_TIMEOUT_MS`; a missing/invalid/non-positive value falls back to
@@ -164,6 +170,26 @@ export function mcpCallTimeoutMs(): number {
return positiveEnv('AI_MCP_CALL_TIMEOUT_MS', DEFAULT_MCP_CALL_TIMEOUT_MS);
}
/**
* `bodyTimeout` (ms) for the EXTERNAL-MCP **SSE** transport ONLY #489. Override
* with `AI_MCP_SSE_BODY_TIMEOUT_MS`; a missing/invalid/non-positive value falls
* back to {@link DEFAULT_MCP_SSE_BODY_TIMEOUT_MS} (10 min).
*
* The SSE transport holds ONE long-lived response body open across many tool
* calls, so undici's `bodyTimeout` (time between body bytes) counts the LEGITIMATE
* silence BETWEEN calls, not just a hung single call. At the tight HTTP silence
* timeout ({@link mcpStreamTimeoutMs}, 1 min) a normal >1-min gap between the
* model's tool calls would break the SSE socket, and the cache would then serve a
* dead client until TTL. So the SSE transport gets its OWN, RAISED bodyTimeout;
* the per-call total cap ({@link mcpCallTimeoutMs}) still bounds a single stuck
* call, and the app-level transport-error retry heals a socket that does break.
* The HTTP (streamable) transport keeps the tight timeout it opens a fresh
* request per call, so idle-between-calls does not apply there.
*/
export function mcpSseBodyTimeoutMs(): number {
return positiveEnv('AI_MCP_SSE_BODY_TIMEOUT_MS', DEFAULT_MCP_SSE_BODY_TIMEOUT_MS);
}
/**
* undici `Agent` options for streaming AI traffic the (generous, finite)
* silence timeouts plus the keep-alive recycle window. Shared by the chat
@@ -105,6 +105,10 @@ export interface ResolvedAiConfig extends Partial<AiProviderSettings> {
// Max context window in tokens; surfaced to the chat header badge as the
// "current / max" denominator. 0/unset = no limit.
chatContextWindow?: number;
// RAW stored context window (::text), BEFORE parsePositiveInt collapses `0` and
// unset to `undefined`. The #490 replay budgeter needs the raw value to honor an
// explicit `0` off-switch distinctly from "unset -> flat default".
chatContextWindowRaw?: string | number;
// Cheap model id for the public-share assistant; reuses the chat creds.
publicShareChatModel?: string;
// Agent-role id whose persona the public-share assistant adopts (empty/unset
@@ -0,0 +1,305 @@
import { Kysely } from 'kysely';
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 { AiChatRunService } from '../../src/core/ai-chat/ai-chat-run.service';
import {
getTestDb,
destroyTestDb,
createWorkspace,
createUser,
createChat,
createMessage,
} from './db';
/**
* #487 commit 4 bidirectional reconcile + owner-write priority, real SQL.
*
* Proves the OBSERVABLE recovery properties against docmost_test:
* - the CONDITIONAL owner-write beats a reconcile stamp, and a stamp never
* clobbers a proper terminal row;
* - a LATE owner-finalize with real content OVERWRITES a reconcile 'aborted'
* stamp (finalizeFailed);
* - each reconcile clause (b message<-run, c stale-run, d historical row) settles
* the stuck row/run, and a LIVE run entry is never touched;
* - the "kill DB on finish" recovery: after the DB comes back, neither the
* message row nor the run row stays stuck.
*/
describe('#487 reconcile + owner-write priority [integration]', () => {
let db: Kysely<any>;
let messageRepo: AiChatMessageRepo;
let runRepo: AiChatRunRepo;
let runService: AiChatRunService;
let workspaceId: string;
let userId: string;
beforeAll(async () => {
db = getTestDb();
messageRepo = new AiChatMessageRepo(db as any);
runRepo = new AiChatRunRepo(db as any);
runService = new AiChatRunService(runRepo, { isCloud: () => false } as never);
workspaceId = (await createWorkspace(db)).id;
userId = (await createUser(db, workspaceId)).id;
});
afterAll(async () => {
await destroyTestDb();
});
const newChat = async () =>
(await createChat(db, { workspaceId, creatorId: userId })).id;
const metaOf = async (id: string): Promise<Record<string, unknown> | null> => {
const row = await messageRepo.findById(id, workspaceId);
return (row?.metadata as Record<string, unknown> | null) ?? null;
};
it('owner finalizeOwner writes a streaming row and CLEARS finalizeFailed', async () => {
const chatId = await newChat();
const m = await createMessage(db, {
workspaceId,
chatId,
role: 'assistant',
status: 'streaming',
metadata: { parts: [] },
});
const wrote = await messageRepo.finalizeOwner(m.id, workspaceId, {
content: 'final answer',
status: 'completed',
metadata: { parts: [{ type: 'text', text: 'final answer' }] },
} as never);
expect(wrote!.status).toBe('completed');
expect((await metaOf(m.id))?.finalizeFailed).toBeUndefined();
});
it('a reconcile stamp NEVER clobbers a proper terminal row (finalizeOwner is a no-op there)', async () => {
const chatId = await newChat();
const m = await createMessage(db, {
workspaceId,
chatId,
role: 'assistant',
status: 'completed',
content: 'real',
metadata: { parts: [] },
});
// The reconcile stamp is onlyIfStreaming -> no-op on a completed row.
const stamped = await messageRepo.stampTerminalIfStreaming(
m.id,
workspaceId,
'aborted',
);
expect(stamped).toBeUndefined();
expect((await messageRepo.findById(m.id, workspaceId))!.status).toBe(
'completed',
);
});
it('LATE owner-finalize with real content OVERWRITES a reconcile aborted stamp', async () => {
const chatId = await newChat();
const m = await createMessage(db, {
workspaceId,
chatId,
role: 'assistant',
status: 'streaming',
metadata: { parts: [{ type: 'text', text: 'partial' }] },
});
// Reconcile stamps it aborted + finalizeFailed (final text lived only in mem).
const stamped = await messageRepo.stampTerminalIfStreaming(
m.id,
workspaceId,
'aborted',
);
expect(stamped!.status).toBe('aborted');
expect((await metaOf(m.id))?.finalizeFailed).toBe(true);
// A LATE owner-write (finalizeFailed=true satisfies the OR) overwrites it with
// real content, clearing the flag — owner-write priority.
const wrote = await messageRepo.finalizeOwner(m.id, workspaceId, {
content: 'the real final answer',
status: 'completed',
metadata: { parts: [{ type: 'text', text: 'the real final answer' }] },
} as never);
expect(wrote!.status).toBe('completed');
expect(wrote!.content).toBe('the real final answer');
expect((await metaOf(m.id))?.finalizeFailed).toBeUndefined();
});
it('clause (c): a stale active run with NO live entry -> aborted; a LIVE entry is untouched', async () => {
// Stale run, NOT owned by this replica (no entry) -> reconcile aborts it.
const staleChat = await newChat();
const stale = await runRepo.insert({
chatId: staleChat,
workspaceId,
createdBy: userId,
status: 'running',
});
await db
.updateTable('aiChatRuns')
.set({ updatedAt: new Date(Date.now() - 60 * 60 * 1000) })
.where('id', '=', stale.id)
.execute();
// A live run OWNED by this replica (beginRun registers an in-memory entry),
// ALSO backdated stale — the "no entry" primary gate must protect it.
const liveChat = await newChat();
const live = await runService.beginRun({
chatId: liveChat,
workspaceId,
userId,
});
await db
.updateTable('aiChatRuns')
.set({ updatedAt: new Date(Date.now() - 60 * 60 * 1000) })
.where('id', '=', live.runId)
.execute();
const aborted = await runService.reconcileStaleRuns(15 * 60 * 1000);
expect(aborted).toBeGreaterThanOrEqual(1);
expect((await runRepo.findById(stale.id, workspaceId))!.status).toBe(
'aborted',
);
// The live entry is NEVER aborted, however stale its row looks.
expect((await runRepo.findById(live.runId, workspaceId))!.status).toBe(
'running',
);
expect(runService.isLocallyActive(live.runId)).toBe(true);
// cleanup the live run
await runService.finalizeRun(live.runId, workspaceId, 'aborted');
});
it('clause (b): a streaming message whose RUN is terminal is stamped by run status (succeeded -> aborted, NOT completed-empty)', async () => {
const chatId = await newChat();
const msg = await createMessage(db, {
workspaceId,
chatId,
role: 'assistant',
status: 'streaming',
metadata: { parts: [] },
});
// A SUCCEEDED run linked to the still-streaming message (the asymmetry).
const run = await runRepo.insert({
chatId,
workspaceId,
createdBy: userId,
status: 'running',
assistantMessageId: msg.id,
});
await runRepo.finalizeIfActive(run.id, workspaceId, {
status: 'succeeded',
error: null,
});
const stuck = await messageRepo.findStreamingWithTerminalRun();
const mine = stuck.find((s) => s.messageId === msg.id);
expect(mine?.runStatus).toBe('succeeded');
// Reconcile clause (b): succeeded run -> message 'aborted' (NOT 'completed'),
// the final text lived only in memory (documented loss), +finalizeFailed.
const status = mine!.runStatus === 'failed' ? 'error' : 'aborted';
await messageRepo.stampTerminalIfStreaming(msg.id, workspaceId, status);
const row = await messageRepo.findById(msg.id, workspaceId);
expect(row!.status).toBe('aborted');
expect((row!.metadata as Record<string, unknown>).finalizeFailed).toBe(true);
});
it('clause (d): a stale streaming row with NO active run on the chat -> aborted+finalizeFailed', async () => {
const chatId = await newChat();
const msg = await createMessage(db, {
workspaceId,
chatId,
role: 'assistant',
status: 'streaming',
metadata: { parts: [] },
});
await db
.updateTable('aiChatMessages')
.set({ updatedAt: new Date(Date.now() - 60 * 60 * 1000) })
.where('id', '=', msg.id)
.execute();
const swept = await messageRepo.sweepStreamingWithoutActiveRun(
15 * 60 * 1000,
);
expect(swept).toBeGreaterThanOrEqual(1);
const row = await messageRepo.findById(msg.id, workspaceId);
expect(row!.status).toBe('aborted');
expect((row!.metadata as Record<string, unknown>).finalizeFailed).toBe(true);
});
it('clause (d) is DOUBLE-GATED: a stale streaming row WITH an active run on the chat is left alone', async () => {
const chatId = await newChat();
const msg = await createMessage(db, {
workspaceId,
chatId,
role: 'assistant',
status: 'streaming',
metadata: { parts: [] },
});
await db
.updateTable('aiChatMessages')
.set({ updatedAt: new Date(Date.now() - 60 * 60 * 1000) })
.where('id', '=', msg.id)
.execute();
// An ACTIVE run on the same chat -> clause (d) must NOT touch the message.
const run = await runRepo.insert({
chatId,
workspaceId,
createdBy: userId,
status: 'running',
});
await messageRepo.sweepStreamingWithoutActiveRun(15 * 60 * 1000);
expect((await messageRepo.findById(msg.id, workspaceId))!.status).toBe(
'streaming',
);
await runRepo.finalizeIfActive(run.id, workspaceId, {
status: 'aborted',
error: null,
});
});
it('"kill DB on finish" recovery: after the DB is back, reconcile leaves NEITHER the row nor the run stuck', async () => {
// Simulate a process that seeded the assistant row + run, then died before
// finalizing EITHER (a mid-turn crash): a streaming message + a running run,
// both stale, with no in-memory entry (fresh service = fresh maps).
const chatId = await newChat();
const msg = await createMessage(db, {
workspaceId,
chatId,
role: 'assistant',
status: 'streaming',
metadata: { parts: [{ type: 'text', text: 'partial' }] },
});
const run = await runRepo.insert({
chatId,
workspaceId,
createdBy: userId,
status: 'running',
assistantMessageId: msg.id,
});
await db
.updateTable('aiChatRuns')
.set({ updatedAt: new Date(Date.now() - 60 * 60 * 1000) })
.where('id', '=', run.id)
.execute();
await db
.updateTable('aiChatMessages')
.set({ updatedAt: new Date(Date.now() - 60 * 60 * 1000) })
.where('id', '=', msg.id)
.execute();
// Reconcile (as the periodic job would): (c) aborts the orphan run, then
// (b) settles the message from the now-terminal run.
await runService.reconcileStaleRuns(15 * 60 * 1000);
const stuck = await messageRepo.findStreamingWithTerminalRun();
for (const s of stuck) {
const status = s.runStatus === 'failed' ? 'error' : 'aborted';
await messageRepo.stampTerminalIfStreaming(s.messageId, s.workspaceId, status);
}
// Neither is stuck: the run is terminal AND the message is terminal.
expect((await runRepo.findById(run.id, workspaceId))!.status).toBe('aborted');
const row = await messageRepo.findById(msg.id, workspaceId);
expect(row!.status).toBe('aborted');
expect((row!.metadata as Record<string, unknown>).finalizeFailed).toBe(true);
});
});
@@ -281,6 +281,52 @@ describe('AiChatRun durable lifecycle [integration]', () => {
});
});
it('#487 finalizeIfActive is CONDITIONAL: a late terminal write cannot clobber the settled status (real SQL)', async () => {
const c = (await createChat(db, { workspaceId, creatorId: userId })).id;
const run = await runRepo.insert({
chatId: c,
workspaceId,
createdBy: userId,
status: 'running',
});
// First terminal write: the run IS active, so it flips + returns the row.
const first = await runRepo.finalizeIfActive(run.id, workspaceId, {
status: 'succeeded',
error: null,
});
expect(first!.status).toBe('succeeded');
expect(first!.finishedAt).toBeTruthy();
// A late/second writer tries to flip it to 'aborted' — the WHERE status IN
// ('pending','running') guard matches NOTHING now, so it is a benign no-op.
const second = await runRepo.finalizeIfActive(run.id, workspaceId, {
status: 'aborted',
error: 'late clobber attempt',
});
expect(second).toBeUndefined();
// The persisted terminal status is UNCHANGED — last-writer-wins is gone.
const row = await runRepo.findById(run.id, workspaceId);
expect(row!.status).toBe('succeeded');
expect(row!.error).toBeNull();
});
it('#487 double-settle through the service collapses to one write at the SQL gate', async () => {
const c = (await createChat(db, { workspaceId, creatorId: userId })).id;
const handle = await service.beginRun({ chatId: c, workspaceId, userId });
// First settle writes 'aborted' via the conditional write.
await service.finalizeRun(handle.runId, workspaceId, 'aborted');
// A late safety-net settle to 'error' is a no-op (row already terminal).
await service.finalizeRun(handle.runId, workspaceId, 'error', 'late');
const row = await runRepo.findById(handle.runId, workspaceId);
expect(row!.status).toBe('aborted');
expect(service.isLocallyActive(handle.runId)).toBe(false);
expect(service.hasZombie(handle.runId)).toBe(false);
});
it('sweepRunning() with NO args (boot sweep / variant C) aborts even a FRESH running run', async () => {
// F1/DECISION C at the SQL level: the unconditional boot sweep has NO
// staleness window, so a run updated just now (a fast restart) is settled too
+174 -78
View File
@@ -8,19 +8,35 @@ real pain (a "which tools fail most?" analysis that confidently answered
Read the **Gotchas** section before you trust any error count.
> **TWO ERAS — check the marker first.** The `tool_calls` shape changed in **#490
> (trace v2)**. A row written by v2 carries `metadata.toolTraceVersion = 2`; older
> rows have no such key. The two shapes store DIFFERENT things (v2 dropped the tool
> OUTPUT from the trace), so **every query below is dual-shape** — branch on the
> marker. **Never compare an aggregate or trend across the era boundary**: a metric
> jump on the cut-over week is an artifact of the shape change, not a behavior
> change.
## TL;DR
- Agent chats live in Postgres, DB `docmost`, tables `ai_chat_*`.
- Each tool invocation is stored as **two** array elements (a `tool-call` part and
a `tool-result` part), so naive counting double-counts.
- **A tool that *throws* writes no result part.** Since the #407 fix its error is
persisted as a dedicated `{toolName, error}` element in `tool_calls` (queryable +
replayed to the model). **Rows written before #407 still drop it** — the error is
nowhere in the DB and shows only in the live UI. So `isError` / `success=false`
scans under-report by design, and pre-#407 thrown errors are invisible.
- To find where agents fail: (1) soft-failure markers in `tool_calls`, (2) the new
`error` field for thrown errors (new rows) / the orphan-gap proxy (old rows),
(3) server logs / the live UI for full stack traces beyond the truncated message.
- **Era marker:** `metadata.toolTraceVersion = 2` ⇒ v2 (#490) row; absent ⇒ legacy row.
- Each tool invocation is stored as **two** consecutive array elements — a
`tool-call` part then an OUTCOME part — so naive counting double-counts.
- **v2 (#490):** outcome is `{toolName, ok: true}` on success, or
`{toolName, error, kind: 'thrown'|'interrupted'}` on failure. The tool **OUTPUT
is NOT in `tool_calls`** any more — it lives once in `metadata.parts` (this
removed a hundreds-of-MB-per-run write duplication). Soft-failure analysis
therefore reads `metadata.parts`, not `tool_calls`.
- **legacy:** outcome is `{toolName, output}` on success; a **thrown** failure is
a `{toolName, error}` element **only on rows after #407**, and is dropped
entirely (silent orphan) on pre-#407 rows.
- **A tool that *throws* writes no result part.** In v2 it is a
`{error, kind:'thrown'}` element; an interrupted/aborted call is a distinct
`{error, kind:'interrupted'}`. `isError`/`success=false` scans read the *output*
and so under-report thrown failures in every era.
- To find where agents fail: (1) soft-failure markers in `metadata.parts` outputs
(v2) / `tool_calls` outputs (legacy), (2) the `error`/`kind` fields for thrown
failures (v2 + post-#407), (3) server logs / the live UI for full stack traces.
## Where the data lives
@@ -53,33 +69,67 @@ are rows in `workspaces`, not separate deployments.
separate `tool` role), `content` (text), `tool_calls` (jsonb array), `metadata`
(jsonb, holds run `error` + rendered `parts`), `status`, `tsv` (full-text index).
## Era marker — check this before every query
```sql
-- how many rows are in each era?
SELECT COALESCE((metadata->>'toolTraceVersion'), 'legacy') AS era, count(*)
FROM ai_chat_messages
WHERE role = 'assistant' AND jsonb_typeof(tool_calls) = 'array'
GROUP BY 1 ORDER BY 2 DESC;
```
- `toolTraceVersion = '2'`**v2** (#490): outcome flags, **no output in the trace**.
- `NULL` (`'legacy'`) → pre-#490: outcome carries the tool `output` inline.
**Do not trend a metric across the cut-over.** The shape change alone shifts counts
(e.g. "elements with `output`" collapses to zero for v2), so a week that straddles
the boundary shows an artifact, not a behavior change. Segment by era, or restrict to
one era, before comparing.
## How tool calls are stored — READ THIS
Tool calls are **not** one-object-per-call. Each logical invocation is split into
two consecutive elements of the `tool_calls` array:
two consecutive elements of the `tool_calls` array — a **call** then an **outcome**.
The outcome shape is era-dependent:
```text
index 0: { "toolName": "getPage", "input": { "pageId": "…" } } ← tool-call (has input, NO output)
index 1: { "toolName": "getPage", "output": { … } } ← tool-result (has output, NO input)
# v2 (#490) — metadata.toolTraceVersion = 2
index 0: { "toolName":"getPage", "input":{...} } ← call (has input)
index 1: { "toolName":"getPage", "ok":true } ← success (NO output here)
or : { "toolName":"getPage", "error":"…", "kind":"thrown" } ← threw
or : { "toolName":"getPage", "error":"…", "kind":"interrupted" } ← aborted mid-step
# legacy — no toolTraceVersion
index 0: { "toolName":"getPage", "input":{...} } ← call (has input, NO output)
index 1: { "toolName":"getPage", "output":{...} } ← success (has output)
or : { "toolName":"getPage", "error":"…" } ← threw (post-#407 only)
```
The keys that appear on an element are `toolName`, `input`, `output`, and — for a
**thrown** failure on rows written after the #407 fix — `error` (the tool's error
message; see the "Hard failures" section below). There is no `state`, no `errorText`,
no `type`. On pre-#407 rows a thrown failure has NO paired result element at all
(silent orphan). Consequences:
The keys that can appear: `toolName`, `input` (call), and on the outcome — **v2:**
`ok` **or** `error`+`kind`; **legacy:** `output` **or** (post-#407) `error`. There is
no `state`, no `errorText`, no `type` in `tool_calls` (those live on `metadata.parts`).
Consequences:
1. **Real invocation count = elements that have `output` or `error`.** Counting every
element double-counts (you get ~2× and a spurious "~50% of every tool has no output").
2. **Pairing:** a call = a `tool-call` part followed by its result part. A success
carries `output`; a thrown failure (post-#407) carries `error` instead. Both carry
`toolName`, so you can group by tool on either.
1. **Real invocation count** — count the OUTCOME elements, not every element (else you
double-count): **v2** = elements with `ok` or `error`; **legacy** = elements with
`output` or `error`.
2. **Pairing:** a call (`input`) is followed by its outcome. `toolName` is on both, so
you can group by tool on either. In v2 the `kind` field separates a real hard-fail
(`thrown`) from an aborted call (`interrupted`) — a distinction legacy rows cannot
make (both are orphans; see below).
3. **The tool OUTPUT is only in `metadata.parts` on v2 rows.** To inspect what a tool
returned (soft-error markers, page bodies) on a v2 row, read the parts
(`part->>'type' LIKE 'tool-%'`, `part->>'state' = 'output-available'`, `part->'output'`),
not `tool_calls`.
## The two classes of failure (and which the DB can see)
### 1. Soft failures — tool RAN and returned an error-shaped result → PERSISTED ✅
These are visible in the `tool-result` `output`. The marker differs per tool:
These are visible in the tool `output`**on v2 rows in `metadata.parts`** (the
`output-available` part's `output`), on **legacy rows in the `tool_calls` outcome
element's `output`**. The marker differs per tool:
| Tool(s) | Error marker in `output` |
| --- | --- |
@@ -91,37 +141,32 @@ These are visible in the `tool-result` `output`. The marker differs per tool:
Note `editPageText` returns `failed: []` on success — filtering on the *presence*
of the key gives false positives; filter on **non-empty**.
### 2. Hard failures — tool THREW → NOW PERSISTED ✅ (since the #407 fix)
### 2. Hard failures — tool THREW → PERSISTED ✅
When a tool throws (the classic one is `patchNode` / `insertNode` / `tableUpdateCell`
`Failed to encode document to Yjs (fromJSON): Unknown node type: undefined`), the
runtime still writes **no `tool-result` part** — the failure is an ai@6 `tool-error`
content part instead. **Since the #407 fix, that error is persisted**: `serializeSteps`
appends a dedicated element `{toolName, error: "<message>"}` right after the failed
call, mirroring how a successful `{toolName, output}` element is appended. So a thrown
error now leaves a queryable `error` field carrying its (truncated) reason, and the
same real text is replayed to the model on the next turn (an `output-error` part with
the real `errorText`, no longer the `'Tool call did not complete.'` placeholder).
runtime writes **no `tool-result` part** — the failure is an ai@6 `tool-error` content
part. How that lands in `tool_calls` depends on the era:
**Cutover caveat — old rows keep the old blind shape.** Rows written **before** this
change have the two-part shape (`call` + `output` only) and simply **drop** thrown
errors, leaving a silent **orphan** (a `call` with no `output` *and* no `error`). Rows
written **after** the fix additionally carry the `error` element. So:
- **v2 (#490):** a `{toolName, error, kind:'thrown'}` outcome element. An interrupted /
aborted mid-step call is a **distinct** `{toolName, error:'Tool call did not
complete.', kind:'interrupted'}` element — so you can tell a real hard-fail from an
abort **directly, without the orphan heuristic**. Query `kind = 'thrown'`.
- **post-#407 legacy:** a `{toolName, error}` element (no `kind`) right after the call.
- **pre-#407 legacy:** the error is **dropped** — a silent **orphan** (a `call` with no
`output` *and* no `error`).
- **New rows:** query the `error` field directly (see the hard-error query below) — no
orphan heuristic needed for thrown failures.
- **Old rows (pre-#407):** the only DB-side proxy is still an **orphan**: a `tool-call`
part with no matching `tool-result` *and* no `error`. Orphans also appear when a run
is **aborted** mid-flight (server restart), so a high-volume tool (`createComment`,
`searchInPage`, `Search_web_search`) shows orphans from aborts, not real errors on
old rows. Treat the orphan gap as an *upper bound*, and cross-check the tool: a gap on
a structural editor (`patchNode`, `insertNode`, `updatePageJson`, `transformPage`) is
almost certainly a thrown Yjs-encode error; a gap on `createComment` is mostly aborts.
The same real error text is replayed to the model on the next turn (an `output-error`
part with the real `errorText`, from `metadata.parts`), in every era.
A note on the aborted-call fallback: a call with **neither** a result **nor** a
`tool-error` (genuinely interrupted mid-step) still replays with the
`'Tool call did not complete.'` placeholder and persists as an orphan — that path is
unchanged, and is distinct from a real thrown error, which now carries `error`.
**Cutover caveat.** Only pre-#407 legacy rows need the orphan proxy: an orphan is a
`tool-call` with no matching outcome. Orphans there also appear when a run is **aborted**
mid-flight (server restart), so a high-volume tool (`createComment`, `searchInPage`,
`Search_web_search`) shows orphans from aborts, not real errors. Treat the orphan gap as
an *upper bound* and cross-check the tool: a gap on a structural editor (`patchNode`,
`insertNode`, `updatePageJson`, `transformPage`) is almost certainly a thrown Yjs-encode
error; a gap on `createComment` is mostly aborts. **On v2 rows this ambiguity is gone**
`kind` labels each outcome.
### 3. Run-level failures → `ai_chat_runs`
@@ -134,22 +179,34 @@ the wild: `Run interrupted by a server restart.` (aborts) and
Run all of these via `docker exec gitmost-postgresql psql -U docmost -d docmost -P pager=off -c "…"`.
**Real invocation count per tool** (result parts only — the correct denominator):
**Real invocation count per tool** (outcome parts only — the correct denominator).
Dual-shape: a v2 outcome has `ok` or `error`; a legacy outcome has `output` or `error`:
```sql
SELECT elem->>'toolName' AS tool, count(*) AS calls
FROM ai_chat_messages m, jsonb_array_elements(m.tool_calls) elem
WHERE jsonb_typeof(m.tool_calls) = 'array' AND elem ? 'output'
WHERE jsonb_typeof(m.tool_calls) = 'array'
AND (elem ? 'ok' OR elem ? 'output' OR elem ? 'error')
GROUP BY 1 ORDER BY 2 DESC;
```
**Soft errors per tool** (everything the DB can honestly see):
**Soft errors per tool.** The soft-error marker lives in the tool OUTPUT — which on
**v2 rows is in `metadata.parts`**, on **legacy rows is in the `tool_calls` outcome
element**. This query UNIONs both eras, projecting each output as `o`:
```sql
WITH res AS (
-- v2 (#490): output is in metadata.parts (output-available tool parts)
SELECT part->>'type' AS tool, part->'output' AS o
FROM ai_chat_messages m, jsonb_array_elements(m.metadata->'parts') part
WHERE (m.metadata->>'toolTraceVersion') = '2'
AND part->>'type' LIKE 'tool-%' AND part->>'state' = 'output-available'
UNION ALL
-- legacy: output is inline in the tool_calls outcome element
SELECT elem->>'toolName' AS tool, elem->'output' AS o
FROM ai_chat_messages m, jsonb_array_elements(m.tool_calls) elem
WHERE jsonb_typeof(m.tool_calls) = 'array' AND elem ? 'output'
WHERE (m.metadata->>'toolTraceVersion') IS NULL
AND jsonb_typeof(m.tool_calls) = 'array' AND elem ? 'output'
)
SELECT tool, count(*) AS calls,
sum(COALESCE(
@@ -167,13 +224,23 @@ FROM res GROUP BY tool HAVING sum(COALESCE(
ORDER BY soft_errors DESC;
```
**`editPageText` failure reasons** (the most common real agent mistake — bad `find`):
Note the v2 `tool` label is the part type (`tool-editPageText`); strip the `tool-`
prefix if you join it against the legacy `toolName`.
**`editPageText` failure reasons** (the most common real agent mistake — bad `find`).
Same dual-shape output source:
```sql
WITH res AS (
SELECT part->'output' AS o
FROM ai_chat_messages m, jsonb_array_elements(m.metadata->'parts') part
WHERE (m.metadata->>'toolTraceVersion') = '2'
AND part->>'type' = 'tool-editPageText' AND part->>'state' = 'output-available'
UNION ALL
SELECT elem->'output' AS o
FROM ai_chat_messages m, jsonb_array_elements(m.tool_calls) elem
WHERE jsonb_typeof(m.tool_calls) = 'array'
WHERE (m.metadata->>'toolTraceVersion') IS NULL
AND jsonb_typeof(m.tool_calls) = 'array'
AND elem->>'toolName' = 'editPageText' AND elem ? 'output'
)
SELECT f->>'reason' AS reason, count(*)
@@ -182,30 +249,43 @@ WHERE jsonb_typeof(o->'failed') = 'array'
GROUP BY 1 ORDER BY 2 DESC;
```
**Hard errors — persisted `error` field per tool (NEW rows, since #407)** — thrown
tool failures now carry their real reason, so query them directly:
**Hard errors — persisted `error` field per tool (v2 + post-#407 rows)** — thrown tool
failures carry their real reason, so query them directly. On **v2** rows exclude the
`interrupted` kind so an aborted call is not counted as a hard-fail:
```sql
SELECT elem->>'toolName' AS tool, count(*) AS thrown_errors,
min(elem->>'error') AS sample_error
FROM ai_chat_messages m, jsonb_array_elements(m.tool_calls) elem
WHERE jsonb_typeof(m.tool_calls) = 'array' AND elem ? 'error'
-- v2 rows label the kind; a legacy error element has no kind (count it).
AND COALESCE(elem->>'kind', 'thrown') = 'thrown'
GROUP BY 1 ORDER BY 2 DESC;
```
Aborted mid-step calls on v2 rows are a distinct, directly countable population:
```sql
SELECT elem->>'toolName' AS tool, count(*) AS interrupted
FROM ai_chat_messages m, jsonb_array_elements(m.tool_calls) elem
WHERE jsonb_typeof(m.tool_calls) = 'array' AND elem->>'kind' = 'interrupted'
GROUP BY 1 ORDER BY 2 DESC;
```
**Hard-error proxy for OLD rows (pre-#407) — orphan gap per tool, WITH a spread column**
(call parts minus result parts, plus how many distinct chats the gap is spread across).
This covers rows written before thrown errors were persisted; on new rows a thrown
failure now has its own `error` element (use the query above) and an orphan means only
a genuinely aborted mid-step call:
(call parts minus outcome parts, plus how many distinct chats the gap is spread across).
This is needed ONLY for pre-#407 legacy rows (v2 and post-#407 rows carry the error /
`kind` directly — use the queries above). The `WHERE` restricts to the legacy era so v2
rows (where an `ok` outcome is not an `output`) never produce phantom orphans:
```sql
WITH parts AS (
SELECT m.chat_id, elem->>'toolName' AS tool,
(elem ? 'input' AND NOT (elem ? 'output')) AS is_call,
(elem ? 'output' OR elem ? 'error') AS is_result
(elem ? 'input' AND NOT (elem ? 'output') AND NOT (elem ? 'ok')) AS is_call,
(elem ? 'output' OR elem ? 'error' OR elem ? 'ok') AS is_result
FROM ai_chat_messages m, jsonb_array_elements(m.tool_calls) elem
WHERE jsonb_typeof(m.tool_calls) = 'array' AND m.role = 'assistant'
AND (m.metadata->>'toolTraceVersion') IS NULL
),
per_chat AS (
SELECT tool, chat_id, sum(is_call::int) - sum(is_result::int) AS gap
@@ -261,11 +341,21 @@ WHERE tsv @@ websearch_to_tsquery('english', 'some phrase') LIMIT 20;
## Don't blow up your context
A single `tool_calls` row can be **300–400 KB** (results embed full page content and
search payloads). Never `SELECT tool_calls` (or `jsonb_pretty(tool_calls)`) raw.
Always project just the keys you need and truncate:
Tool outputs embed full page content and search payloads (hundreds of KB per row).
On **legacy** rows they are in `tool_calls`; on **v2** rows they moved to
`metadata->'parts'` (the `tool_calls` trace itself is now small). Never `SELECT
tool_calls` / `metadata` (or `jsonb_pretty(...)`) raw — project just the keys you need
and truncate:
```sql
-- v2: outputs live in metadata.parts
SELECT part->>'type',
left(regexp_replace((part->'output')::text, '\s+', ' ', 'g'), 200)
FROM ai_chat_messages m, jsonb_array_elements(m.metadata->'parts') part
WHERE (m.metadata->>'toolTraceVersion') = '2'
AND part->>'state' = 'output-available' LIMIT 5;
-- legacy: outputs live in tool_calls
SELECT elem->>'toolName',
left(regexp_replace((elem->'output')::text, '\s+', ' ', 'g'), 200)
FROM ai_chat_messages m, jsonb_array_elements(m.tool_calls) elem
@@ -280,26 +370,32 @@ docker compose -p gitmost logs -f --tail=100 # whole stack
```
Logging is `json-file`, `max-size=10m max-file=5` → ~50 MB retained, then rotated,
and **wiped on container recreate**. Since the #407 fix, thrown-tool error text is
**persisted in the `error` field** of `tool_calls` (see the hard-error query above), so
you no longer depend on live logs for it. Logs/live UI remain useful for **pre-#407
rows** (whose thrown errors were dropped) and for full stack traces beyond the
truncated stored message. A per-tool `tool_calls_total{tool,status}` metric to
VictoriaMetrics is still a possible future add for aggregate dashboards.
and **wiped on container recreate**. Thrown-tool error text is **persisted** — in the
`error` field of `tool_calls` (v2 `kind:'thrown'` / post-#407 legacy) — so you no longer
depend on live logs for it. Logs/live UI remain useful for **pre-#407 rows** (whose
thrown errors were dropped) and for full stack traces beyond the truncated stored
message. A per-tool `tool_calls_total{tool,status}` metric to VictoriaMetrics is still a
possible future add for aggregate dashboards.
## Gotchas checklist
- [ ] Counting every `tool_calls` element → **overcount**. Count `output` elements; add `error` elements for thrown failures (new rows), but don't count both as invocations.
- [ ] `isError` / `success=false` ≈ 0 does **not** mean "no errors" — thrown errors are a separate `error` element (new rows) or dropped entirely (pre-#407 rows).
- [ ] Thrown errors persist only on rows written **after the #407 fix** — pre-#407 rows still drop them (orphan only). Mind the cutover when trending over time.
- [ ] **Check `metadata.toolTraceVersion` first.** v2 (`= 2`) has no output in `tool_calls`; legacy has it inline. Never trend a metric across the era boundary.
- [ ] Counting every `tool_calls` element → **overcount**. Count OUTCOME elements — v2: `ok` or `error`; legacy: `output` or `error` — never both call+outcome as invocations.
- [ ] `isError` / `success=false` ≈ 0 does **not** mean "no errors" — thrown errors are an `error` element (v2 `kind:'thrown'` / post-#407), not in the output.
- [ ] **v2:** soft-error markers (the tool output) are in `metadata.parts`, NOT `tool_calls`. Legacy: they are in the `tool_calls` outcome `output`.
- [ ] **v2:** `kind` splits a real hard-fail (`thrown`) from an aborted call (`interrupted`) directly — no orphan heuristic needed. The orphan gap is a pre-#407-legacy-only proxy.
- [ ] `editPageText.failed` is `[]` on success — test for **non-empty**, not presence.
- [ ] Orphan gap on OLD rows mixes thrown errors **and** aborted runs — split by tool. On NEW rows a thrown error is its own `error` element, so a gap ≈ aborted call.
- [ ] `aborted` runs = server restarts, `failed` runs = provider overload — not agent mistakes.
- [ ] Never dump a raw `tool_calls` cell — it can be hundreds of KB.
- [ ] Logs are ephemeral (≤50 MB, wiped on recreate) — grab hard-error text live.
- [ ] Never dump a raw `tool_calls` **or** `metadata.parts` cell — outputs are hundreds of KB.
- [ ] Logs are ephemeral (≤50 MB, wiped on recreate) — grab pre-#407 hard-error text live.
## Snapshot (2026-07-07, illustrative — rerun the queries for current numbers)
> All rows in this snapshot predate #490, so they are **legacy-era** (outputs inline in
> `tool_calls`, orphan proxy for thrown errors). Do not trend these numbers against v2
> rows — segment by `toolTraceVersion` first.
- 226 chats, 732 messages, 46 runs; ~4 400 real tool invocations.
- Soft errors (persisted): `editPageText` 4/79 (bad/non-unique `find`) + 9 markdown-in-`find` warnings; `semanticSearch` 3/4 (`unavailable`); `Habr_update_draft_from_docmost` 1/2 (`doc` sent as object, not string).
- Missing-result proxy, read WITH the spread column:
+58 -20
View File
@@ -59,6 +59,39 @@ import {
mergeFootnoteDefinitions,
} from "../lib/transforms.js";
// Max concurrent per-page comment fetches in checkNewComments (#490). The scan is
// O(N) independent REST reads over the working set; running them one-at-a-time made
// a large space linear in round-trips. A small cap parallelizes without hammering
// the server (or exhausting sockets). 6 is a conservative middle of the 5–8 band.
const COMMENT_SCAN_CONCURRENCY = 6;
/**
* Map `items` through `fn` with at most `limit` in flight, preserving INPUT ORDER
* in the returned array. A tiny bounded pool (no p-limit dependency): `limit`
* workers pull the next index off a shared cursor until the list is drained.
*/
async function mapWithConcurrency<T, R>(
items: readonly T[],
limit: number,
fn: (item: T, index: number) => Promise<R>,
): Promise<R[]> {
const results = new Array<R>(items.length);
let cursor = 0;
const worker = async (): Promise<void> => {
for (;;) {
const i = cursor++;
if (i >= items.length) return;
results[i] = await fn(items[i], i);
}
};
const workers = Array.from(
{ length: Math.max(1, Math.min(limit, items.length)) },
() => worker(),
);
await Promise.all(workers);
return results;
}
// Public method surface of CommentsMixin (issue #450) — a NAMED type so the factory
// return type is expressible in the emitted .d.ts (the anonymous mixin class
// carries the base's protected shared state, which would otherwise trip TS4094).
@@ -655,27 +688,32 @@ export function CommentsMixin<TBase extends GConstructor<DocmostClientContext>>(
parentPageId,
);
// 2. Fetch comments for each page, keep ones created after since
const results: any[] = [];
for (const page of pagesInScope) {
try {
// Full feed (incl. resolved): a "new comments since" scan reports all
// recent activity; the active-only filter is scoped to listComments.
const comments = (await this.listComments(page.id, true)).items;
const newComments = comments.filter(
(c: any) => new Date(c.createdAt) > sinceDate,
);
if (newComments.length > 0) {
results.push({
pageId: page.id,
pageTitle: page.title,
comments: newComments,
});
// 2. Fetch comments for each page, keep ones created after since. Runs with
// bounded concurrency (#490) instead of one-at-a-time — the per-page reads are
// independent, so a large working set no longer costs O(N) serial round-trips.
// Order is preserved (mapWithConcurrency keeps input order), so the output is
// deterministic regardless of which fetch finishes first.
const perPage = await mapWithConcurrency(
pagesInScope,
COMMENT_SCAN_CONCURRENCY,
async (page: any) => {
try {
// Full feed (incl. resolved): a "new comments since" scan reports all
// recent activity; the active-only filter is scoped to listComments.
const comments = (await this.listComments(page.id, true)).items;
const newComments = comments.filter(
(c: any) => new Date(c.createdAt) > sinceDate,
);
return newComments.length > 0
? { pageId: page.id, pageTitle: page.title, comments: newComments }
: null;
} catch (e: any) {
// Skip pages with errors (e.g. deleted between calls)
return null;
}
} catch (e: any) {
// Skip pages with errors (e.g. deleted between calls)
}
}
},
);
const results: any[] = perPage.filter((r): r is any => r !== null);
const totalNewComments = results.reduce(
(sum, r) => sum + r.comments.length,
+63 -2
View File
@@ -170,6 +170,43 @@ export abstract class DocmostClientContext {
// cached conversion can never leak across identities. See getpage-cache.ts.
protected getPageCache = new GetPageConversionCache();
// #487: an OPTIONAL abort signal the in-app tool host sets before each tool
// call (a composite of the turn's Stop signal + a per-call wall-clock cap). It
// is checked at safe-points BETWEEN the sequential HTTP calls of a paginated
// read (paginateAll) and just before the atomic collab commit of a write (the
// mutatePage/replacePage/mutateLiveContentUnlocked seams), so a Stop / cap
// stops the NEXT network call from STARTING. An already-started single call may
// still land — a documented limitation (#487).
//
// SINGLE-WRITER by phase-1 assumption: exactly one DocmostClient is built per
// turn and shared by every tool call; the host sets this per call and does NOT
// restore the prior value on unwind (set-and-leave) — a fresh client per turn
// plus overwrite-by-the-next-call keeps it correct, and leaving a settled
// call's signal in place is what makes a discarded race-loser throw on its
// next safe-point. If the model emits PARALLEL in-app
// tool calls they share this one field, so the per-call CAP of one call is not
// guaranteed to bound another's in-flight pagination — but every composite the
// host sets carries the SAME turn Stop signal, so a Stop still aborts whichever
// signal is current. #487.
protected toolAbortSignal: AbortSignal | null = null;
/**
* #487: set (or clear with null) the in-app tool abort signal governing the
* NEXT client call's safe-points. The host wraps each in-app tool call: it sets
* the composite (Stop + per-call cap) here before invoking the tool and leaves
* it in place afterwards (set-and-leave, NOT restored) the next call
* overwrites it, and a fresh client is built per turn. Public so the
* server-side tool wrapper can reach it; harmless (a no-op) when never set.
*/
public setToolAbortSignal(signal: AbortSignal | null): void {
this.toolAbortSignal = signal;
}
/** #487: the abort signal currently governing this client's safe-points. */
public getToolAbortSignal(): AbortSignal | null {
return this.toolAbortSignal;
}
// Two construction forms:
// - new DocmostClient(config) // discriminated union (current)
// - new DocmostClient(baseURL, email, password) // legacy positional creds
@@ -571,6 +608,10 @@ export abstract class DocmostClientContext {
this.onMetricFn?.("collab_connect_timeouts_total", 1),
});
try {
// #487 PRE-COMMIT safe-point (reentrant twin of mutatePageContent): a
// Stop/cap after acquiring the session but before the atomic write skips
// this commit. Same limitation applies (stops the NEXT commit only).
this.toolAbortSignal?.throwIfAborted();
return await session.mutate(transform);
} catch (e) {
// Drop the session on any failure so the next call reconnects fresh.
@@ -602,6 +643,11 @@ export abstract class DocmostClientContext {
let truncated = false;
for (let page = 0; page < MAX_PAGES; page++) {
// #487 safe-point: a Stop (or the in-app tool per-call cap) that fires
// BETWEEN sequential page fetches must stop the NEXT request from starting
// — a read tool that would otherwise paginate for minutes is interrupted
// here. throwIfAborted() rejects with the signal's reason.
this.toolAbortSignal?.throwIfAborted();
const payload: Record<string, any> = {
...basePayload,
limit: clampedLimit,
@@ -709,7 +755,15 @@ export abstract class DocmostClientContext {
// #486: on a rejected collab-WS handshake, invalidate + refresh the token and
// retry the write once (symmetric to the HTTP-401 reauth path).
return this.writeWithCollabAuthRetry(collabToken, (token) =>
mutatePageContent(pageUuid, token, apiUrl, transform),
// #487: thread the in-app tool signal to mutatePageContent's pre-commit
// safe-point so a Stop/cap during the connect/lock window skips the write.
mutatePageContent(
pageUuid,
token,
apiUrl,
transform,
this.toolAbortSignal ?? undefined,
),
);
}
@@ -733,7 +787,14 @@ export abstract class DocmostClientContext {
// #486: on a rejected collab-WS handshake, invalidate + refresh the token and
// retry the write once (symmetric to the HTTP-401 reauth path).
return this.writeWithCollabAuthRetry(collabToken, (token) =>
replacePageContent(pageUuid, doc, token, apiUrl),
// #487: same pre-commit safe-point as mutatePage, for full-document writes.
replacePageContent(
pageUuid,
doc,
token,
apiUrl,
this.toolAbortSignal ?? undefined,
),
);
}
+7
View File
@@ -30,6 +30,13 @@ export { destroyAllSessions } from "./lib/collab-session.js";
// internals directly; it goes through loadDocmostMcp()).
export { SHARED_TOOL_SPECS } from "./tool-specs.js";
export type { SharedToolSpec } from "./tool-specs.js";
// #489 — write-class registry consumed by the in-app external-MCP retry gate.
export {
SHARED_TOOL_WRITE_CLASS,
isRetryableWriteClass,
assertEverySpecDeclaresWriteClass,
} from "./tool-specs.js";
export type { ToolWriteClass } from "./tool-specs.js";
// Re-export the build-time REGISTRY_STAMP (issue #447): a deterministic hash of
// the tool-specs registry content, generated into src/registry-stamp.generated.ts
+16
View File
@@ -254,6 +254,12 @@ export async function mutatePageContent(
collabToken: string,
baseUrl: string,
transform: (liveDoc: any) => any | null,
// #487: optional abort signal carrying the turn's Stop + the in-app tool
// per-call cap. Checked as the PRE-COMMIT safe-point below (after the session
// is acquired, immediately before the atomic read->write), so a Stop that
// arrives during the connect/lock window stops THIS write from landing. See the
// limitation note at the check.
signal?: AbortSignal,
): Promise<MutationResult> {
return withPageLock(pageId, async () => {
if (process.env.DEBUG) {
@@ -266,6 +272,13 @@ export async function mutatePageContent(
const session = await acquireCollabSession(pageId, collabToken, baseUrl);
try {
// #487 PRE-COMMIT safe-point: if the turn was Stopped (or the in-app tool
// per-call cap fired) after we acquired the collab session but before the
// atomic write, throw NOW so this commit never runs. KNOWN LIMITATION
// (#487): this only stops THIS commit — a write tool that already committed
// an EARLIER call this turn leaves that op applied. Cancel guarantees "no
// NEW commit starts", NOT "the write didn't land".
signal?.throwIfAborted();
return await session.mutate(transform);
} catch (e) {
// Drop the session on any failure so the next call reconnects fresh (this
@@ -291,6 +304,8 @@ export async function replacePageContent(
prosemirrorDoc: any,
collabToken: string,
baseUrl: string,
// #487: threaded straight to mutatePageContent's pre-commit safe-point.
signal?: AbortSignal,
): Promise<MutationResult> {
// Fail fast on a bad document instead of deferring the failure into the
// collaboration write (where TiptapTransformer.toYdoc(undefined) used to
@@ -307,6 +322,7 @@ export async function replacePageContent(
collabToken,
baseUrl,
() => prosemirrorDoc,
signal,
);
}
+110
View File
@@ -127,6 +127,18 @@ export interface SharedToolSpec {
mcpName: string;
/** camelCase key in the ai-SDK tools object (the in-app layer). */
inAppKey: string;
/**
* Write-class of the tool (#489), declared on EVERY spec (a registration-time
* assert enforces completeness; `satisfies Record<string, SharedToolSpec>`
* makes it a compile error to omit). 'readOnly' = a pure read that mutates
* NOTHING durable, so it is safe to auto-retry once after a transport break.
* 'write' = anything that mutates a page/comment/share/diagram/etc a
* transport error is INDETERMINATE (the server may have applied it before the
* connection reset), so it is NEVER blind-retried (a retry would double-apply,
* the #435 incident class). Consumed by the external-MCP retry path
* (mcp-clients.service.ts) to gate its single auto-retry.
*/
writeClass: 'readOnly' | 'write';
/** Single canonical model-facing description used by both layers. */
description: string;
/**
@@ -240,6 +252,7 @@ export const SHARED_TOOL_SPECS = {
getWorkspace: {
mcpName: 'getWorkspace',
inAppKey: 'getWorkspace',
writeClass: 'readOnly',
description: 'Fetch metadata about the current workspace (name, settings).',
tier: 'core',
catalogLine: 'getWorkspace — fetch current workspace metadata (name, settings).',
@@ -249,6 +262,7 @@ export const SHARED_TOOL_SPECS = {
listSpaces: {
mcpName: 'listSpaces',
inAppKey: 'listSpaces',
writeClass: 'readOnly',
description:
'List the spaces the current user can access. Returns the array of ' +
'spaces (id, name, slug, ...).',
@@ -260,6 +274,7 @@ export const SHARED_TOOL_SPECS = {
listShares: {
mcpName: 'listShares',
inAppKey: 'listShares',
writeClass: 'readOnly',
description:
'List all public shares in the workspace with page titles and public URLs.',
tier: 'deferred',
@@ -272,6 +287,7 @@ export const SHARED_TOOL_SPECS = {
getPageJson: {
mcpName: 'getPageJson',
inAppKey: 'getPageJson',
writeClass: 'readOnly',
description:
'Get page details with the raw ProseMirror JSON content (lossless: ' +
'includes block ids, callouts, tables, link/image attributes) plus the ' +
@@ -289,6 +305,7 @@ export const SHARED_TOOL_SPECS = {
getOutline: {
mcpName: 'getOutline',
inAppKey: 'getOutline',
writeClass: 'readOnly',
description:
"Return a COMPACT outline of a page's top-level blocks ({index, type, " +
'id, level, firstText}; tables add rows/cols/header; lists add item ' +
@@ -309,6 +326,7 @@ export const SHARED_TOOL_SPECS = {
getNode: {
mcpName: 'getNode',
inAppKey: 'getNode',
writeClass: 'readOnly',
description:
"Fetch a single block for editing. `nodeId` is a block id from the page " +
'outline or page-JSON view (works for headings/paragraphs/callouts/images), OR ' +
@@ -350,6 +368,7 @@ export const SHARED_TOOL_SPECS = {
searchInPage: {
mcpName: 'searchInPage',
inAppKey: 'searchInPage',
writeClass: 'readOnly',
description:
'Find every occurrence of a string (or regex) INSIDE one page and get ' +
'WHERE each is — instead of pulling blocks one-by-one with getNode. ' +
@@ -413,6 +432,7 @@ export const SHARED_TOOL_SPECS = {
deleteNode: {
mcpName: 'deleteNode',
inAppKey: 'deleteNode',
writeClass: 'write',
description:
'Remove a single block by its attrs.id (from the page outline or ' +
'page-JSON view) WITHOUT resending the whole document.',
@@ -438,6 +458,7 @@ export const SHARED_TOOL_SPECS = {
patchNode: {
mcpName: 'patchNode',
inAppKey: 'patchNode',
writeClass: 'write',
description:
'Replace a single content block identified by its attrs.id, WITHOUT ' +
'resending the whole document; the replacement keeps the same block id. ' +
@@ -505,6 +526,7 @@ export const SHARED_TOOL_SPECS = {
insertNode: {
mcpName: 'insertNode',
inAppKey: 'insertNode',
writeClass: 'write',
description:
'Insert content before/after another block (by attrs.id or anchor text) ' +
'or append it at the end (top level). For before/after you MUST provide ' +
@@ -597,6 +619,7 @@ export const SHARED_TOOL_SPECS = {
sharePage: {
mcpName: 'sharePage',
inAppKey: 'sharePage',
writeClass: 'write',
// CANONICAL: merges the MCP copy's URL-format + idempotency detail with the
// in-app copy's reversibility note; keeps the security framing both had.
description:
@@ -626,6 +649,7 @@ export const SHARED_TOOL_SPECS = {
unsharePage: {
mcpName: 'unsharePage',
inAppKey: 'unsharePage',
writeClass: 'write',
description: 'Remove the public share of a page (revokes the public URL).',
tier: 'deferred',
catalogLine: "unsharePage — revoke a page's public share (removes the public URL).",
@@ -640,6 +664,7 @@ export const SHARED_TOOL_SPECS = {
diffPageVersions: {
mcpName: 'diffPageVersions',
inAppKey: 'diffPageVersions',
writeClass: 'readOnly',
description:
'Diff two versions of a page and return a Docmost-equivalent change set ' +
'(inserted/deleted text, integrity counts for images/links/tables/' +
@@ -672,6 +697,7 @@ export const SHARED_TOOL_SPECS = {
listPageHistory: {
mcpName: 'listPageHistory',
inAppKey: 'listPageHistory',
writeClass: 'readOnly',
description:
"List a page's saved versions (Docmost auto-snapshots on every save), " +
'newest first, cursor-paginated. Returns { items, nextCursor }; each ' +
@@ -693,6 +719,7 @@ export const SHARED_TOOL_SPECS = {
restorePageVersion: {
mcpName: 'restorePageVersion',
inAppKey: 'restorePageVersion',
writeClass: 'write',
description:
'Restore a page to a saved version: writes that version\'s content back ' +
'as the page\'s current content (Docmost has no restore endpoint, so ' +
@@ -713,6 +740,7 @@ export const SHARED_TOOL_SPECS = {
importPageMarkdown: {
mcpName: 'importPageMarkdown',
inAppKey: 'importPageMarkdown',
writeClass: 'write',
// IN-APP ONLY (issue #411): the external /mcp surface no longer exposes
// importPageMarkdown — the registry loop in index.ts skips inAppOnly specs,
// so this stays available to the in-app agent (round-tripping an EXPORTED
@@ -742,6 +770,7 @@ export const SHARED_TOOL_SPECS = {
copyPageContent: {
mcpName: 'copyPageContent',
inAppKey: 'copyPageContent',
writeClass: 'write',
description:
"Replace targetPageId's content with a copy of sourcePageId's content, " +
'entirely server-side — the document is NOT sent through the model. The ' +
@@ -770,6 +799,7 @@ export const SHARED_TOOL_SPECS = {
editPageText: {
mcpName: 'editPageText',
inAppKey: 'editPageText',
writeClass: 'write',
description:
"Surgical find/replace inside a page's text, preserving all block " +
'ids and marks. A find MAY cross bold/italic/link boundaries; the ' +
@@ -819,6 +849,7 @@ export const SHARED_TOOL_SPECS = {
stashPage: {
mcpName: 'stashPage',
inAppKey: 'stashPage',
writeClass: 'readOnly',
description:
'Serialize a whole page (the full ProseMirror JSON, as getPageJson ' +
'returns) into an ephemeral in-memory blob and return ONLY a short ' +
@@ -880,6 +911,7 @@ export const SHARED_TOOL_SPECS = {
getPage: {
mcpName: 'getPage',
inAppKey: 'getPage',
writeClass: 'readOnly',
description:
'Fetch a single page as Markdown by its id. Returns the page title and ' +
'its Markdown content. The converter is canonical (round-trips text and ' +
@@ -919,6 +951,7 @@ export const SHARED_TOOL_SPECS = {
listPages: {
mcpName: 'listPages',
inAppKey: 'listPages',
writeClass: 'readOnly',
description:
'List the most recent pages (ordered by updatedAt, descending), ' +
'optionally scoped to a single space. Returns a bounded list (default ' +
@@ -965,6 +998,7 @@ export const SHARED_TOOL_SPECS = {
getTree: {
mcpName: 'getTree',
inAppKey: 'getTree',
writeClass: 'readOnly',
description:
"Get a space's page hierarchy (or one subtree) as a nested tree in a " +
'SINGLE request — completely and without loss. Each node is ' +
@@ -1009,6 +1043,7 @@ export const SHARED_TOOL_SPECS = {
getPageContext: {
mcpName: 'getPageContext',
inAppKey: 'getPageContext',
writeClass: 'readOnly',
description:
'Given a pageId, get its LOCATION and immediate surroundings (metadata ' +
'only, no page content) in one call — answers "where am I / what is ' +
@@ -1038,6 +1073,7 @@ export const SHARED_TOOL_SPECS = {
createPage: {
mcpName: 'createPage',
inAppKey: 'createPage',
writeClass: 'write',
description:
'Create a new page with a Markdown body in a space, optionally under a ' +
'parent page (omit parentPageId to create at the space root). Returns ' +
@@ -1088,6 +1124,7 @@ export const SHARED_TOOL_SPECS = {
movePage: {
mcpName: 'movePage',
inAppKey: 'movePage',
writeClass: 'write',
description:
'Move a page under a new parent page, or to the space root when no ' +
'parent is given. Reversible: move it back at any time.',
@@ -1181,6 +1218,7 @@ export const SHARED_TOOL_SPECS = {
renamePage: {
mcpName: 'renamePage',
inAppKey: 'renamePage',
writeClass: 'write',
description:
'Rename a page (change its title only; the body is untouched, never ' +
'resent). Reversible: rename back at any time.',
@@ -1202,6 +1240,7 @@ export const SHARED_TOOL_SPECS = {
deletePage: {
mcpName: 'deletePage',
inAppKey: 'deletePage',
writeClass: 'write',
description:
'Move a page to the trash — SOFT delete only: the page can be restored ' +
'from trash and nothing is ever permanently deleted.',
@@ -1235,6 +1274,7 @@ export const SHARED_TOOL_SPECS = {
updatePageJson: {
mcpName: 'updatePageJson',
inAppKey: 'updatePageJson',
writeClass: 'write',
description:
"Replace a page's content with a raw ProseMirror JSON document (lossless " +
'write: preserves the block ids, callouts, tables and attributes you pass ' +
@@ -1291,6 +1331,7 @@ export const SHARED_TOOL_SPECS = {
updatePageMarkdown: {
mcpName: 'updatePageMarkdown',
inAppKey: 'updatePageMarkdown',
writeClass: 'write',
description:
"Replace a page's body with new Markdown content (and optionally its " +
'title). The whole body is re-imported from the markdown (block ids ' +
@@ -1325,6 +1366,7 @@ export const SHARED_TOOL_SPECS = {
exportPageMarkdown: {
mcpName: 'exportPageMarkdown',
inAppKey: 'exportPageMarkdown',
writeClass: 'readOnly',
// CANONICAL: the MCP copy (a strict superset of the terse in-app wording).
description:
'Export a page to a single self-contained Docmost-flavoured Markdown ' +
@@ -1373,6 +1415,7 @@ export const SHARED_TOOL_SPECS = {
createComment: {
mcpName: 'createComment',
inAppKey: 'createComment',
writeClass: 'write',
// CANONICAL: the in-app copy (the more-maintained one). It keeps the same
// rules as the MCP copy — inline-only, top-level requires a `selection`, no
// page-level comments, replies inherit the anchor, suggestedText must be
@@ -1505,6 +1548,7 @@ export const SHARED_TOOL_SPECS = {
listComments: {
mcpName: 'listComments',
inAppKey: 'listComments',
writeClass: 'readOnly',
// CANONICAL: the two copies are near-identical; the MCP copy is the
// superset (it keeps the "(pagination is handled internally)" note the
// in-app copy dropped), so it is used verbatim.
@@ -1532,6 +1576,7 @@ export const SHARED_TOOL_SPECS = {
resolveComment: {
mcpName: 'resolveComment',
inAppKey: 'resolveComment',
writeClass: 'write',
// CANONICAL: the MCP copy's richer wording, minus its reference
// to `deleteComment` (a sibling tool that does NOT exist in the in-app
// layer) — rephrased transport-neutrally per the registry convention.
@@ -1572,6 +1617,7 @@ export const SHARED_TOOL_SPECS = {
checkNewComments: {
mcpName: 'checkNewComments',
inAppKey: 'checkNewComments',
writeClass: 'readOnly',
// CANONICAL: the MCP copy (the more detailed of the two). The MCP layer's
// execute-side guard that rejects an unparseable `since` timestamp stays in
// its execute body (per-layer logic), not in the shared schema.
@@ -1651,6 +1697,7 @@ export const SHARED_TOOL_SPECS = {
tableInsertRow: {
mcpName: 'tableInsertRow',
inAppKey: 'tableInsertRow',
writeClass: 'write',
description:
'Insert a row of plain-text cells into a table. `table` is `#<index>` ' +
'from the page outline, or a block id inside it. `cells` is the text per ' +
@@ -1684,6 +1731,7 @@ export const SHARED_TOOL_SPECS = {
tableDeleteRow: {
mcpName: 'tableDeleteRow',
inAppKey: 'tableDeleteRow',
writeClass: 'write',
description:
'Delete the row at 0-based `index` from a table (`table` is `#<index>` ' +
'from the page outline, or a block id inside it). Refuses to delete the ' +
@@ -1707,6 +1755,7 @@ export const SHARED_TOOL_SPECS = {
tableUpdateCell: {
mcpName: 'tableUpdateCell',
inAppKey: 'tableUpdateCell',
writeClass: 'write',
description:
'Set the plain-text content of cell [row, col] (0-based) in a table ' +
'(`table` is `#<index>` from the page outline, or a block id inside it). ' +
@@ -1747,6 +1796,7 @@ export const SHARED_TOOL_SPECS = {
insertFootnote: {
mcpName: 'insertFootnote',
inAppKey: 'insertFootnote',
writeClass: 'write',
description:
'Insert an AUTHOR-INLINE footnote: you specify only WHERE (anchorText) ' +
'and WHAT (text). The footnote marker is placed right after anchorText in ' +
@@ -1784,6 +1834,7 @@ export const SHARED_TOOL_SPECS = {
insertImage: {
mcpName: 'insertImage',
inAppKey: 'insertImage',
writeClass: 'write',
description:
'Download an image from a web (http/https) URL and insert it into ' +
'a page in one step. By default ' +
@@ -1828,6 +1879,7 @@ export const SHARED_TOOL_SPECS = {
replaceImage: {
mcpName: 'replaceImage',
inAppKey: 'replaceImage',
writeClass: 'write',
description:
'Replace an existing image on a page with a new image fetched from a web ' +
'(http/https) URL: uploads the new file as a NEW ' +
@@ -1866,6 +1918,7 @@ export const SHARED_TOOL_SPECS = {
drawioGet: {
mcpName: 'drawioGet',
inAppKey: 'drawioGet',
writeClass: 'readOnly',
description:
'Read a draw.io diagram on a page as mxGraph XML (default) or as its raw ' +
'`.drawio.svg`. `node` is the drawio node\'s attrs.id (from getOutline / ' +
@@ -1899,6 +1952,7 @@ export const SHARED_TOOL_SPECS = {
drawioCreate: {
mcpName: 'drawioCreate',
inAppKey: 'drawioCreate',
writeClass: 'write',
description:
'Create a draw.io diagram from mxGraph XML and insert it as a diagram ' +
'block. `xml` is a bare `<mxGraphModel>` OR a list of `<mxCell>` elements ' +
@@ -1969,6 +2023,7 @@ export const SHARED_TOOL_SPECS = {
drawioUpdate: {
mcpName: 'drawioUpdate',
inAppKey: 'drawioUpdate',
writeClass: 'write',
description:
'Replace a draw.io diagram\'s content with new mxGraph XML (same lint ' +
'pipeline as drawioCreate). `baseHash` is MANDATORY: pass the hash from ' +
@@ -2020,6 +2075,7 @@ export const SHARED_TOOL_SPECS = {
drawioEditCells: {
mcpName: 'drawioEditCells',
inAppKey: 'drawioEditCells',
writeClass: 'write',
description:
'Make TARGETED, id-based edits to an existing draw.io diagram instead of ' +
'resending the whole XML (a full-XML diff is fragile — draw.io reorders ' +
@@ -2078,6 +2134,7 @@ export const SHARED_TOOL_SPECS = {
drawioFromGraph: {
mcpName: 'drawioFromGraph',
inAppKey: 'drawioFromGraph',
writeClass: 'write',
description:
'Build a draw.io diagram from a SEMANTIC graph — you describe nodes, groups ' +
'and edges by MEANING and the server picks every coordinate, color and icon ' +
@@ -2202,6 +2259,7 @@ export const SHARED_TOOL_SPECS = {
drawioFromMermaid: {
mcpName: 'drawioFromMermaid',
inAppKey: 'drawioFromMermaid',
writeClass: 'write',
description:
'Convert Mermaid `flowchart` text into an EDITABLE draw.io diagram (LLMs ' +
'write Mermaid reliably). Best for STANDARD flowcharts/decision trees: ' +
@@ -2248,6 +2306,7 @@ export const SHARED_TOOL_SPECS = {
drawioShapes: {
mcpName: 'drawioShapes',
inAppKey: 'drawioShapes',
writeClass: 'readOnly',
description:
'Look up VERIFIED draw.io stencil style-strings so you never guess a ' +
'`shape=mxgraph.*` name (a wrong name renders as an EMPTY BOX). Searches a ' +
@@ -2294,6 +2353,7 @@ export const SHARED_TOOL_SPECS = {
drawioGuide: {
mcpName: 'drawioGuide',
inAppKey: 'drawioGuide',
writeClass: 'readOnly',
description:
'Progressive-disclosure draw.io authoring reference. Call with a `section` ' +
'to pull one focused, <=4KB chapter instead of bloating context: ' +
@@ -2323,3 +2383,53 @@ export const SHARED_TOOL_SPECS = {
inlineBothHosts: true,
},
} satisfies Record<string, SharedToolSpec>;
// --- write-class registry (#489) ------------------------------------------
/** A tool's retry-safety class. 'readOnly' may be auto-retried once after a
* transport break; 'write' is indeterminate and must never be blind-retried. */
export type ToolWriteClass = 'readOnly' | 'write';
/**
* Name write-class map for the shared registry, keyed by mcpName (=== inAppKey).
* The external-MCP retry path (mcp-clients.service.ts) looks a tool up here by its
* RAW (un-namespaced) name to decide whether a transport failure may be retried.
* A tool NOT in this map (a third-party external MCP tool) is treated as 'write'
* by the consumer the safe default (never blind-retry an unknown tool).
*/
export const SHARED_TOOL_WRITE_CLASS: Record<string, ToolWriteClass> =
Object.fromEntries(
Object.values(SHARED_TOOL_SPECS).map((spec) => [spec.mcpName, spec.writeClass]),
);
/** Whether a write-class permits a single automatic retry after a transport
* break. Only a pure read is retry-safe; everything mutating is indeterminate. */
export function isRetryableWriteClass(
writeClass: ToolWriteClass | undefined,
): boolean {
return writeClass === 'readOnly';
}
/**
* Registration-time assert (#489): EVERY spec must declare a valid write-class.
* `satisfies Record<string, SharedToolSpec>` already makes an omission a compile
* error, but this guards a raw/cast construction path and documents the invariant
* at the point of use. Runs once on import both hosts import this module, so
* both get the check. Throws (fails startup) rather than silently mis-gating a
* retry in production.
*/
export function assertEverySpecDeclaresWriteClass(): void {
for (const [key, spec] of Object.entries(SHARED_TOOL_SPECS)) {
const wc = (spec as SharedToolSpec).writeClass;
if (wc !== 'readOnly' && wc !== 'write') {
throw new Error(
`tool-specs: spec "${key}" must declare writeClass ('readOnly' | 'write'), got ${JSON.stringify(
wc,
)}`,
);
}
}
}
// Enforce at module load (registration time) on both hosts.
assertEverySpecDeclaresWriteClass();
@@ -0,0 +1,143 @@
// #487 commit 1 — the in-app tool cancellation safe-point inside paginateAll.
//
// The in-app tool host sets a composite abort signal on the client
// (setToolAbortSignal) before each tool call; paginateAll checks it at a
// safe-point BEFORE every sequential page fetch, so a Stop that lands mid-read
// stops the NEXT HTTP request from STARTING (a read tool can no longer paginate
// for minutes past a Stop). This pins the HONEST observable property against the
// REAL client + a real HTTP server: "after Stop, no NEW request starts".
import { test, after } from "node:test";
import assert from "node:assert/strict";
import http from "node:http";
import { DocmostClient } from "../../build/client.js";
function readBody(req) {
return new Promise((resolve) => {
let raw = "";
req.on("data", (c) => (raw += c));
req.on("end", () => resolve(raw));
});
}
function sendJson(res, status, obj, extra = {}) {
res.writeHead(status, { "Content-Type": "application/json", ...extra });
res.end(JSON.stringify(obj));
}
const openServers = [];
async function spawn(handler) {
const server = await new Promise((resolve) => {
const s = http.createServer(handler);
s.listen(0, "127.0.0.1", () => resolve(s));
});
openServers.push(server);
const { port } = server.address();
return { baseURL: `http://127.0.0.1:${port}/api` };
}
after(async () => {
await Promise.all(openServers.map((s) => new Promise((r) => s.close(r))));
});
function handleLogin(req, res) {
if (req.url === "/api/auth/login") {
sendJson(res, 200, { success: true }, {
"Set-Cookie": "authToken=t; Path=/; HttpOnly",
});
return true;
}
return false;
}
// A Stop that lands DURING pagination: the server aborts the client signal as it
// serves page 1 (more pages remain). The loop's next safe-point must throw before
// the page-2 request is sent.
test("paginateAll stops the NEXT request when the signal aborts mid-pagination", async () => {
let requests = 0;
const ac = new AbortController();
const { baseURL } = await spawn(async (req, res) => {
await readBody(req);
if (handleLogin(req, res)) return;
if (req.url === "/api/spaces") {
requests++;
// Simulate a user Stop that lands while page 1 is in flight.
if (requests === 1) ac.abort(new Error("user stop"));
sendJson(res, 200, {
success: true,
data: {
items: [{ id: `p${requests}` }],
meta: { hasNextPage: true, nextCursor: `c${requests}` },
},
});
return;
}
sendJson(res, 404, {});
});
const client = new DocmostClient(baseURL, "user@example.com", "pw");
client.setToolAbortSignal(ac.signal);
await assert.rejects(
() => client.paginateAll("/spaces", {}),
/user stop/,
"the aborted safe-point rejects with the signal's reason",
);
assert.equal(requests, 1, "page 2 never started after the Stop");
});
// A Stop that is already in effect before the read starts: zero requests fire.
test("paginateAll starts no request when the signal is already aborted", async () => {
let requests = 0;
const { baseURL } = await spawn(async (req, res) => {
await readBody(req);
if (handleLogin(req, res)) return;
if (req.url === "/api/spaces") {
requests++;
sendJson(res, 200, {
success: true,
data: { items: [], meta: { hasNextPage: false, nextCursor: null } },
});
return;
}
sendJson(res, 404, {});
});
const client = new DocmostClient(baseURL, "user@example.com", "pw");
// Warm the auth so ensureAuthenticated does not itself POST after the abort.
await client.ensureAuthenticated();
const ac = new AbortController();
ac.abort(new Error("already stopped"));
client.setToolAbortSignal(ac.signal);
await assert.rejects(() => client.paginateAll("/spaces", {}), /already stopped/);
assert.equal(requests, 0, "no /spaces request started once already aborted");
});
// Without a tool signal (default), pagination is unaffected — the safe-point is a
// pure no-op, so pre-#487 behaviour is byte-identical.
test("paginateAll is unaffected when no tool signal is set", async () => {
let requests = 0;
const PAGES = {
"": { items: [{ id: "a" }], nextCursor: "c1" },
c1: { items: [{ id: "b" }], nextCursor: null },
};
const { baseURL } = await spawn(async (req, res) => {
const raw = await readBody(req);
if (handleLogin(req, res)) return;
if (req.url === "/api/spaces") {
requests++;
const body = JSON.parse(raw || "{}");
const page = PAGES[body.cursor ?? ""] ?? { items: [], nextCursor: null };
sendJson(res, 200, {
success: true,
data: {
items: page.items,
meta: { hasNextPage: page.nextCursor != null, nextCursor: page.nextCursor },
},
});
return;
}
sendJson(res, 404, {});
});
const client = new DocmostClient(baseURL, "user@example.com", "pw");
const all = await client.paginateAll("/spaces", {});
assert.equal(requests, 2, "both pages fetched with no signal set");
assert.deepEqual(all.map((p) => p.id), ["a", "b"]);
});
@@ -442,3 +442,71 @@ test("checkNewComments subtree includes the root without a separate getPageRaw",
assert.equal(result.checkedPages, 2, "root + one descendant scanned");
assert.equal(result.totalNewComments, 1, "the root's fresh comment found");
});
// -----------------------------------------------------------------------------
// 6) checkNewComments parallelism (#490): the per-page comment fetches run with
// bounded concurrency (not one-at-a-time), and the results still preserve the
// page order deterministically regardless of which fetch finishes first.
// -----------------------------------------------------------------------------
test("checkNewComments fetches pages concurrently (bounded) and preserves order", async () => {
// A subtree with 12 descendants so the scan has plenty to parallelize.
const NODES = [{ id: "parent", title: "Parent", parentPageId: null, hasChildren: true }];
for (let i = 0; i < 12; i++) {
NODES.push({ id: `k${i}`, title: `Kid ${i}`, parentPageId: "parent", hasChildren: false });
}
let inFlight = 0;
let maxInFlight = 0;
const { baseURL } = await spawn(async (req, res) => {
const raw = await readBody(req);
if (handleLogin(req, res)) return;
if (req.url === "/api/pages/tree") {
sendJson(res, 200, { success: true, data: { items: NODES } });
return;
}
if (req.url === "/api/comments") {
const body = JSON.parse(raw || "{}");
inFlight++;
maxInFlight = Math.max(maxInFlight, inFlight);
// Hold the response briefly so concurrent fetches actually overlap.
setTimeout(() => {
inFlight--;
// Every page carries one fresh comment so ordering is observable.
sendJson(res, 200, {
success: true,
data: {
items: [
{ id: `c-${body.pageId}`, createdAt: "2030-01-01T00:00:00.000Z", content: null },
],
meta: { nextCursor: null },
},
});
}, 25);
return;
}
sendJson(res, 404, {});
});
const client = new DocmostClient(baseURL, "user@example.com", "pw");
const result = await client.checkNewComments(
"space-1",
"2020-01-01T00:00:00.000Z",
"parent",
);
// 13 pages (parent + 12 kids) were scanned; each had a fresh comment.
assert.equal(result.checkedPages, 13, "all pages scanned");
assert.equal(result.totalNewComments, 13, "one fresh comment per page");
// Parallelism: more than one request was in flight at once, but never above the
// cap (6). A serial implementation would show maxInFlight === 1.
assert.ok(maxInFlight > 1, `expected concurrent fetches, saw max ${maxInFlight}`);
assert.ok(maxInFlight <= 6, `concurrency must be bounded, saw ${maxInFlight}`);
// Deterministic order: results follow the page-enumeration order (parent first).
assert.equal(result.comments[0].pageId, "parent", "results preserve page order");
assert.deepEqual(
result.comments.map((r) => r.pageId),
["parent", ...Array.from({ length: 12 }, (_, i) => `k${i}`)],
"result order matches the enumeration order regardless of finish order",
);
});
@@ -0,0 +1,164 @@
// #487 F4 — the WRITE-side cancellation safe-point.
//
// Every content-mutating collab write (collaboration.mutatePageContent and the
// reentrant twin client.mutateLiveContentUnlocked used by replaceImage) checks
// the in-app tool abort signal at a PRE-COMMIT safe-point — after the collab
// session is acquired but immediately BEFORE the atomic read->write
// (session.mutate). So a Stop (or the per-call cap) that lands during the
// connect/lock window stops THIS write from landing: no new commit starts once
// aborted. paginate-abort-safepoint.test.mjs pins the READ half; this pins the
// integrity-critical WRITE half — remove the `throwIfAborted()` and the transform
// would run and the doc would be mutated past a Stop.
//
// There is no collab server in the unit env, so we swap the provider factory
// (__setCollabProviderFactory) for a fake that reports an immediate successful
// sync. That makes acquireCollabSession SUCCEED and hand back a live, ready
// session, so the ONLY thing standing between the call and session.mutate is the
// safe-point under test. The transform is instrumented to prove it never runs.
import { test, afterEach } from "node:test";
import assert from "node:assert/strict";
import { mutatePageContent } from "../../build/lib/collaboration.js";
import {
__setCollabProviderFactory,
destroyAllSessions,
} from "../../build/lib/collab-session.js";
import { DocmostClient } from "../../build/client.js";
const BASE_URL = "http://127.0.0.1:1/api";
// mutatePageContent locks via withPageLock, which demands a canonical page UUID
// (resolve-then-lock invariant, #260/#449); the unlocked twin does not.
const PAGE_UUID = "11111111-1111-4111-8111-111111111111";
// A fake HocuspocusProvider that immediately reports a successful initial sync so
// CollabSession.open() resolves to a ready session, and stays "synced" with zero
// unsynced changes so a reached session.mutate() would resolve at once. It speaks
// only the tiny CollabProviderLike surface the session depends on.
function syncedProviderFactory(config) {
// Fire the initial-sync callback so open() settles as ready.
config.onSynced();
return {
synced: true,
unsyncedChanges: 0,
destroy() {},
on() {},
off() {},
};
}
// Disable the session cache so every acquire opens (and the failure path destroys)
// its own ephemeral session — no cross-test session reuse.
process.env.MCP_COLLAB_SESSION_IDLE_MS = "0";
afterEach(() => {
__setCollabProviderFactory(null); // restore the real factory
destroyAllSessions();
});
// --- collaboration.mutatePageContent (the page-locked write path) -------------
test("mutatePageContent rejects at the pre-commit safe-point BEFORE session.mutate when the signal is already aborted", async () => {
__setCollabProviderFactory(syncedProviderFactory);
let transformCalls = 0;
const ac = new AbortController();
ac.abort(new Error("user stop"));
await assert.rejects(
() =>
mutatePageContent(
PAGE_UUID,
"collab-jwt",
BASE_URL,
(liveDoc) => {
transformCalls++;
return liveDoc;
},
ac.signal,
),
/user stop/,
"the aborted safe-point rejects with the signal's reason before committing",
);
assert.equal(
transformCalls,
0,
"the transform (and therefore session.mutate) must NEVER run once aborted",
);
});
test("mutatePageContent (control) DOES reach session.mutate and invoke the transform when the signal is live", async () => {
__setCollabProviderFactory(syncedProviderFactory);
let transformCalls = 0;
const ac = new AbortController(); // never aborted
const result = await mutatePageContent(
PAGE_UUID,
"collab-jwt",
BASE_URL,
(liveDoc) => {
transformCalls++;
return null; // null -> no-op write; still proves the transform was invoked
},
ac.signal,
);
assert.equal(
transformCalls,
1,
"with a live signal the safe-point is a no-op and session.mutate runs the transform",
);
assert.ok(result && result.verify, "a MutationResult is returned");
});
// --- client.mutateLiveContentUnlocked (the reentrant twin, replaceImage) ------
test("mutateLiveContentUnlocked rejects at the pre-commit safe-point BEFORE session.mutate when the tool signal is already aborted", async () => {
__setCollabProviderFactory(syncedProviderFactory);
const client = new DocmostClient({
apiUrl: BASE_URL,
getToken: async () => "access",
getCollabToken: async () => "collab-jwt",
});
let transformCalls = 0;
const ac = new AbortController();
ac.abort(new Error("cap fired"));
client.setToolAbortSignal(ac.signal);
await assert.rejects(
() =>
client.mutateLiveContentUnlocked("page-1", "collab-jwt", (liveDoc) => {
transformCalls++;
return liveDoc;
}),
/cap fired/,
"the aborted safe-point rejects with the signal's reason before committing",
);
assert.equal(
transformCalls,
0,
"the transform (and therefore session.mutate) must NEVER run once aborted",
);
});
test("mutateLiveContentUnlocked (control) DOES reach session.mutate and invoke the transform when the tool signal is live", async () => {
__setCollabProviderFactory(syncedProviderFactory);
const client = new DocmostClient({
apiUrl: BASE_URL,
getToken: async () => "access",
getCollabToken: async () => "collab-jwt",
});
let transformCalls = 0;
client.setToolAbortSignal(new AbortController().signal); // live
const result = await client.mutateLiveContentUnlocked(
"page-1",
"collab-jwt",
(liveDoc) => {
transformCalls++;
return null;
},
);
assert.equal(
transformCalls,
1,
"with a live signal the safe-point is a no-op and session.mutate runs the transform",
);
assert.ok(result && result.verify, "a MutationResult is returned");
});
+41 -1
View File
@@ -2,7 +2,12 @@ import { test } from "node:test";
import assert from "node:assert/strict";
import { z } from "zod";
import { SHARED_TOOL_SPECS } from "../../build/tool-specs.js";
import {
SHARED_TOOL_SPECS,
SHARED_TOOL_WRITE_CLASS,
isRetryableWriteClass,
assertEverySpecDeclaresWriteClass,
} from "../../build/tool-specs.js";
// The shared registry is consumed by BOTH the zod-v3 MCP server and the zod-v4
// in-app AI-SDK service, so every spec must carry the cross-layer wiring
@@ -43,6 +48,41 @@ test("mcpName and inAppKey are each unique across the registry", () => {
}
});
// #489 — every spec must declare its write-class so the external-MCP retry path
// can gate a single auto-retry ONLY on a pure read (a blind retry of a write =
// double-apply). The declaration is enforced at registration time.
test("#489: every spec declares a valid writeClass ('readOnly' | 'write')", () => {
for (const [key, spec] of Object.entries(SHARED_TOOL_SPECS)) {
assert.ok(
spec.writeClass === "readOnly" || spec.writeClass === "write",
`${key}: missing/invalid writeClass: ${JSON.stringify(spec.writeClass)}`,
);
}
// The registration-time assert must not throw for the shipped registry.
assert.doesNotThrow(() => assertEverySpecDeclaresWriteClass());
});
test("#489: SHARED_TOOL_WRITE_CLASS maps every mcpName to its class; helper gates on readOnly", () => {
const specs = Object.values(SHARED_TOOL_SPECS);
assert.equal(Object.keys(SHARED_TOOL_WRITE_CLASS).length, specs.length);
for (const spec of specs) {
assert.equal(SHARED_TOOL_WRITE_CLASS[spec.mcpName], spec.writeClass);
}
// Only a readOnly tool is retry-eligible; a write tool and an unknown tool are not.
assert.equal(isRetryableWriteClass("readOnly"), true);
assert.equal(isRetryableWriteClass("write"), false);
assert.equal(isRetryableWriteClass(undefined), false);
});
test("#489: representative reads are readOnly and representative writes are write", () => {
for (const name of ["getPage", "getTree", "searchInPage", "listComments"]) {
assert.equal(SHARED_TOOL_SPECS[name].writeClass, "readOnly", `${name} should be readOnly`);
}
for (const name of ["patchNode", "createPage", "deletePage", "createComment", "drawioCreate"]) {
assert.equal(SHARED_TOOL_SPECS[name].writeClass, "write", `${name} should be write`);
}
});
test("buildShape (when present) returns a usable ZodRawShape with a real zod", () => {
for (const [key, spec] of Object.entries(SHARED_TOOL_SPECS)) {
if (!spec.buildShape) continue;
+19
View File
@@ -0,0 +1,19 @@
{
"name": "@docmost/token-estimate",
"version": "0.1.0",
"description": "Shared, provider-agnostic token estimator (chars/2.5) used by the AI-chat client counter and the server history-replay budgeter, so the two never diverge.",
"private": true,
"main": "./dist/index.js",
"types": "./dist/index.d.ts",
"scripts": {
"build": "tsc",
"watch": "tsc --watch",
"test": "vitest run",
"test:watch": "vitest"
},
"license": "MIT",
"devDependencies": {
"typescript": "^5.0.0",
"vitest": "4.1.6"
}
}
+31
View File
@@ -0,0 +1,31 @@
import { describe, it, expect } from 'vitest';
import { estimateTokens, CHARS_PER_TOKEN } from './index';
describe('estimateTokens (shared chars/2.5)', () => {
it('returns 0 for empty / nullish input', () => {
expect(estimateTokens('')).toBe(0);
expect(estimateTokens(null)).toBe(0);
expect(estimateTokens(undefined)).toBe(0);
});
it('uses the chars/2.5 ratio, ceiled', () => {
expect(CHARS_PER_TOKEN).toBe(2.5);
// 5 chars / 2.5 = 2
expect(estimateTokens('abcde')).toBe(2);
// any non-empty string is at least 1 token (ceil)
expect(estimateTokens('a')).toBe(1);
// 100 chars / 2.5 = 40
expect(estimateTokens('x'.repeat(100))).toBe(40);
});
it('counts Cyrillic ~2x higher than the old chars/4 rule (no undercount)', () => {
const cyr = 'привет мир как дела'; // 19 chars
expect(estimateTokens(cyr)).toBe(Math.ceil(19 / 2.5)); // 8
expect(estimateTokens(cyr)).toBeGreaterThan(Math.ceil(19 / 4)); // > 5
});
it('is deterministic / byte-stable (same input => same output)', () => {
const s = 'the quick brown fox';
expect(estimateTokens(s)).toBe(estimateTokens(s));
});
});
+35
View File
@@ -0,0 +1,35 @@
/**
* Shared, provider-agnostic token estimator (#490).
*
* No provider exposes an exact tokenizer we can afford to run on the hot path (a
* real BPE pass is O(n²)-ish, bloats the client bundle, and is wrong for
* Gemini/Ollama anyway), so both the client's in-body counter AND the server's
* history-replay budgeter use this ONE cheap chars-based heuristic. Keeping it in
* a single shared module is deliberate: two independent estimators drift, and then
* "the badge shows 60%" while "the budgeter already trimmed" the exact confusion
* this package prevents.
*
* Ratio: **chars / 2.5**. Most content here is Cyrillic, where a token is ~2.5
* characters; the common English `chars/4` rule of thumb UNDER-counts Cyrillic by
* ~2×, which for a budget check is the dangerous direction (it lets the context
* overflow). 2.5 slightly over-estimates pure English/code, which is the SAFE
* direction for a budget. This is an estimate, never an exact count the
* authoritative figure is always the provider's reported usage; the estimate is
* for UI affordances, the delta of not-yet-sent messages, and deciding what to
* trim.
*/
/** Characters per token for the shared estimate. See the module comment. */
export const CHARS_PER_TOKEN = 2.5;
/**
* Rough token estimate for a piece of text (chars / {@link CHARS_PER_TOKEN}).
* Returns 0 for empty/nullish input, and ceils so any non-empty text counts as at
* least one token. Pure and deterministic (byte-stable), so the same text always
* yields the same estimate which the server budgeter relies on to keep replay
* trimming stable turn to turn (provider prompt-cache friendliness).
*/
export function estimateTokens(text: string | null | undefined): number {
if (!text) return 0;
return Math.ceil(text.length / CHARS_PER_TOKEN);
}
+16
View File
@@ -0,0 +1,16 @@
{
"compilerOptions": {
"target": "ES2022",
"module": "CommonJS",
"moduleResolution": "Node",
"outDir": "./dist",
"rootDir": "./src",
"strict": true,
"esModuleInterop": true,
"skipLibCheck": true,
"forceConsistentCasingInFileNames": true,
"declaration": true
},
"include": ["src/**/*"],
"exclude": ["src/**/*.test.ts"]
}
+15
View File
@@ -287,6 +287,9 @@ importers:
'@docmost/prosemirror-markdown':
specifier: workspace:*
version: link:../../packages/prosemirror-markdown
'@docmost/token-estimate':
specifier: workspace:*
version: link:../../packages/token-estimate
'@excalidraw/excalidraw':
specifier: 0.18.0-3a5ef40
version: 0.18.0-3a5ef40(@types/react-dom@18.3.1)(@types/react@18.3.12)(react-dom@18.3.1(react@18.3.1))(react@18.3.1)
@@ -561,6 +564,9 @@ importers:
'@docmost/prosemirror-markdown':
specifier: workspace:*
version: link:../../packages/prosemirror-markdown
'@docmost/token-estimate':
specifier: workspace:*
version: link:../../packages/token-estimate
'@fastify/compress':
specifier: ^9.0.0
version: 9.0.0
@@ -1148,6 +1154,15 @@ importers:
specifier: 4.1.6
version: 4.1.6(@opentelemetry/api@1.9.0)(@types/node@20.19.43)(@vitest/coverage-v8@4.1.6)(happy-dom@20.8.9)(jsdom@25.0.0)(vite@8.0.5(@types/node@20.19.43)(esbuild@0.28.0)(jiti@2.4.2)(less@4.2.0)(sugarss@5.0.1(postcss@8.5.14))(terser@5.39.0)(tsx@4.21.0)(yaml@2.8.3))
packages/token-estimate:
devDependencies:
typescript:
specifier: ^5.0.0
version: 5.9.3
vitest:
specifier: 4.1.6
version: 4.1.6(@opentelemetry/api@1.9.0)(@types/node@25.5.0)(@vitest/coverage-v8@4.1.6)(happy-dom@20.8.9)(jsdom@27.4.0(@noble/hashes@2.0.1))(vite@8.0.5(@types/node@25.5.0)(esbuild@0.28.0)(jiti@2.4.2)(less@4.2.0)(sugarss@5.0.1(postcss@8.5.14))(terser@5.39.0)(tsx@4.21.0)(yaml@2.8.3))
packages:
'@aashutoshrathi/word-wrap@1.2.6':