Compare commits

..

11 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
45 changed files with 4900 additions and 2916 deletions
+18 -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.
+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
+17
View File
@@ -336,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",
@@ -86,11 +86,19 @@ const MIN_HEIGHT = 400;
// Margin kept between the window and the viewport edges while dragging.
const EDGE_MARGIN = 8;
// #184 phase 1.5 / #430 / #488: the degraded-poll fallback. The window owns only
// a DUMB 2.5s timer, gated by an armed flag; the THREAD's run-lifecycle FSM owns
// arm/disarm AND the inactivity cap that turns a stuck run into a `stalled` banner
// (#488 commit 4a — the cap moved into the thread so polling->stalled is a single
// FSM transition; the window no longer silently stops polling at the cap).
// #184 phase 1.5 / #430: backstop for the degraded-poll fallback. The poll is
// armed when a resume attempt could not attach to the live run and disarmed by the
// thread on settle / local stream; this cap is the ONLY backstop against an endless
// tick (a stuck 'streaming' row before the boot-sweep, or a user-tail 204 with no
// run).
//
// #430: measured from RUN ACTIVITY, not from arm-time. A real autonomous run takes
// 11-25 min — longer than a fixed 10-min-from-start cap, which used to cut the poll
// off mid-run. Instead we cap on INACTIVITY: keep polling as long as the run is
// still making progress (its persisted rows keep changing), and only give up after
// this long with NO new activity. A genuinely stuck run produces no row changes, so
// the idle cap still bounds it; a long-but-progressing run polls to completion.
const DEGRADED_POLL_IDLE_MAX_MS = 10 * 60_000;
/** Compact token formatter: 1.2M / 3.4k / 950. */
function formatTokens(n: number): string {
@@ -251,13 +259,17 @@ export default function AiChatWindow() {
[roles],
);
// #184 phase 1.5 / #488: degraded-poll fallback. ChatThread's FSM arms this via
// onResumeFallback(true) when it enters a poll-bearing recovery (attach 204 /
// starved finish / stop) and disarms it on settle / local stream / stalled. The
// window owns ONLY the dumb 2.5s timer; the THREAD owns arm/disarm AND the
// inactivity cap (a stuck run -> the thread's `stalled` banner disarms this).
// #184 phase 1.5: degraded-poll fallback (replaces the F4/F5/F7 latches). When
// ChatThread could not attach to a still-running run it arms this via
// onResumeFallback(true); the thread disarms it on settle / local stream. The
// window only OWNS the timer (armedAtRef stamps when it was armed for the cap).
const [degradedPoll, setDegradedPoll] = useState(false);
// #430: timestamp of the LAST run activity while the poll is armed — stamped on
// arm and re-stamped whenever the polled rows change (see the effect below). The
// idle cap is measured from this, so a long-but-progressing run keeps polling.
const lastActivityAtRef = useRef(0);
const onResumeFallback = useCallback((active: boolean): void => {
if (active) lastActivityAtRef.current = Date.now();
setDegradedPoll(active);
}, []);
// Reset the degraded poll whenever the open chat changes: it is scoped to the
@@ -269,17 +281,33 @@ export default function AiChatWindow() {
const { data: messageRows, isLoading: messagesLoading } =
useAiChatMessagesQuery(
activeChatId ?? undefined,
// DELIBERATELY DUMB: poll every 2.5s WHILE ARMED, otherwise off. NO error
// checks (TanStack resets fetchFailureCount each fetch; the poll must survive
// a server restart), NO tail checks, NO cap here — the settled/stalled/idle-cap
// semantics all live in ChatThread's FSM, which disarms via onResumeFallback.
() => (degradedPoll === true ? 2500 : false),
// DELIBERATELY DUMB (invariant 8 / task 2.4): poll every 2.5s while armed
// and while the run is still active (#430: under the INACTIVITY cap, not a
// fixed-from-start cap); otherwise off. NO error checks (TanStack v5 resets
// fetchFailureCount each fetch, so consecutive errors are not expressible —
// and the poll must survive a server restart) and NO tail checks (the
// settled/local-stream semantics live in ChatThread, which disarms via
// onResumeFallback(false)). The idle cap is the only backstop.
() =>
degradedPoll === true &&
Date.now() - lastActivityAtRef.current < DEGRADED_POLL_IDLE_MAX_MS
? 2500
: false,
// #344: gate on windowOpen too — no message history is fetched (and no
// degraded poll runs) while the window is closed; it loads when the window
// opens with an active chat.
windowOpen,
);
// #430: re-stamp the activity clock whenever the polled rows change while the
// poll is armed. TanStack keeps the same `messageRows` reference across refetches
// that return deep-equal data (structural sharing), so a new reference means the
// run genuinely progressed — which extends the inactivity cap above. A stuck run
// yields no reference change, so the cap eventually fires and stops the poll.
useEffect(() => {
if (degradedPoll) lastActivityAtRef.current = Date.now();
}, [degradedPoll, messageRows]);
// #184 reconnect-and-live-follow. Whether detached agent runs are enabled for
// this workspace. When the feature is off no runs are ever created, so the
// resume attempt would only ever 204; gating ChatThread's resume on it avoids a
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
@@ -57,25 +57,6 @@ export async function stopRun(
return req.data;
}
/**
* #488: the run-fact — "is a run active on this chat?" — first-class from the
* server (POST /ai-chat/run). Called on mount to seed the client FSM's run-fact
* and to VERIFY after a supersede mismatch (an observer following a superseded
* run asks for the latest run and follows it). Returns the latest run row (with
* its `id` and `status`) and its projected assistant message, or `run: null` when
* the chat has never had a run. Owner-gated server-side.
*/
export async function getRun(chatId: string): Promise<{
run: { id: string; status: string } | null;
message: IAiChatMessageRow | null;
}> {
const req = await api.post<{
run: { id: string; status: string } | null;
message: IAiChatMessageRow | null;
}>("/ai-chat/run", { chatId });
return req.data;
}
/**
* Resolve the chat bound to a document (the current user's most-recent chat
* created on that page), or null when there is none. Drives auto-open-on-page.
@@ -1,183 +0,0 @@
# AI-chat run-lifecycle FSM — design spec (#488)
This is the written design that `run-fsm.ts` implements. It ships in the PR (issue
#488 commit 1: "the spec is written FIRST and enters the PR"). It has four parts:
(1) the event × state transition table, (2) the map of every `chat-thread.tsx` ref
to {FSM state | FSM context | stays data}, (3) the run-fact protocol, (4) the
invariants.
The reducer is a **pure function** `reduce(machine, event) → machine`. The returned
machine carries the **command effects** for that transition; a thin runtime in
`chat-thread.tsx` dispatches events and executes effects. Because it is pure, the
whole machine is enumerable and unit-tested directly (event × state → next state is
the observable property) — see `run-fsm.test.ts`.
---
## 1. Event × state transition table
Phases: `idle | sending | streaming | attaching | reconnecting(attempt,failed) |
polling(reason) | stalled | stopping | superseding | error(kind)`.
Context (orthogonal): `epoch`, `ownership: local|observer`, `runFact: {runId}|null`,
`liveFollow` (are we following a live run we locally streamed — the reconnect
ladder — vs a one-shot mount-attach resume? both are `observer`, but a live-follow
drop RE-ENTERS the ladder (#488 commit 3) while a mount-resume drop polls).
Legend: **†** = command-transition (bumps `epoch`, I1). Effects in `[…]`.
| Event (source) | From phase(s) | → To phase | Effects / ctx |
|---|---|---|---|
| `SEND_LOCAL` (user send) | idle, error, polling, stalled, reconnecting | sending **†** | `[cancelReconnect, disarmPoll]`, ownership=local |
| `STREAM_START{runId}` (SDK `start` metadata) | sending, attaching, reconnecting, superseding | streaming | `[cancelReconnect, disarmPoll]`, runFact←runId |
| `FINISH_CLEAN` (onFinish clean) | streaming, … | idle | `[disarmPoll, cancelReconnect]`, runFact←null |
| `FINISH_ABORT` (onFinish isAbort) | streaming, stopping | idle | `[disarmPoll, cancelReconnect]`, runFact←null (I4 exits stopping by this DATA) |
| `FINISH_DISCONNECT` (observer, NOT liveFollow) | streaming(observer) | polling(disconnect-visible) | `[armPoll]` (a mount-resume drop polls) |
| `FINISH_DISCONNECT{hasVisibleContent}` (local drop OR liveFollow) | streaming | reconnecting(1) **†** *iff runFact\|liveFollow* | `[scheduleReconnect(1)]` (+`armPoll` if visible), ownership=observer, liveFollow=true (commit 3: repeatable) |
| `FINISH_DISCONNECT` (no runFact, not liveFollow) | streaming | idle | runFact←null (plain terminal "connection lost") |
| `STREAM_INCOMPLETE{reason}` (observer starved/torn clean finish) | streaming(observer) | polling(reason) | `[armPoll(reason)]` |
| `FINISH_ERROR{kind}` (onFinish isError) | any | error(kind) | `[disarmPoll, cancelReconnect]`, runFact←null |
| `STREAM_START{runId}` (first assistant frame of a local turn) | sending | streaming | runFact←runId, `[cancelReconnect, disarmPoll]` |
| `ATTACH_START{runId}` (mount resume) | **idle only** (F2) | attaching **†** | `[resumeStream]`, ownership=observer, runFact←runId; ignored from any non-idle phase |
| `ATTACH_LIVE` (attach GET 2xx) | attaching | streaming | — |
| `ATTACH_NONE` (attach GET 204/err/throw) | attaching | polling(attach-none) | `[armPoll(attach-none)]` |
| `RECONNECT_ATTEMPT{n}` (backoff timer) | reconnecting | reconnecting(n) **†** | `[resumeStream]` |
| `RECONNECT_ATTACHED` (reconnect GET 2xx) | reconnecting | streaming | `[cancelReconnect, disarmPoll]`**counter reset** (commit 3) |
| `RECONNECT_NONE` (reconnect GET 204/err), attempt<MAX | reconnecting | reconnecting(n+1) **†** | `[armPoll(attach-none), scheduleReconnect(n+1)]` |
| `RECONNECT_NONE`, attempt=MAX | reconnecting | reconnecting(MAX, failed) | `[armPoll(reconnect-exhausted)]` |
| `RETRY` (manual, failed banner) | reconnecting(failed) | reconnecting(1) **†** | `[resumeStream]` |
| `RETRY` (manual, stalled banner) | stalled | polling(attach-none) **†** | `[armPoll]` |
| `POLL_TERMINAL` (settled tail merged) | polling, reconnecting, stopping | idle | `[disarmPoll, cancelReconnect]`, runFact←null (I4) |
| `POLL_IDLE_CAP` (inactivity cap) | polling, reconnecting | stalled | `[disarmPoll, cancelReconnect]` (commit 4a — no more silent) |
| `RUN_FACT{null}` (POST /run → null/terminal, 204) | reconnecting/attaching/polling/stopping | idle | `[cancelReconnect, disarmPoll]`, runFact←null (I3 fresh-negative gate) |
| `RUN_FACT{runId}` | any | (same) | runFact←runId (pessimism toward an attempt) |
| `STOP_REQUESTED` (user Stop) | streaming, reconnecting, polling | stopping **†** | `[stopRun, abortAttach, cancelReconnect, armPoll]` (poll drives the terminal — I4 exit by data) |
| `SUPERSEDE_REQUESTED{targetRunId}` (interrupt+send) | streaming, reconnecting, polling, error | superseding **†** | `[supersede(target), cancelReconnect, disarmPoll]` |
| `SUPERSEDE_READY{runId}` (CAS ok) | superseding | streaming | ownership=local, runFact←runId |
| `SUPERSEDE_MISMATCH{currentRunId}` (409 SUPERSEDE_TARGET_MISMATCH) | superseding | error(supersede-mismatch) | `[postRun(verify)]`, runFact←currentRunId |
| `SUPERSEDE_TIMEOUT` (409 SUPERSEDE_TIMEOUT) | superseding | error(supersede-timeout) | — (composer keeps text; no auto-retry) |
| `SUPERSEDE_INVALID` (409 SUPERSEDE_INVALID) | superseding | error(supersede-invalid) | — |
| `RUN_ALREADY_ACTIVE` (409 A_RUN_ALREADY_ACTIVE, plain POST) | sending | error(run-already-active) | — (composer offers supersede; NO auto-retry) |
| `DISPOSE` (unmount) | any | idle **†** | `[abortAttach, cancelReconnect, disarmPoll]` (I1/I5 — epoch++ kills late callbacks) |
**`stopping` honors any finish (re-review MEDIUM):** BEFORE the epoch filter, a
stream finish (`FINISH_*`/`STREAM_INCOMPLETE`) arriving in phase `stopping` exits
`stopping -> idle` regardless of generation. A plain Stop has no successor stream,
so the aborted stream's finish IS the expected end (I4 exit by data) — and it
carries the PRE-stop generation (STOP_REQUESTED bumped the epoch), so the filter
would otherwise strand the machine in `stopping` (no idle-cap covers it). The filter
stays in force for `superseding` (that is the F1 supersede drop).
**Epoch filter (I1):** the reducer then drops any event carrying an `epoch` that
does not equal the current `ctx.epoch`. Outcome events (`STREAM_START`, `ATTACH_*`,
`RECONNECT_*`, `SUPERSEDE_*`, **`FINISH_*`/`STREAM_INCOMPLETE`**, `RUN_FACT`) are
stamped with the generation the corresponding STREAM started under (the runtime
holds a per-owned-stream `turnEpoch`); trigger events (user actions, fresh
disconnects) carry no epoch. **F1:** this is what makes a SUPERSEDED stream's late
`onFinish` (a dead stream A closing after the CAS started stream B) get dropped, so
A cannot drive the live new run into a false reconnect or reset its run-fact. The
supersede path additionally ABORTS A and starts B only from A's onFinish (a
microtask), because ai@6 `AbstractChat.makeRequest` corrupts overlapping streams
(A's `finally` reads then nulls the shared `activeResponse`).
**Removed events (scope-cut, internal review):** `RUN_SUPERSEDED` (a ghost feature —
never dispatched; the observer-superseded case is handled by the degraded poll,
which follows the latest rows regardless of runId), `RECONNECT_BEGIN` (reconnect is
entered by `FINISH_DISCONNECT`), and `POLL_ACTIVITY` (the window's activity clock was
removed when the idle-cap moved into the thread). The reducer and this table now
share exactly the dispatched event set.
### 409-code → event map (the real #487 contract consumed here)
| Server response | Event dispatched | error kind → banner |
|---|---|---|
| 409 `A_RUN_ALREADY_ACTIVE` (plain POST) | `RUN_ALREADY_ACTIVE` | run-already-active → "already answering / interrupt & send" |
| 409 `SUPERSEDE_TARGET_MISMATCH` (+ body.runId) | `SUPERSEDE_MISMATCH{currentRunId}` | supersede-mismatch → verify via /run |
| 409 `SUPERSEDE_TIMEOUT` | `SUPERSEDE_TIMEOUT` | supersede-timeout → "couldn't interrupt in time, resend" |
| 409 `SUPERSEDE_INVALID` | `SUPERSEDE_INVALID` | supersede-invalid → "couldn't interrupt this run" |
| 503 `A_RUN_BEGIN_FAILED` | `FINISH_ERROR{begin-failed}` | begin-failed → "could not start, temporary" |
---
## 2. Ref-map — every `chat-thread.tsx` ref → its new home (MIGRATION RESOLVED)
The migration is COMPLETE: the 13 run-lifecycle FLAGS below are GONE from
`chat-thread.tsx` (collapsed into FSM phase/ctx/effects, or deleted). What remains
are identity/data mirrors, effect-owned controllers/timers, and ONE React-liveness
bit — none of which is a run-lifecycle flag, so the post-merge "no new flags" rule
holds. **Pending column: empty.**
| # | Old ref | Resolved to | Where now |
|---|---|---|---|
| 1 | `reconcileTailRef` | **FSM phase** | reconcile-merge gated on `phase ∈ {polling, reconnecting, stopping}` |
| 2 | `noStreamHandledRef` | **FSM epoch (I1)** | the attach outcome's epoch guard drops the stale/second outcome |
| 3 | `onNoActiveStreamRef` | **FSM event** | transport → `handleAttachOutcome` dispatches `ATTACH_NONE`/`RECONNECT_NONE` |
| 4 | `onReconnectAttachedRef` | **FSM event** | transport dispatches `ATTACH_LIVE` / `RECONNECT_ATTACHED` |
| 5 | `resumedTurnRef` + `resumedTurn` state | **FSM ctx `ownership`** | `ownership==='observer'` ⇒ never flush; hides "Send now" |
| 6 | `reconnectStateRef` + `reconnectState` state | **FSM phase** | `reconnecting(attempt,failed)` renders the banner |
| 7 | `reconnectTimerRef` | **effect-owned timer** | owned by `scheduleReconnect`/`cancelReconnect` effects (not a flag) |
| 8 | `flushOnAbortRef` | **DELETED** | the stop→flush dance is replaced by the CAS supersede (commit 5) |
| 9 | `interruptNextSendRef` | **DELETED** | the server injects the interrupt note from the supersede itself |
| 10 | `supersedeRetryRef` | **DELETED** (commit 5) | the client 409 retry ladder is gone; CAS supersede replaces it |
| 11 | `stopPendingRef` | **FSM phase `stopping`** | the deferred stop fires from the chat-id adoption effect while `stopping` |
| 12 | `mountedRef` | **retained (React liveness)** | orthogonal to run-lifecycle; gates imperative onFinish side-effects post-unmount. Epoch (I1) handles stale COMMAND-outcomes; DISPOSE bumps it |
| 13 | `attemptResumeRef` | **FSM `ATTACH_START` + run-fact** | mount arms attach ONLY on a confirmed active run (commit 4b: streaming-tail status, or POST /run for a user tail) |
| 14 | `stripRef` | **data** (attachStrategy) | strip+replay detail; the `resumeStream` effect reads it |
| 15 | `strippedRowRef` | **data** (attachStrategy) | the anchor row |
| 16 | `attachAbortRef` | **effect-owned controller** | aborted by the `abortAttach` effect in cleanup (I5) |
| 17–25 | `chatIdRef`, `openPageRef`, `getEditorSelectionRef`, `roleIdRef`, `stableIdRef`, `queuedRef`, `sendMessageRef`, `statusRef`, `lastForwardedChatIdRef` | **data** (identity/send mirrors) | unchanged — not lifecycle flags |
| NEW | `pendingSupersedeRef` | **data** (send-plumbing) | the runId injected into the next `POST /stream {supersede}`; the single replacement for the 3 DELETED one-shots (#8/#9/#10) — net −2 refs |
| NEW | `idleCapTimerRef` | **effect-owned timer** | the stalled inactivity cap → `POLL_IDLE_CAP` (commit 4a); not a flag |
Net: the 13 lifecycle flags (#1#13) are eliminated: **8** → FSM phase/ctx/epoch/event
(#1#6, #11, #13), **3** deleted (#8/#9/#10), **`reconnectTimerRef` (#7)** becomes an
effect-owned controller, and **`mountedRef` (#12)** is retained as React liveness
(8 + 3 + 1 + 1 = 13). (`attachAbortRef` (#16) is outside the #1#13 set — it was
already an effect-owned controller.) Two effect-owned timers + one send-plumbing data
ref are added — none is a boolean lifecycle latch.
---
## 3. Run-fact protocol (`runFact: {runId} | null`) — I3
"A run is active" is first-class from the SERVER, not inferred from an assistant
message. Sources, in the order they update `ctx.runFact`:
1. **Init (mount):** `POST /ai-chat/run { chatId }``{ run, message }`. A `run`
with a non-terminal `status` seeds `runFact = { runId: run.id }`; a null/terminal
run seeds `null`. This is what arms the resume attempt (`ATTACH_START`) — the
attempt is armed ONLY on a positive fact (commit 4b: a user-tail with no active
run no longer arms a pointless poll on every open).
2. **Live update:** the `start` stream metadata carries `runId``STREAM_START{runId}`.
3. **Attach outcomes:** `ATTACH_LIVE` (2xx) confirms active; a 204 on a non-stripped
path is an authoritative NEGATIVE fact → the runtime dispatches `RUN_FACT{null}`,
which cancels recovery (I3 fresh-negative gate).
4. **Poll (future resume-stack iteration #491):** the delta will carry the run field;
until then the poll drives to a terminal ROW, dispatched as `POLL_TERMINAL`.
Pessimism rule: a stale-but-positive fact PERMITS entering recovery (attach); the
204 then cuts it. A fresh negative fact gates recovery OUT immediately.
---
## 4. Invariants
- **I1 — Epoch (generation counter).** Every command-emitting transition bumps
`ctx.epoch`; every async outcome event carries its issuing epoch; the reducer
drops stale-epoch outcomes. Replaces the one-shot-ref zoo (`noStreamHandledRef`,
the flush/interrupt/supersede one-shots, the `mountedRef` late-callback gate).
- **I2 — Ownership is context, not state.** `local | observer` is orthogonal to the
transport phase. The queue flushes ONLY under local ownership; an observer
following a detached run never flushes (was `resumedTurnRef`).
- **I3 — Run-fact is first-class from the server.** Reconnect is entered by the
run-fact, not by an assistant message (commit 2). A fresh negative fact cancels
recovery.
- **I4 — Exit `stopping` by DATA.** A terminal row / negative run-fact / terminal
finish exits `stopping`, never the stopRun HTTP response (which returns after the
abort but before finalization — keying off it would unlock the composer on a 409).
- **I5 — Dispose protocol.** Command controllers (attach GET, POST /stream, POST
/run) are effect-owned and aborted in cleanup (`abortAttach` on `DISPOSE`), not
render-phase refs. A client abort of an already-sent POST does not cancel the
server action, so disarming on unmount is safe.
- **attachStrategy** (strip+replay today) is behind the `resumeStream` effect; the
resume-stack iteration (#491) swaps it to tail-only WITHOUT touching the FSM.
- **Queue** stays a data structure; flush/interrupt decisions are transitions.
@@ -1,461 +0,0 @@
import { describe, it, expect } from "vitest";
import {
reduce,
initialMachine,
reconnectDelayMs,
RECONNECT_MAX_ATTEMPTS,
type Machine,
type Effect,
type Event,
} from "./run-fsm";
// Drive a sequence of events through the reducer, returning the final machine.
function run(m: Machine, ...events: Event[]): Machine {
return events.reduce(reduce, m);
}
function withRunFact(runId = "run-1"): Machine {
return {
...initialMachine(),
ctx: { epoch: 0, ownership: "local", runFact: { runId }, liveFollow: false },
};
}
function effectTypes(m: Machine): string[] {
return m.effects.map((e) => e.type);
}
function hasEffect(m: Machine, type: Effect["type"]): boolean {
return m.effects.some((e) => e.type === type);
}
describe("run-fsm — epoch invariant (I1)", () => {
it("drops an outcome carrying a stale epoch", () => {
// A command bumps the epoch; an outcome stamped with the OLD epoch is dropped.
const m0 = reduce(initialMachine(), { type: "ATTACH_START", runId: "r" }); // epoch 0->1, attaching
expect(m0.ctx.epoch).toBe(1);
expect(m0.phase.name).toBe("attaching");
// A late ATTACH_LIVE from a SUPERSEDED attempt (epoch 0) must NOT drive us.
const stale = reduce(m0, { type: "ATTACH_LIVE", epoch: 0 });
expect(stale.phase.name).toBe("attaching");
expect(stale.effects).toEqual([]);
});
it("applies an outcome carrying the current epoch", () => {
const m0 = reduce(initialMachine(), { type: "ATTACH_START", runId: "r" });
const live = reduce(m0, { type: "ATTACH_LIVE", epoch: m0.ctx.epoch });
expect(live.phase.name).toBe("streaming");
});
it("an outcome with no epoch is never dropped (trigger events)", () => {
const m0 = reduce(initialMachine(), { type: "ATTACH_START", runId: "r" });
const disposed = reduce(m0, { type: "DISPOSE" });
expect(disposed.phase.name).toBe("idle");
expect(hasEffect(disposed, "abortAttach")).toBe(true);
});
it("every command-transition increments the epoch exactly once", () => {
let m = initialMachine();
const before = m.ctx.epoch;
m = reduce(m, { type: "SEND_LOCAL" });
expect(m.ctx.epoch).toBe(before + 1);
m = reduce(m, { type: "STOP_REQUESTED" });
expect(m.ctx.epoch).toBe(before + 2);
});
});
describe("run-fsm — local turn", () => {
it("SEND_LOCAL → sending, local ownership, cancels recovery", () => {
const m = reduce(withRunFact(), { type: "SEND_LOCAL" });
expect(m.phase.name).toBe("sending");
expect(m.ctx.ownership).toBe("local");
expect(effectTypes(m)).toEqual(
expect.arrayContaining(["cancelReconnect", "disarmPoll"]),
);
});
it("STREAM_START adopts the runId into the run-fact and goes streaming", () => {
const m = run(initialMachine(), { type: "SEND_LOCAL" });
const s = reduce(m, { type: "STREAM_START", runId: "run-9", epoch: m.ctx.epoch });
expect(s.phase.name).toBe("streaming");
expect(s.ctx.runFact).toEqual({ runId: "run-9" });
});
it("FINISH_CLEAN → idle, run-fact cleared, poll/reconnect disarmed", () => {
const streaming = run(initialMachine(), { type: "SEND_LOCAL" }, { type: "STREAM_START", runId: "r" });
const done = reduce(streaming, { type: "FINISH_CLEAN" });
expect(done.phase.name).toBe("idle");
expect(done.ctx.runFact).toBeNull();
});
});
// #488 commit 2 — SSE break BEFORE the first assistant frame must still recover.
describe("run-fsm — commit 2: reconnect by run-fact, not by assistant message", () => {
it("FINISH_DISCONNECT with an active run-fact → reconnecting (even with no visible content)", () => {
// Setup-phase break: no assistant frame yet, but a run-fact exists.
const streaming = withRunFact("run-2");
const m = reduce(streaming, {
type: "FINISH_DISCONNECT",
hasVisibleContent: false,
epoch: streaming.ctx.epoch,
});
expect(m.phase.name).toBe("reconnecting");
if (m.phase.name === "reconnecting") expect(m.phase.attempt).toBe(1);
expect(m.ctx.ownership).toBe("observer");
expect(hasEffect(m, "scheduleReconnect")).toBe(true);
// No visible content -> no poll arm yet (the reconnect ladder rebuilds it).
expect(hasEffect(m, "armPoll")).toBe(false);
});
it("FINISH_DISCONNECT WITH visible content also arms the poll", () => {
const m = reduce(withRunFact("run-2"), {
type: "FINISH_DISCONNECT",
hasVisibleContent: true,
epoch: 0,
});
expect(m.phase.name).toBe("reconnecting");
expect(hasEffect(m, "armPoll")).toBe(true);
});
it("FINISH_DISCONNECT with NO run-fact → idle (plain connection-lost)", () => {
const m = reduce(initialMachine(), {
type: "FINISH_DISCONNECT",
hasVisibleContent: true,
epoch: 0,
});
expect(m.phase.name).toBe("idle");
});
});
// #488 commit 3 — a SECOND break after a successful re-attach starts a NEW ladder.
describe("run-fsm — commit 3: repeated reconnect cycles", () => {
it("two breaks in a row produce two reconnect cycles (counter resets on attach)", () => {
let m = withRunFact("run-3");
// First break -> reconnecting(1).
m = reduce(m, { type: "FINISH_DISCONNECT", hasVisibleContent: false, epoch: m.ctx.epoch });
expect(m.phase.name).toBe("reconnecting");
// Attempt fires, re-attaches live.
m = reduce(m, { type: "RECONNECT_ATTEMPT", attempt: 1, epoch: m.ctx.epoch });
m = reduce(m, { type: "RECONNECT_ATTACHED", epoch: m.ctx.epoch });
expect(m.phase.name).toBe("streaming");
// SECOND break: the counter was reset, so a fresh ladder starts at attempt 1
// (the old one-shot !wasResumed gate would have sent this to silent poll).
m = reduce(m, { type: "FINISH_DISCONNECT", hasVisibleContent: false, epoch: m.ctx.epoch });
expect(m.phase.name).toBe("reconnecting");
if (m.phase.name === "reconnecting") expect(m.phase.attempt).toBe(1);
expect(hasEffect(m, "scheduleReconnect")).toBe(true);
});
it("a MOUNT-attach observer drop falls to POLL, not the reconnect ladder", () => {
// Distinguishes commit 3 from a one-shot resume: an observer that never
// live-followed (liveFollow false) polls on a drop.
let m = reduce(initialMachine(), { type: "ATTACH_START", runId: "r" });
m = reduce(m, { type: "ATTACH_LIVE", epoch: m.ctx.epoch });
expect(m.ctx.ownership).toBe("observer");
expect(m.ctx.liveFollow).toBe(false);
m = reduce(m, { type: "FINISH_DISCONNECT", hasVisibleContent: true, epoch: m.ctx.epoch });
expect(m.phase.name).toBe("polling");
expect(hasEffect(m, "armPoll")).toBe(true);
});
it("STREAM_INCOMPLETE (observer starved/torn finish) → polling", () => {
let m = reduce(initialMachine(), { type: "ATTACH_START", runId: "r" });
m = reduce(m, { type: "ATTACH_LIVE", epoch: m.ctx.epoch });
m = reduce(m, { type: "STREAM_INCOMPLETE", reason: "starved", epoch: m.ctx.epoch });
expect(m.phase).toEqual({ name: "polling", reason: "starved" });
expect(hasEffect(m, "armPoll")).toBe(true);
});
it("liveFollow is set on the first local drop and kept across a re-attach", () => {
let m = withRunFact("run-3");
m = reduce(m, { type: "FINISH_DISCONNECT", hasVisibleContent: false, epoch: m.ctx.epoch });
expect(m.ctx.liveFollow).toBe(true);
m = reduce(m, { type: "RECONNECT_ATTEMPT", attempt: 1, epoch: m.ctx.epoch });
m = reduce(m, { type: "RECONNECT_ATTACHED", epoch: m.ctx.epoch });
expect(m.ctx.liveFollow).toBe(true); // kept — so a second drop reconnects
// A clean finish clears it.
m = reduce(m, { type: "FINISH_CLEAN", epoch: m.ctx.epoch });
expect(m.ctx.liveFollow).toBe(false);
});
it("RECONNECT_NONE backs off through the ladder, then fails at the cap", () => {
let m = withRunFact("run-3");
m = reduce(m, { type: "FINISH_DISCONNECT", hasVisibleContent: false, epoch: m.ctx.epoch });
for (let n = 1; n < RECONNECT_MAX_ATTEMPTS; n++) {
m = reduce(m, { type: "RECONNECT_ATTEMPT", attempt: n, epoch: m.ctx.epoch });
m = reduce(m, { type: "RECONNECT_NONE", epoch: m.ctx.epoch });
expect(m.phase.name).toBe("reconnecting");
if (m.phase.name === "reconnecting") {
expect(m.phase.attempt).toBe(n + 1);
expect(m.phase.failed).toBe(false);
}
// The belt-and-suspenders poll is armed each failed attempt.
expect(hasEffect(m, "armPoll")).toBe(true);
}
// Final attempt fails -> failed banner (Retry), poll armed.
m = reduce(m, { type: "RECONNECT_ATTEMPT", attempt: RECONNECT_MAX_ATTEMPTS, epoch: m.ctx.epoch });
m = reduce(m, { type: "RECONNECT_NONE", epoch: m.ctx.epoch });
expect(m.phase.name).toBe("reconnecting");
if (m.phase.name === "reconnecting") expect(m.phase.failed).toBe(true);
// RETRY restarts at attempt 1.
m = reduce(m, { type: "RETRY" });
expect(m.phase.name).toBe("reconnecting");
if (m.phase.name === "reconnecting") {
expect(m.phase.attempt).toBe(1);
expect(m.phase.failed).toBe(false);
}
expect(hasEffect(m, "resumeStream")).toBe(true);
});
it("reconnectDelayMs is the exponential backoff 1s,2s,4s,8s,16s", () => {
expect([1, 2, 3, 4, 5].map(reconnectDelayMs)).toEqual([1000, 2000, 4000, 8000, 16000]);
});
});
// #488 commit 4 — polling stalled-state + user-tail gating.
describe("run-fsm — commit 4: stalled + run-fact gating", () => {
it("POLL_IDLE_CAP: polling → stalled with a banner (poll disarmed), not silent", () => {
let m = reduce(withRunFact(), { type: "ATTACH_START", runId: "r" });
m = reduce(m, { type: "ATTACH_NONE", epoch: m.ctx.epoch });
expect(m.phase.name).toBe("polling");
m = reduce(m, { type: "POLL_IDLE_CAP" });
expect(m.phase.name).toBe("stalled");
expect(hasEffect(m, "disarmPoll")).toBe(true);
});
it("RETRY from stalled re-arms the poll", () => {
let m = reduce(withRunFact(), { type: "ATTACH_START", runId: "r" });
m = reduce(m, { type: "ATTACH_NONE", epoch: m.ctx.epoch });
m = reduce(m, { type: "POLL_IDLE_CAP" });
m = reduce(m, { type: "RETRY" });
expect(m.phase.name).toBe("polling");
expect(hasEffect(m, "armPoll")).toBe(true);
});
it("a fresh NEGATIVE run-fact while attaching cancels recovery (user-tail, no active run)", () => {
// The mount POST /run returns no active run: attaching → idle, no poll armed.
let m = reduce(withRunFact(), { type: "ATTACH_START", runId: "r" });
m = reduce(m, { type: "RUN_FACT", runFact: null, epoch: m.ctx.epoch });
expect(m.phase.name).toBe("idle");
expect(m.ctx.runFact).toBeNull();
expect(hasEffect(m, "disarmPoll")).toBe(true);
});
it("a negative run-fact while polling stops the poll", () => {
let m = reduce(withRunFact(), { type: "ATTACH_START", runId: "r" });
m = reduce(m, { type: "ATTACH_NONE", epoch: m.ctx.epoch });
m = reduce(m, { type: "RUN_FACT", runFact: null, epoch: m.ctx.epoch });
expect(m.phase.name).toBe("idle");
});
it("POLL_TERMINAL settles polling → idle (I4 data-driven exit)", () => {
let m = reduce(withRunFact(), { type: "ATTACH_START", runId: "r" });
m = reduce(m, { type: "ATTACH_NONE", epoch: m.ctx.epoch });
m = reduce(m, { type: "POLL_TERMINAL" });
expect(m.phase.name).toBe("idle");
expect(m.ctx.runFact).toBeNull();
});
});
// #488 commit 5 — error classification + supersede CAS transitions.
describe("run-fsm — commit 5: supersede CAS + error classification", () => {
it("SUPERSEDE_REQUESTED → superseding, fires the CAS effect, bumps epoch", () => {
const streaming = withRunFact("run-old");
const m = reduce(streaming, { type: "SUPERSEDE_REQUESTED", targetRunId: "run-old" });
expect(m.phase.name).toBe("superseding");
expect(m.ctx.epoch).toBe(streaming.ctx.epoch + 1);
const sup = m.effects.find((e) => e.type === "supersede");
expect(sup).toEqual({ type: "supersede", targetRunId: "run-old" });
});
it("SUPERSEDE_READY → streaming as the new local owner", () => {
let m = reduce(withRunFact("run-old"), { type: "SUPERSEDE_REQUESTED", targetRunId: "run-old" });
m = reduce(m, { type: "SUPERSEDE_READY", runId: "run-new", epoch: m.ctx.epoch });
expect(m.phase.name).toBe("streaming");
expect(m.ctx.ownership).toBe("local");
expect(m.ctx.runFact).toEqual({ runId: "run-new" });
});
it("SUPERSEDE_MISMATCH → error(supersede-mismatch) + verify via /run (no blind banner)", () => {
let m = reduce(withRunFact("run-old"), { type: "SUPERSEDE_REQUESTED", targetRunId: "run-old" });
m = reduce(m, { type: "SUPERSEDE_MISMATCH", currentRunId: "run-x", epoch: m.ctx.epoch });
expect(m.phase).toEqual({ name: "error", kind: "supersede-mismatch" });
expect(hasEffect(m, "postRun")).toBe(true);
expect(m.ctx.runFact).toEqual({ runId: "run-x" });
});
it("SUPERSEDE_TIMEOUT → error(supersede-timeout), no auto-retry effect", () => {
let m = reduce(withRunFact("run-old"), { type: "SUPERSEDE_REQUESTED", targetRunId: "run-old" });
m = reduce(m, { type: "SUPERSEDE_TIMEOUT", epoch: m.ctx.epoch });
expect(m.phase).toEqual({ name: "error", kind: "supersede-timeout" });
expect(m.effects).toEqual([]);
});
it("SUPERSEDE_INVALID → error(supersede-invalid)", () => {
let m = reduce(withRunFact("run-old"), { type: "SUPERSEDE_REQUESTED", targetRunId: "run-old" });
m = reduce(m, { type: "SUPERSEDE_INVALID", epoch: m.ctx.epoch });
expect(m.phase).toEqual({ name: "error", kind: "supersede-invalid" });
});
it("a stale SUPERSEDE outcome from a superseded epoch is dropped", () => {
let m = reduce(withRunFact("run-old"), { type: "SUPERSEDE_REQUESTED", targetRunId: "run-old" });
const supersedingEpoch = m.ctx.epoch;
// The user retriggers, bumping the epoch again.
m = reduce(m, { type: "SUPERSEDE_REQUESTED", targetRunId: "run-old" });
// The first CAS's late TIMEOUT (old epoch) must NOT knock us out of superseding.
const late = reduce(m, { type: "SUPERSEDE_TIMEOUT", epoch: supersedingEpoch });
expect(late.phase.name).toBe("superseding");
});
it("RUN_ALREADY_ACTIVE (plain POST gate) → error(run-already-active), no retry effect", () => {
const m = reduce(run(initialMachine(), { type: "SEND_LOCAL" }), { type: "RUN_ALREADY_ACTIVE" });
expect(m.phase).toEqual({ name: "error", kind: "run-already-active" });
expect(m.effects).toEqual([]);
});
});
// #488 F2 — a late mount `getRun → ATTACH_START` must not hijack a local turn.
describe("run-fsm — F2: ATTACH_START only from idle", () => {
it("ATTACH_START from a local `sending` turn is ignored (no observer hijack)", () => {
const sending = reduce(initialMachine(), { type: "SEND_LOCAL" }); // idle -> sending, local
const m = reduce(sending, { type: "ATTACH_START", runId: "r" });
expect(m.phase.name).toBe("sending");
expect(m.ctx.ownership).toBe("local"); // NOT flipped to observer
expect(m.effects).toEqual([]); // no resumeStream
});
it("ATTACH_START from idle attaches as normal", () => {
const m = reduce(initialMachine(), { type: "ATTACH_START", runId: "r" });
expect(m.phase.name).toBe("attaching");
expect(m.ctx.ownership).toBe("observer");
expect(hasEffect(m, "resumeStream")).toBe(true);
});
});
describe("run-fsm — stop (I4: exit by data)", () => {
it("STOP_REQUESTED → stopping, fires stopRun + abortAttach, no data-independent exit", () => {
const m = reduce(withRunFact(), { type: "STOP_REQUESTED" });
expect(m.phase.name).toBe("stopping");
expect(effectTypes(m)).toEqual(expect.arrayContaining(["stopRun", "abortAttach"]));
});
it("stopping exits on the aborted stream's finish carrying the PRE-STOP epoch", () => {
// MEDIUM (#488 re-review): STOP_REQUESTED is a command that BUMPS the epoch, but
// the runtime stamps the aborted stream's onFinish with the stream's START (pre-
// stop) generation — exactly what the component sends. `stopping` must HONOR
// that finish regardless of generation (no idle-cap covers `stopping`).
// MUTATION-VERIFY: remove the honor-in-`stopping` branch and this hangs in
// `stopping` (the epoch filter drops the pre-stop finish) -> red.
const preStopEpoch = withRunFact().ctx.epoch; // E1 (the stream's start epoch)
let m = reduce(withRunFact(), { type: "STOP_REQUESTED" }); // E1 -> E2, stopping
expect(m.ctx.epoch).toBe(preStopEpoch + 1);
m = reduce(m, { type: "FINISH_ABORT", epoch: preStopEpoch }); // NOT the current epoch
expect(m.phase.name).toBe("idle");
expect(m.ctx.runFact).toBeNull();
});
it("stopping exits on a clean finish carrying the pre-stop epoch too", () => {
const preStopEpoch = withRunFact().ctx.epoch;
let m = reduce(withRunFact(), { type: "STOP_REQUESTED" });
m = reduce(m, { type: "FINISH_CLEAN", epoch: preStopEpoch });
expect(m.phase.name).toBe("idle");
});
it("stopping exits on a negative run-fact (data)", () => {
let m = reduce(withRunFact(), { type: "STOP_REQUESTED" });
m = reduce(m, { type: "RUN_FACT", runFact: null, epoch: m.ctx.epoch });
expect(m.phase.name).toBe("idle");
});
// Review #4: `stopping` arms the poll but had no inactivity backstop.
it("review-4: POLL_IDLE_CAP in `stopping` exits to idle (bounded), NOT stalled", () => {
let m = reduce(withRunFact(), { type: "STOP_REQUESTED" });
expect(m.phase.name).toBe("stopping");
expect(hasEffect(m, "armPoll")).toBe(true);
// MUTATION-VERIFY: drop the `stopping` branch in POLL_IDLE_CAP and this hangs
// in `stopping` (poll forever) -> red.
m = reduce(m, { type: "POLL_IDLE_CAP" });
expect(m.phase.name).toBe("idle");
expect(hasEffect(m, "disarmPoll")).toBe(true);
expect(m.ctx.ownership).toBe("local");
});
});
// Review #1: positive attach outcomes must be guarded by the SOURCE phase — the
// epoch filter alone is insufficient because POLL_TERMINAL uses to() (no epoch
// bump) and does not abort the in-flight GET.
describe("run-fsm — review-1: attach outcomes guarded by source phase", () => {
it("a late RECONNECT_ATTACHED after POLL_TERMINAL stays idle (no phantom streaming)", () => {
let m = withRunFact("run-1");
m = reduce(m, { type: "FINISH_DISCONNECT", hasVisibleContent: true, epoch: m.ctx.epoch });
m = reduce(m, { type: "RECONNECT_ATTEMPT", attempt: 1, epoch: m.ctx.epoch }); // attach GET
const epoch = m.ctx.epoch;
// The armed degraded poll reaches the terminal row FIRST (epoch unchanged).
m = reduce(m, { type: "POLL_TERMINAL" });
expect(m.phase.name).toBe("idle");
expect(m.ctx.epoch).toBe(epoch); // POLL_TERMINAL did NOT bump the epoch
// The slow GET returns live 2xx under the SAME epoch — must NOT resurrect.
m = reduce(m, { type: "RECONNECT_ATTACHED", epoch });
expect(m.phase.name).toBe("idle");
});
it("a late ATTACH_LIVE / ATTACH_NONE after leaving `attaching` is ignored", () => {
let m = reduce(initialMachine(), { type: "ATTACH_START", runId: "r" });
const epoch = m.ctx.epoch;
m = reduce(m, { type: "ATTACH_NONE", epoch }); // attaching -> polling
m = reduce(m, { type: "POLL_TERMINAL" }); // -> idle (epoch unchanged)
expect(m.phase.name).toBe("idle");
m = reduce(m, { type: "ATTACH_LIVE", epoch }); // late 2xx, same epoch
expect(m.phase.name).toBe("idle");
// And a late ATTACH_NONE (not `attaching`) is a no-op too.
m = reduce(m, { type: "ATTACH_NONE", epoch });
expect(m.phase.name).toBe("idle");
});
});
// Review #2: every terminal transition resets ownership to local.
describe("run-fsm — review-2: terminal transitions reset ownership to local", () => {
const observer = (): Machine => {
let m = reduce(initialMachine(), { type: "ATTACH_START", runId: "r" });
m = reduce(m, { type: "ATTACH_LIVE", epoch: m.ctx.epoch });
expect(m.ctx.ownership).toBe("observer");
return m;
};
it("FINISH_CLEAN resets ownership", () => {
const m = reduce(observer(), { type: "FINISH_CLEAN", epoch: observer().ctx.epoch });
expect(m.ctx.ownership).toBe("local");
});
it("FINISH_ERROR / POLL_TERMINAL / RUN_FACT(null) reset ownership", () => {
let o = observer();
expect(reduce(o, { type: "FINISH_ERROR", kind: "stream", epoch: o.ctx.epoch }).ctx.ownership).toBe("local");
// POLL_TERMINAL from an observer polling phase
let p = reduce(observer(), { type: "STREAM_INCOMPLETE", reason: "starved", epoch: observer().ctx.epoch });
expect(reduce(p, { type: "POLL_TERMINAL" }).ctx.ownership).toBe("local");
// RUN_FACT(null) from an observer attaching phase
let a = reduce(initialMachine(), { type: "ATTACH_START", runId: "r" });
expect(reduce(a, { type: "RUN_FACT", runFact: null, epoch: a.ctx.epoch }).ctx.ownership).toBe("local");
});
});
describe("run-fsm — ownership (I2) is context, orthogonal to phase", () => {
it("attach/reconnect set observer; send/supersede-ready set local", () => {
let m = reduce(initialMachine(), { type: "ATTACH_START", runId: "r" });
expect(m.ctx.ownership).toBe("observer");
m = reduce(m, { type: "ATTACH_LIVE", epoch: m.ctx.epoch });
expect(m.phase.name).toBe("streaming");
expect(m.ctx.ownership).toBe("observer"); // still observing a detached run
// A local send flips ownership back to local.
m = reduce(m, { type: "SEND_LOCAL" });
expect(m.ctx.ownership).toBe("local");
});
});
describe("run-fsm — dispose (I5)", () => {
it("DISPOSE from any phase aborts controllers and bumps epoch", () => {
let m = reduce(withRunFact(), { type: "ATTACH_START", runId: "r" });
const before = m.ctx.epoch;
m = reduce(m, { type: "DISPOSE" });
expect(m.phase.name).toBe("idle");
expect(m.ctx.epoch).toBe(before + 1);
expect(effectTypes(m)).toEqual(
expect.arrayContaining(["abortAttach", "cancelReconnect", "disarmPoll"]),
);
});
});
@@ -1,595 +0,0 @@
/**
* Run-lifecycle finite state machine for a single AI-chat thread (#488).
*
* ============================================================================
* WHY THIS EXISTS
* ----------------------------------------------------------------------------
* The resume/reconnect/poll/stop/supersede lifecycle used to be spread across
* ~26 `useRef` one-shot flags in `chat-thread.tsx`, each disarmed "on every
* path". Ownerless flag combinations produced silent UI freezes, and every fix
* added another ref (the #381 -> #432 -> #456 spiral). This module replaces that
* ref-zoo with ONE pure reducer whose transitions are enumerable and unit-
* testable in isolation (event x state -> next state is the observable property).
*
* The reducer is PURE: it owns no timers, no fetches, no React state. It maps
* `(machine, event) -> machine`, where the returned machine carries the list of
* COMMAND EFFECTS to run for that transition. A thin runtime in `chat-thread.tsx`
* dispatches events (from SDK callbacks / HTTP outcomes) and executes the
* effects (attach GET, POST /stream, POST /run, POST /stop, backoff timers,
* poll arm/disarm). The runtime lives in a THREAD, not the window, so a late SDK
* callback dies with the owner (kills the "event from a dead view" class, #161).
*
* ============================================================================
* INVARIANTS (see run-fsm.spec.md for the full spec + tables)
* ----------------------------------------------------------------------------
* I1 EPOCH (generation counter). Commands (`resumeStream`, `postRun`, `stop`,
* `supersede`, `scheduleReconnect`) are async; their outcomes arrive on the
* SAME SDK/HTTP callbacks. Every command-emitting transition increments
* `ctx.epoch`; every OUTCOME event carries the epoch it was issued under;
* the reducer DROPS an outcome whose epoch != the current epoch. This is
* what the one-shot-ref zoo used to approximate by hand.
* I2 OWNERSHIP is a CONTEXT FIELD (`'local' | 'observer'`), not a state —
* orthogonal to the transport phase. The queue is flushed ONLY by a local
* owner (an observer following a detached run never flushes).
* I3 RUN-FACT ("a run is active") is first-class from the server: `runFact`
* holds the server-confirmed active run id (POST /run on mount, the `start`
* metadata runId, attach outcomes). Reconnect is entered by the RUN-FACT,
* not by the presence of an assistant message (#488 commit 2). A fresh
* negative fact (null) cancels reconnect immediately.
* I4 Exit `stopping` by DATA (a terminal row / negative run-fact), NEVER by the
* stopRun HTTP response (which returns after abort, before finalization).
* I5 Command controllers are effect-owned (abort in cleanup), NOT render-phase
* refs — expressed here as the `abortAttach` effect on disposing transitions.
* ============================================================================
*/
// ---------------------------------------------------------------------------
// Phases (the transport lifecycle). Ownership / runFact are CONTEXT, not here.
// ---------------------------------------------------------------------------
/** Why the degraded poll is the active recovery. */
export type PollReason =
| "attach-none" // mount attach returned 204 / error — nothing live to attach
| "starved" // a resumed finish carried no visible content
| "disconnect-visible" // a live disconnect WITH on-screen content — poll to terminal
| "reconnect-exhausted"; // the live re-attach ladder gave up
/** The classified error kind (drives the banner text + composer behavior). */
export type ErrorKind =
| "stream" // a generic provider/network stream error (useChat error)
| "run-already-active" // 409 A_RUN_ALREADY_ACTIVE (a plain POST hit the gate)
| "supersede-mismatch" // 409 SUPERSEDE_TARGET_MISMATCH (CAS target moved)
| "supersede-timeout" // 409 SUPERSEDE_TIMEOUT (old run did not settle in W)
| "supersede-invalid" // 409 SUPERSEDE_INVALID (bad supersede target)
| "begin-failed"; // 503 A_RUN_BEGIN_FAILED (could not start the run)
export type Phase =
| { name: "idle" }
| { name: "sending" } // local POST in flight, before the first frame
| { name: "streaming" } // receiving frames
| { name: "attaching" } // mount-time attach GET in flight
| { name: "reconnecting"; attempt: number; failed: boolean }
| { name: "polling"; reason: PollReason }
| { name: "stalled" } // poll hit the inactivity cap — banner + Retry
| { name: "stopping" }
| { name: "superseding" }
| { name: "error"; kind: ErrorKind };
export type Ownership = "local" | "observer";
/** The server-confirmed active run, or null when no run is active. */
export type RunFact = { runId: string } | null;
export interface Ctx {
/** I1: generation counter — every command-transition increments it. */
epoch: number;
/** I2: does THIS client own the turn's writes (local streamer) or observe? */
ownership: Ownership;
/** I3: the server-confirmed active run. */
runFact: RunFact;
/**
* Are we FOLLOWING a live run we were locally streaming (the reconnect ladder),
* as opposed to a one-shot mount-attach resume? Both are `ownership: 'observer'`,
* but they recover DIFFERENTLY on a drop: a live-follow drop RE-ENTERS the
* reconnect ladder (#488 commit 3 — the second break after a successful re-attach
* must reconnect again, not fall to silent poll), while a mount-resume drop falls
* to the degraded poll. This is the ctx bit that separates the two WITHOUT a new
* component ref (it is why commit 3 needs the FSM, not a surgical patch).
*/
liveFollow: boolean;
}
export interface Machine {
phase: Phase;
ctx: Ctx;
/** Command effects to run for the transition that produced THIS machine.
* The runtime executes them and does not read them again. */
effects: Effect[];
}
// ---------------------------------------------------------------------------
// Command effects (the reducer's only side-channel — executed by the runtime).
// ---------------------------------------------------------------------------
export type Effect =
/** POST /run to (re)establish or verify the run-fact. `reason` is diagnostic. */
| { type: "postRun"; reason: "mount" | "verify" }
/** Trigger the SDK `resumeStream()` (attach GET via prepareReconnectToStream). */
| { type: "resumeStream" }
/** Schedule a reconnect attempt after a backoff, then dispatch RECONNECT_ATTEMPT. */
| { type: "scheduleReconnect"; attempt: number; delayMs: number }
/** Cancel any pending reconnect backoff timer. */
| { type: "cancelReconnect" }
/** Arm the degraded poll (the window's dumb timer follows the run in the DB). */
| { type: "armPoll"; reason: PollReason }
/** Disarm the degraded poll. */
| { type: "disarmPoll" }
/** POST /stop the chat's active run (authoritative detached-run stop). */
| { type: "stopRun" }
/** POST /stream { supersede: { runId } } — the CAS "interrupt and send now". */
| { type: "supersede"; targetRunId: string }
/** Abort the in-flight attach/reconnect GET controller (dispose / observer stop). */
| { type: "abortAttach" };
// ---------------------------------------------------------------------------
// Events. An OUTCOME event MAY carry `epoch`; if it does and it does not equal
// the current epoch, the reducer drops it (I1). Trigger events (user actions,
// fresh disconnects) carry no epoch and are never dropped.
// ---------------------------------------------------------------------------
export type Event =
// -- local turn --
| { type: "SEND_LOCAL" }
| { type: "STREAM_START"; runId?: string; epoch?: number }
/** An OBSERVER's attached stream ended WITHOUT reaching terminal (a starved
* clean replay, or a torn resume) — fall to the degraded poll to drive the row
* to its real terminal state. (A live-follow drop uses FINISH_DISCONNECT.) */
| { type: "STREAM_INCOMPLETE"; reason: PollReason; epoch?: number }
| { type: "FINISH_CLEAN"; epoch?: number }
| { type: "FINISH_ABORT"; epoch?: number }
| { type: "FINISH_DISCONNECT"; hasVisibleContent: boolean; epoch?: number }
| { type: "FINISH_ERROR"; kind: ErrorKind; epoch?: number }
// -- mount attach (resume) --
| { type: "ATTACH_START"; runId?: string }
| { type: "ATTACH_LIVE"; epoch?: number }
| { type: "ATTACH_NONE"; epoch?: number }
// -- reconnect after a live disconnect (entered by FINISH_DISCONNECT, #488 c2) --
| { type: "RECONNECT_ATTEMPT"; attempt: number; epoch?: number }
| { type: "RECONNECT_ATTACHED"; epoch?: number }
| { type: "RECONNECT_NONE"; epoch?: number }
| { type: "RETRY" }
// -- degraded poll --
| { type: "POLL_TERMINAL" }
| { type: "POLL_IDLE_CAP" }
// -- run-fact (server-confirmed active run) --
| { type: "RUN_FACT"; runFact: RunFact; epoch?: number }
// -- stop --
| { type: "STOP_REQUESTED" }
// -- supersede (CAS) --
| { type: "SUPERSEDE_REQUESTED"; targetRunId: string }
| { type: "SUPERSEDE_READY"; runId?: string; epoch?: number }
| { type: "SUPERSEDE_MISMATCH"; currentRunId?: string; epoch?: number }
| { type: "SUPERSEDE_TIMEOUT"; epoch?: number }
| { type: "SUPERSEDE_INVALID"; epoch?: number }
| { type: "RUN_ALREADY_ACTIVE" }
// -- lifecycle --
| { type: "DISPOSE" };
export const RECONNECT_MAX_ATTEMPTS = 5;
export const RECONNECT_BASE_DELAY_MS = 1000;
/** Backoff before attempt N (1-based): 1s, 2s, 4s, 8s, 16s. */
export function reconnectDelayMs(attempt: number): number {
return RECONNECT_BASE_DELAY_MS * 2 ** (attempt - 1);
}
// ---------------------------------------------------------------------------
// Constructors / helpers.
// ---------------------------------------------------------------------------
export function initialMachine(overrides?: Partial<Ctx>): Machine {
return {
phase: { name: "idle" },
ctx: { epoch: 0, ownership: "local", runFact: null, liveFollow: false, ...overrides },
effects: [],
};
}
/** Build a machine result: a phase, optional ctx patch, and effects. Empty
* effects by default. Never mutates the input. */
function to(
m: Machine,
phase: Phase,
opts?: { ctx?: Partial<Ctx>; effects?: Effect[] },
): Machine {
return {
phase,
ctx: { ...m.ctx, ...(opts?.ctx ?? {}) },
effects: opts?.effects ?? [],
};
}
/** No transition: keep the phase, clear effects (so a re-run does not re-fire). */
function stay(m: Machine): Machine {
return { phase: m.phase, ctx: m.ctx, effects: [] };
}
/** A command-transition: same as `to` but bumps the epoch (I1). Any outcome
* event issued under the old epoch is dropped once this lands. */
function command(
m: Machine,
phase: Phase,
effects: Effect[],
ctx?: Partial<Ctx>,
): Machine {
return {
phase,
ctx: { ...m.ctx, ...(ctx ?? {}), epoch: m.ctx.epoch + 1 },
effects,
};
}
// ---------------------------------------------------------------------------
// The pure reducer.
// ---------------------------------------------------------------------------
/** The terminal stream-finish events (one turn's stream ended). */
function isFinishEvent(event: Event): boolean {
return (
event.type === "FINISH_ABORT" ||
event.type === "FINISH_CLEAN" ||
event.type === "FINISH_DISCONNECT" ||
event.type === "FINISH_ERROR" ||
event.type === "STREAM_INCOMPLETE"
);
}
export function reduce(m: Machine, event: Event): Machine {
// MEDIUM (#488 re-review): honor ANY stream finish in `stopping` regardless of
// generation. A plain user Stop has NO successor stream — the aborted stream's
// finish IS the expected end of the stop, so exit `stopping -> idle` by that DATA
// (I4). The epoch filter below must NOT drop it: STOP_REQUESTED bumped the epoch,
// but the finish carries the PRE-stop generation (the runtime stamps it with the
// stream's start epoch), so I1 would otherwise strand the machine in `stopping`
// forever (no idle-cap covers `stopping`). The epoch filter stays in force for
// `superseding` (a successor B owns) — that is the F1 supersede drop.
if (m.phase.name === "stopping" && isFinishEvent(event)) {
return to(m, { name: "idle" }, {
// Reset ownership to local on this terminal transition (review #2): otherwise
// an observer-stop leaves ownership 'observer' and hides "Send now" forever.
ctx: { runFact: null, liveFollow: false, ownership: "local" },
effects: [{ type: "disarmPoll" }, { type: "cancelReconnect" }],
});
}
// I1: drop a stale outcome (an event issued under a superseded epoch).
if ("epoch" in event && event.epoch !== undefined && event.epoch !== m.ctx.epoch) {
return stay(m);
}
switch (event.type) {
// ---- local turn ----------------------------------------------------
case "SEND_LOCAL":
// A local send owns the view: leave any recovery, become the local
// streamer, disarm poll/reconnect. epoch++ so a late recovery outcome
// from the previous phase is dropped.
return command(
m,
{ name: "sending" },
[{ type: "cancelReconnect" }, { type: "disarmPoll" }],
{ ownership: "local", liveFollow: false },
);
case "STREAM_INCOMPLETE":
// An OBSERVER's attached stream ended incomplete (starved / torn) — follow
// the run to terminal via the degraded poll.
return to(m, { name: "polling", reason: event.reason }, {
effects: [{ type: "armPoll", reason: event.reason }],
});
case "STREAM_START": {
// First frame arrived. Adopt the run-fact runId if present. sending ->
// streaming; a reconnect/attach that just went live also lands here.
const runFact = event.runId ? { runId: event.runId } : m.ctx.runFact;
return to(m, { name: "streaming" }, {
ctx: { runFact },
effects: [{ type: "cancelReconnect" }, { type: "disarmPoll" }],
});
}
case "FINISH_CLEAN":
// A clean terminal outcome. The run is done — clear the run-fact and go
// idle. (The queue flush is a component concern gated by ownership; the
// FSM only models the phase.) Review #2: reset ownership to local so a
// just-finished observer-attach turn re-exposes "Send now" for the queue.
return to(m, { name: "idle" }, {
ctx: { runFact: null, liveFollow: false, ownership: "local" },
effects: [{ type: "disarmPoll" }, { type: "cancelReconnect" }],
});
case "FINISH_ABORT":
// A user Stop / intentional abort finished. If we were stopping, the
// terminal data has now arrived (I4) — go idle. The run-fact is cleared.
return to(m, { name: "idle" }, {
ctx: { runFact: null, liveFollow: false, ownership: "local" },
effects: [{ type: "disarmPoll" }, { type: "cancelReconnect" }],
});
case "FINISH_DISCONNECT":
// A LIVE SSE drop. Recovery depends on WHO we are (I2 + liveFollow):
// - a mount-attach OBSERVER (a one-shot resume, NOT live-follow) that drops
// -> the degraded poll drives the row to terminal from the DB.
if (m.ctx.ownership === "observer" && !m.ctx.liveFollow) {
return to(m, { name: "polling", reason: "disconnect-visible" }, {
effects: [{ type: "armPoll", reason: "disconnect-visible" }],
});
}
// - a LOCAL live turn (first drop) OR a live-follow re-attach (a SUBSEQUENT
// drop) -> (re-)enter the reconnect ladder. #488 commit 3: allowed
// REPEATEDLY — `liveFollow` is kept across a successful re-attach, so the
// second break reconnects again instead of falling to silent poll.
// #488 commit 2: gated on the RUN-FACT (or an existing live-follow), NOT on
// the presence of an assistant message — a setup-phase break still recovers.
// - visible content already on screen -> keep it, ALSO poll to terminal
// (a full replay could clobber the fuller live tail);
// - no visible content -> the reconnect ladder rebuilds it.
if (m.ctx.runFact || m.ctx.liveFollow) {
const effects: Effect[] = [
{ type: "scheduleReconnect", attempt: 1, delayMs: reconnectDelayMs(1) },
];
if (event.hasVisibleContent) effects.push({ type: "armPoll", reason: "disconnect-visible" });
return command(m, { name: "reconnecting", attempt: 1, failed: false }, effects, {
ownership: "observer",
liveFollow: true,
});
}
// No run to recover: a plain disconnect. Surface the terminal notice.
return to(m, { name: "idle" }, {
ctx: { runFact: null, liveFollow: false, ownership: "local" },
});
case "FINISH_ERROR":
return to(m, { name: "error", kind: event.kind }, {
ctx: { runFact: null, liveFollow: false, ownership: "local" },
effects: [{ type: "disarmPoll" }, { type: "cancelReconnect" }],
});
// ---- mount attach (resume) ----------------------------------------
case "ATTACH_START":
// A reopened tab attaches to a still-running run: observer ownership.
// #488 F2: ONLY from idle. The mount `getRun` round-trip resolves async, and
// a local send may have started meanwhile (phase `sending`, ownership local);
// a late ATTACH_START must NOT hijack that local turn into an observer-attach
// (queue would stop flushing, "Send now" would hide). Guarding in the reducer
// covers every dispatch source.
if (m.phase.name !== "idle") return stay(m);
return command(m, { name: "attaching" }, [{ type: "resumeStream" }], {
ownership: "observer",
runFact: event.runId ? { runId: event.runId } : m.ctx.runFact,
});
case "ATTACH_LIVE":
// The attach GET returned a live 2xx stream — follow it as an observer.
// Review #1: guard by SOURCE phase. The epoch filter alone is not enough — a
// POLL_TERMINAL uses to() (no epoch bump) and does not abort the in-flight
// GET, so a slow 2xx landing after the machine already left `attaching` (e.g.
// the armed poll saw the terminal row -> idle) would resurrect a settled run
// into a phantom `streaming`. Only enter streaming FROM `attaching`.
if (m.phase.name !== "attaching") return stay(m);
return to(m, { name: "streaming" });
case "ATTACH_NONE":
// 204 / non-2xx / throw: nothing live to attach. Arm the degraded poll to
// follow the run to terminal from the DB. This is a soft-negative run-fact
// (204 on a non-stripped path is authoritative-negative; the runtime may
// pass a RUN_FACT null separately). Keep the run-fact as-is here.
// Review #1: guard by source phase for consistency (a late outcome after the
// machine already left `attaching` must not re-arm a poll).
if (m.phase.name !== "attaching") return stay(m);
return to(m, { name: "polling", reason: "attach-none" }, {
effects: [{ type: "armPoll", reason: "attach-none" }],
});
// ---- reconnect after a live disconnect ----------------------------
case "RECONNECT_ATTEMPT":
// A scheduled backoff fired — fire the attach GET. epoch++ so the previous
// attempt's late outcome cannot drive this one.
if (m.phase.name !== "reconnecting") return stay(m);
return command(
m,
{ name: "reconnecting", attempt: event.attempt, failed: false },
[{ type: "resumeStream" }],
);
case "RECONNECT_ATTACHED":
// #488 commit 3: a live re-attach succeeded. Reset to streaming — the
// attempt counter is dropped, so a LATER disconnect can start a fresh
// ladder from attempt 1 (the old one-shot `!wasResumed` gate forbade a
// second cycle, sending the second break to silent poll).
// Review #1: guard by SOURCE phase. The armed degraded poll can reach the
// terminal row (POLL_TERMINAL -> idle, via to(), NO epoch bump, GET not
// aborted) BEFORE a slow reconnect GET returns 2xx; without this guard that
// late RECONNECT_ATTACHED (same epoch) would resurrect a settled run into a
// phantom `streaming`. Only re-enter streaming FROM `reconnecting`.
if (m.phase.name !== "reconnecting") return stay(m);
return to(m, { name: "streaming" }, {
effects: [{ type: "cancelReconnect" }, { type: "disarmPoll" }],
});
case "RECONNECT_NONE": {
// 204 / error during a reconnect attempt. Arm the degraded poll as the
// belt-and-suspenders fallback, then either back off to the next attempt
// or, at the cap, surface the manual Retry ("failed").
if (m.phase.name !== "reconnecting") return stay(m);
const attempt = m.phase.attempt;
if (attempt < RECONNECT_MAX_ATTEMPTS) {
return command(
m,
{ name: "reconnecting", attempt: attempt + 1, failed: false },
[
{ type: "armPoll", reason: "attach-none" },
{ type: "scheduleReconnect", attempt: attempt + 1, delayMs: reconnectDelayMs(attempt + 1) },
],
);
}
return to(m, { name: "reconnecting", attempt, failed: true }, {
effects: [{ type: "armPoll", reason: "reconnect-exhausted" }],
});
}
case "RETRY":
// Manual Retry from the "failed" reconnect banner OR the stalled banner.
if (m.phase.name === "reconnecting" && m.phase.failed) {
return command(
m,
{ name: "reconnecting", attempt: 1, failed: false },
[{ type: "resumeStream" }],
);
}
if (m.phase.name === "stalled") {
// Re-arm the poll to try to catch the run up again.
return command(m, { name: "polling", reason: "attach-none" }, [
{ type: "armPoll", reason: "attach-none" },
]);
}
return stay(m);
// ---- degraded poll -------------------------------------------------
case "POLL_TERMINAL":
// The run reached a terminal row via the poll (or the reconcile merge). Go
// idle and disarm everything (I4: this is a DATA-driven exit, incl. exit
// from `stopping`). Review #2: reset ownership to local.
return to(m, { name: "idle" }, {
ctx: { runFact: null, liveFollow: false, ownership: "local" },
effects: [{ type: "disarmPoll" }, { type: "cancelReconnect" }],
});
case "POLL_IDLE_CAP":
// Review #4: `stopping` also arms the poll (STOP_REQUESTED) but has NO other
// backstop — an observer-stop with no SDK stream to fire onFinish, whose
// server stop never drives the run terminal, would poll the DB forever. Give
// it a bounded exit: cap -> idle + disarm (NOT `stalled`; Stop was already
// pressed, so there is nothing for the user to retry).
if (m.phase.name === "stopping") {
return to(m, { name: "idle" }, {
ctx: { runFact: null, liveFollow: false, ownership: "local" },
effects: [{ type: "disarmPoll" }, { type: "cancelReconnect" }],
});
}
// #488 commit 4a: the poll hit the inactivity cap. Instead of going SILENT
// (the old "forever half-done answer"), surface a stalled banner + Retry.
if (m.phase.name !== "polling" && m.phase.name !== "reconnecting") return stay(m);
return to(m, { name: "stalled" }, {
effects: [{ type: "disarmPoll" }, { type: "cancelReconnect" }],
});
// ---- run-fact ------------------------------------------------------
case "RUN_FACT": {
const runFact = event.runFact;
// A fresh NEGATIVE fact (no active run) cancels recovery immediately (I3):
// there is nothing to reconnect to / poll for.
if (!runFact) {
if (
m.phase.name === "reconnecting" ||
m.phase.name === "attaching" ||
m.phase.name === "polling" ||
m.phase.name === "stopping"
) {
return to(m, { name: "idle" }, {
// Review #2: reset ownership to local on this terminal transition.
ctx: { runFact: null, liveFollow: false, ownership: "local" },
effects: [{ type: "cancelReconnect" }, { type: "disarmPoll" }],
});
}
return to(m, m.phase, { ctx: { runFact: null } });
}
// A positive fact just updates the context (pessimism toward an attempt: a
// stale-but-positive fact permits entering recovery; a 204 will cut it).
return to(m, m.phase, { ctx: { runFact } });
}
// ---- stop ----------------------------------------------------------
case "STOP_REQUESTED":
// Authoritative stop of a detached run. Enter `stopping` and fire stopRun +
// abort the local/attach reader. ALSO arm the poll so the terminal row is
// observed — the exit is by DATA (I4: a terminal row / negative run-fact),
// never by the stopRun HTTP response (which returns after abort, before
// finalization). For a local turn the aborted stream's onFinish (ANY finish)
// is HONORED in `stopping` at the top of reduce() — regardless of generation
// — and exits to idle; the armed poll is the fallback for an observer stop
// with no local onFinish.
return command(
m,
{ name: "stopping" },
[
{ type: "stopRun" },
{ type: "abortAttach" },
{ type: "cancelReconnect" },
{ type: "armPoll", reason: "attach-none" },
],
);
// ---- supersede (CAS) ----------------------------------------------
case "SUPERSEDE_REQUESTED":
// "Interrupt and send now": CAS POST /stream { supersede }. epoch++ so a
// late outcome of the interrupted run is dropped.
return command(
m,
{ name: "superseding" },
[{ type: "supersede", targetRunId: event.targetRunId }, { type: "cancelReconnect" }, { type: "disarmPoll" }],
);
case "SUPERSEDE_READY": {
// CAS succeeded (old run stopped/settled, slot taken, new run begun). We
// are now the local streamer of the NEW run. Adopt its runId if provided.
const runFact = event.runId ? { runId: event.runId } : m.ctx.runFact;
return to(m, { name: "streaming" }, {
ctx: { ownership: "local", runFact, liveFollow: false },
});
}
case "SUPERSEDE_MISMATCH":
// The active run moved between the click and the CAS. Per the spec: verify
// via /run rather than blindly banner — the mismatch may be our own already-
// superseded run. Surface a classified error AND fire a run-fact verify.
return to(m, { name: "error", kind: "supersede-mismatch" }, {
ctx: { runFact: event.currentRunId ? { runId: event.currentRunId } : m.ctx.runFact },
effects: [{ type: "postRun", reason: "verify" }],
});
case "SUPERSEDE_TIMEOUT":
// The old run did not settle within W. Nothing persisted; the composer keeps
// its text. Classified error, NO auto-retry (the old client retry ladder is
// removed in #488 commit 5).
return to(m, { name: "error", kind: "supersede-timeout" });
case "SUPERSEDE_INVALID":
return to(m, { name: "error", kind: "supersede-invalid" });
case "RUN_ALREADY_ACTIVE":
// A plain POST hit the one-active-run gate. NO auto-retry — the composer
// offers "interrupt and send" (supersede) instead.
return to(m, { name: "error", kind: "run-already-active" });
// ---- lifecycle -----------------------------------------------------
case "DISPOSE":
// Unmount: abort in-flight controllers, drop timers, and bump the epoch so
// NO late callback can drive this (now dead) machine (I5).
return command(
m,
{ name: "idle" },
[
{ type: "abortAttach" },
{ type: "cancelReconnect" },
{ type: "disarmPoll" },
],
{ liveFollow: false },
);
default: {
// Exhaustiveness guard.
const _never: never = event;
void _never;
return stay(m);
}
}
}
@@ -3,7 +3,6 @@ import {
resolveAdoptedChatId,
newlyAddedChatIds,
extractServerChatId,
extractRunId,
} from "./adopt-chat-id";
describe("resolveAdoptedChatId", () => {
@@ -71,17 +70,3 @@ describe("extractServerChatId", () => {
expect(extractServerChatId(undefined)).toBeUndefined();
});
});
describe("extractRunId", () => {
it("reads a string runId from the start metadata", () => {
expect(extractRunId({ metadata: { runId: "run-1" } })).toBe("run-1");
});
it("returns undefined when runId is absent", () => {
expect(extractRunId({ metadata: { chatId: "c" } })).toBeUndefined();
expect(extractRunId({})).toBeUndefined();
expect(extractRunId(undefined)).toBeUndefined();
});
it("returns undefined for a non-string runId", () => {
expect(extractRunId({ metadata: { runId: 7 } })).toBeUndefined();
});
});
@@ -56,20 +56,6 @@ export function extractServerChatId(
return typeof m?.chatId === "string" ? m.chatId : undefined;
}
/**
* #488: read the authoritative RUN id off a streaming assistant message. The
* server attaches it as `message.metadata.runId` on the `start` part when a run
* wraps the turn (see server `chatStreamMetadata`, #184/#487). This is the live
* run-fact update the client FSM adopts (mirrors `extractServerChatId`). Returns
* it only when it is a string; undefined otherwise.
*/
export function extractRunId(
message: { metadata?: unknown } | undefined,
): string | undefined {
const m = message?.metadata as { runId?: string } | undefined;
return typeof m?.runId === "string" ? m.runId : undefined;
}
/**
* The deduped set of ids present in `afterIds` but not in `beforeIds`. A
* paginated/flatMapped list can repeat the same id, so dedupe: one genuinely-new
@@ -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";
@@ -42,51 +42,6 @@ describe("describeChatError", () => {
);
});
// #488 commit 5: the #487 concurrency-gate / supersede 409s. FULL real bodies:
// a ConflictException(object) whose response is serialized verbatim, carrying a
// `code` and statusCode 409. Each must classify to a human text, not raw JSON.
it("classifies A_RUN_ALREADY_ACTIVE (409) as already-answering, not raw JSON", () => {
const body =
'{"message":"A run is already active for this chat","code":"A_RUN_ALREADY_ACTIVE","statusCode":409}';
expect(describeChatError(body, t).title).toBe(
"The agent is already answering",
);
// Never leaks the raw code as the detail.
expect(describeChatError(body, t).detail).not.toContain("A_RUN_ALREADY_ACTIVE");
});
it("classifies SUPERSEDE_TARGET_MISMATCH (409) as run-changed", () => {
const body =
'{"message":"active run does not match the supersede target","code":"SUPERSEDE_TARGET_MISMATCH","runId":"run-x","statusCode":409}';
expect(describeChatError(body, t).title).toBe(
"Couldn't interrupt — the run changed",
);
});
it("classifies SUPERSEDE_TIMEOUT (409) as couldn't-interrupt-in-time", () => {
const body =
'{"message":"the run did not settle within the supersede window","code":"SUPERSEDE_TIMEOUT","statusCode":409}';
expect(describeChatError(body, t).title).toBe("Couldn't interrupt in time");
});
it("classifies SUPERSEDE_INVALID (409) as couldn't-interrupt-that-run", () => {
const body =
'{"message":"supervise requires chatId","code":"SUPERSEDE_INVALID","statusCode":409}';
expect(describeChatError(body, t).title).toBe(
"Couldn't interrupt that run",
);
});
it("ORDER GUARD: A_RUN_ALREADY_ACTIVE wins over any generic status branch", () => {
// Even though the body could superficially look 4xx-ish, the code branch runs
// first, so it is never mislabeled by a generic status heading.
const body =
'{"message":"conflict","code":"A_RUN_ALREADY_ACTIVE","statusCode":409}';
const view = describeChatError(body, t);
expect(view.title).not.toBe("Something went wrong");
expect(view.title).not.toBe("AI provider not configured");
});
it("classifies a dropped connection (ECONNRESET) as a lost-connection error", () => {
expect(
describeChatError("Cannot connect to API: read ECONNRESET", t).title,
@@ -39,44 +39,6 @@ export function describeChatError(
};
}
// #488 commit 5: the #487 concurrency-gate / supersede 409s. These arrive as a
// ConflictException(object) body carrying a `code` (and statusCode 409). They
// MUST be classified by `code` STRICTLY BEFORE any generic status branch, or the
// user sees the raw JSON `{"code":"A_RUN_ALREADY_ACTIVE",…}`. The code strings
// are the real #487 server contract (ai-chat.controller.ts) — do not invent.
if (/"code"\s*:\s*"A_RUN_ALREADY_ACTIVE"/.test(msg)) {
return {
title: t("The agent is already answering"),
detail: t(
"This chat already has a run in progress. Wait for it to finish, or interrupt it and send now.",
),
};
}
if (/"code"\s*:\s*"SUPERSEDE_TARGET_MISMATCH"/.test(msg)) {
return {
title: t("Couldn't interrupt — the run changed"),
detail: t(
"The run you tried to interrupt is no longer the active one. Check the latest answer and try again.",
),
};
}
if (/"code"\s*:\s*"SUPERSEDE_TIMEOUT"/.test(msg)) {
return {
title: t("Couldn't interrupt in time"),
detail: t(
"The previous run didn't stop in time. Nothing was sent — try sending again.",
),
};
}
if (/"code"\s*:\s*"SUPERSEDE_INVALID"/.test(msg)) {
return {
title: t("Couldn't interrupt that run"),
detail: t(
"The run to interrupt doesn't belong to this chat. Reload and try again.",
),
};
}
if (/"statusCode"\s*:\s*403\b/.test(msg)) {
return {
title: t("AI chat is disabled"),
+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"
}
@@ -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);
});
});
@@ -97,8 +97,14 @@ describe('AiChatService.stream run-lifecycle safety net (#184)', () => {
};
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 = {
@@ -181,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 = {
@@ -287,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;
@@ -314,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 () => {
@@ -369,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);
});
});
/**
@@ -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.
@@ -1440,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(),
@@ -1448,7 +1755,7 @@ 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' })),
@@ -1487,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 = {
@@ -1570,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' }]);
});
});
/**
+547 -104
View File
@@ -14,6 +14,7 @@ import {
convertToModelMessages,
stepCountIs,
type UIMessage,
type ModelMessage,
type LanguageModel,
} from 'ai';
import { AiService } from '../../integrations/ai/ai.service';
@@ -54,6 +55,12 @@ import {
type SelectionContext,
} from './tools/current-page.util';
import { roleModelOverride } from './roles/role-model-config';
import {
resolveReplayBudget,
isContextOverflowError,
trimHistoryForReplay,
REPLAY_AGGRESSIVE_FRACTION,
} from './history-budget';
import {
startSseHeartbeat,
stripStreamingHopByHopHeaders,
@@ -126,6 +133,15 @@ const STEP_LIMIT_NO_ANSWER_MARKER =
const OUTPUT_DEGENERATION_ERROR =
'Output degeneration detected (repeated token loop)';
// Prefix recorded on the assistant row when the provider rejected the turn for
// CONTEXT OVERFLOW (#490): the replayed history exceeded the model's window. The
// row is ALSO stamped `metadata.replayOverflow` so the NEXT turn's budgeter trims
// aggressively (the reactive recovery — the overflowing turn had no usage signal
// to trigger preventive trimming, so the classified 400 is what un-bricks it).
export const CONTEXT_OVERFLOW_ERROR_PREFIX =
'Диалог превысил контекстное окно модели; история будет агрессивно ' +
'сокращена на следующем ходу.';
/**
* Compute the step-budget warning text (#444), or '' when this step is outside
* the warning band. The warning fires on steps
@@ -881,6 +897,21 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
const freshPage = await this.pageRepo.findById(pageId);
// Page deleted during the turn (or somehow foreign) => don't write.
if (!freshPage || freshPage.workspaceId !== workspace.id) return;
// Fast-path (#490): if a snapshot already exists at THIS page version
// (same updated_at instant), its content is already current — skip the full
// Markdown export + upsert entirely. A turn that did NOT touch the open page
// (the common case) thus does no snapshot work. This mirrors the read-side
// fast path in detectPageChange (sameInstant): both trust that a page edit
// bumps updated_at. When the agent (or a human) DID edit the page this turn,
// updated_at advanced, so this does not match and we re-export as before.
const existing = await this.aiChatPageSnapshotRepo.findByChatPage(
chatId,
pageId,
workspace.id,
);
if (existing && sameInstant(existing.pageUpdatedAt, freshPage.updatedAt)) {
return;
}
const currentMd = await this.tools.exportPageMarkdown(
user,
sessionId,
@@ -920,10 +951,17 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
// supplied or the supplied one does not belong to this workspace.
let isNewChat = false;
let chatId = body.chatId;
// Persisted chat-level metadata bag (#490): read once here so the deferred-tool
// activation set can be seeded from the previous turn. Undefined for a new chat.
let chatMetadata: Record<string, unknown> | undefined;
if (chatId) {
const existing = await this.aiChatRepo.findById(chatId, workspace.id);
if (!existing) {
chatId = undefined;
} else {
chatMetadata = (existing.metadata ?? undefined) as
| Record<string, unknown>
| undefined;
}
}
// The open page the client sent is attacker-controllable — BOTH its id and
@@ -1042,7 +1080,58 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
const incoming = lastUserMessage(body.messages);
const incomingText = uiMessageText(incoming);
// Persist the user message before contacting the model.
// #489: sanitize client-supplied parts ON RECEIPT. The client only ever
// sends `sendMessage({ text })` (a single text part); there is no
// file/attachment path. Any other part — most dangerously a tool-part in
// `input-available` state — is untrusted data that, once persisted to
// `metadata.parts` verbatim, is REPLAYED through convertToModelMessages on
// every later turn. A malformed tool-part makes that conversion throw,
// 500-ing every future turn of the chat forever ("bricked"). Drop any
// non-whitelisted part with a warn.
const sanitizedParts = sanitizeUserParts(incoming?.parts, (type) =>
this.logger.warn(
`Dropping unsupported user message part '${type}' on chat ${chatId}`,
),
);
// #489: rebuild the conversation from persisted history (not the client
// payload) and CONVERT it to model messages BEFORE persisting the user row.
// Load the OLD history (WITHOUT the new row) and append the incoming turn in
// memory for the conversion. This makes the insert happen only after a
// successful conversion, so a conversion failure cannot leave a DUPLICATE
// user row behind on the client's retry (the "bricked chat" that accreted a
// dup on every 500). `findAllByChat` returns chronological order (oldest ->
// newest) and keeps a 5000-row memory-safety backstop (on overflow it keeps
// the NEWEST rows and logs a warning); that is a safety net far above any
// realistic chat, not a conversational limit.
const oldHistory = await this.aiChatMessageRepo.findAllByChat(
chatId,
workspace.id,
);
const uiMessages: Array<Omit<UIMessage, 'id'> & { id: string }> = [
...oldHistory.map(rowToUiMessage),
{
id: 'pending-user',
role: 'user',
parts: (sanitizedParts && sanitizedParts.length > 0
? sanitizedParts
: textPart(incomingText)) as UIMessage['parts'],
},
];
// convertToModelMessages is async in ai@6.0.134 (returns Promise<ModelMessage[]>).
// Resilient (#489): a single poisoned row in the OLD history is isolated via
// per-row conversion and degraded to plain text with a "[tool context
// omitted]" marker rather than 500-ing the whole turn (silent loss of tool
// context is not acceptable — the model must see the truncation).
let messages = await convertHistoryResilient(uiMessages, (index, err) =>
this.logger.warn(
`Degraded unconvertible history row ${index} on chat ${chatId} to text: ${
err instanceof Error ? err.message : 'unknown error'
}`,
),
);
// Persist the user message only AFTER a successful conversion (#489).
await this.aiChatMessageRepo.insert({
chatId,
workspaceId: workspace.id,
@@ -1050,31 +1139,21 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
role: 'user',
content: incomingText,
// jsonb column: UIMessage parts are JSON-serializable at runtime but not
// structurally `JsonValue`, so cast through unknown.
metadata: (incoming?.parts ? { parts: incoming.parts } : null) as never,
// structurally `JsonValue`, so cast through unknown. Persist the SANITIZED
// parts (never the raw client parts) so the row is always convertible.
metadata: (sanitizedParts ? { parts: sanitizedParts } : null) as never,
});
// Rebuild the conversation from persisted history (not the client payload),
// so the model always sees the authoritative server-side transcript. Load
// the FULL history in chronological order (oldest -> newest, incl. the user
// message just inserted above) so NO turns are dropped — there is no
// recent-tail window anymore. `findAllByChat` keeps a 5000-row memory-safety
// backstop (on overflow it keeps the NEWEST rows and logs a warning); that
// is a safety net far above any realistic chat, not a conversational limit.
const history = await this.aiChatMessageRepo.findAllByChat(
chatId,
workspace.id,
);
const uiMessages = history.map(rowToUiMessage);
// convertToModelMessages is async in ai@6.0.134 (returns Promise<ModelMessage[]>).
const messages = await convertToModelMessages(uiMessages);
// Interrupt-resume detection (#198): the client "send now" flag is only a
// hint — confirm it against the persisted history (the preceding assistant
// turn must really be aborted/streaming) so a spoofed flag cannot inject the
// interrupt note onto an ordinary turn. The partial output the model needs is
// already in `messages` (the aborted assistant row replays via findRecent).
const interrupted = isInterruptResume(history, body.interrupted);
// Append the new user turn (shape-only) so index -2 is the prior assistant.
const interrupted = isInterruptResume(
[...oldHistory, { role: 'user', status: null, metadata: null }],
body.interrupted,
);
// Per-turn page-change detection (#274): if the open page was hand-edited by
// the user since the agent's last turn ended, compute the unified diff so the
@@ -1093,6 +1172,58 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
// Here we only need the admin-configured system prompt.
const resolved = await this.aiSettings.resolve(workspace.id);
// History-replay token budget (#490). The full conversation is replayed to
// the provider every turn, so a long chat eventually 400s on the context
// window — forever. Bound the REPLAYED history (never the persisted rows).
// PRIMARY signal is the provider's own fact: the last turn's contextTokens.
const replayBudget = resolveReplayBudget(resolved?.chatContextWindowRaw);
if (replayBudget.usedDefault) {
// The default fires precisely for installs with NO configured window —
// the ones that hit terminal overflow. Warn so it is observable.
this.logger.warn(
`AI chat (chat ${chatId}): no chatContextWindow configured; ` +
`applying the default replay budget (${replayBudget.thresholdTokens} tokens).`,
);
}
// Last turn's provider-reported context size (authoritative when present).
const priorContextTokens = lastAssistantContextTokens(oldHistory);
// Reactive recovery (#490): if the LAST turn was rejected for context
// overflow (stamped by onError), trim AGGRESSIVELY this turn — the
// overflowing turn produced no usage signal, so a normal-threshold trim may
// not shrink enough to fit. This is what un-bricks a chat that just 400'd.
const priorOverflowed = lastAssistantReplayOverflow(oldHistory);
const effectiveThreshold =
priorOverflowed && replayBudget.thresholdTokens != null
? Math.floor(
replayBudget.thresholdTokens * REPLAY_AGGRESSIVE_FRACTION,
)
: replayBudget.thresholdTokens;
if (priorOverflowed) {
this.logger.warn(
`AI chat (chat ${chatId}): previous turn hit context overflow; ` +
`applying aggressive replay budget (${effectiveThreshold} tokens).`,
);
}
const preTrim = trimHistoryForReplay(
messages,
effectiveThreshold,
// A prior OVERFLOW means the provider count is stale/absent — force the
// char-estimate path by ignoring priorContextTokens on recovery.
priorOverflowed ? undefined : priorContextTokens,
);
messages = preTrim.messages;
// Observability (#490): record the budgeter's decision on the turn so the UI
// can surface "replay truncated at N tokens". Threaded into flushAssistant.
let replayTrimmedToTokens: number | undefined = preTrim.trimmed
? preTrim.estimatedTokens
: undefined;
if (preTrim.trimmed) {
this.logger.log(
`AI chat (chat ${chatId}): replay history trimmed to ~${preTrim.estimatedTokens} ` +
`tokens (budget ${replayBudget.thresholdTokens}).`,
);
}
// Build the external MCP toolset FIRST so the system prompt can carry each
// connected server's admin-authored guidance (#180). Merge in admin-
// configured external MCP tools (web search, etc.; §6.8). A down/slow
@@ -1284,10 +1415,19 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
// tools + ALL external MCP tools), computed from the ACTUAL toolset so an
// external tool is loadable by its namespaced name. loadTools rejects any
// name outside this set.
const activatedTools = new Set<string>();
const validDeferredNames = new Set<string>(
Object.keys(baseTools).filter((k) => !CORE_TOOL_SET.has(k)),
);
// #490: seed the activation set from the chat's PERSISTED set so the model
// does not re-run loadTools every turn to re-activate the same tools. Only
// when deferred loading is enabled, and ALWAYS intersected with the CURRENT
// valid deferred names — an allowlist/role change must never resurrect a tool
// that no longer exists (prepareAgentStep would get a phantom active name).
const activatedTools = new Set<string>(
deferredEnabled
? seedActivatedTools(chatMetadata, validDeferredNames)
: [],
);
// Add the loadTools meta-tool ONLY when the feature is enabled; when off the
// toolset and behavior are exactly as before.
const tools = deferredEnabled
@@ -1297,6 +1437,39 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
}
: baseTools;
// #490: persist the (deterministically ordered) activation set back onto the
// chat metadata at turn end, so the NEXT turn seeds from it. Once-guarded and
// skipped when nothing new was activated (the set equals its seed) so an
// ordinary turn adds no extra write. Preserves other metadata keys.
let activatedToolsPersisted = false;
const persistActivatedTools = async (): Promise<void> => {
if (!deferredEnabled || activatedToolsPersisted || !chatId) return;
activatedToolsPersisted = true;
const current = [...activatedTools].sort();
const seeded = seedActivatedTools(chatMetadata, validDeferredNames).sort();
if (current.length === 0 || current.join('') === seeded.join('')) {
return; // nothing new activated -> no write
}
try {
await this.aiChatRepo.update(
chatId,
{
metadata: {
...(chatMetadata ?? {}),
activatedTools: current,
},
} as never,
workspace.id,
);
} catch (err) {
this.logger.warn(
`Failed to persist activated tools (chat ${chatId}): ${
err instanceof Error ? err.message : 'unknown error'
}`,
);
}
};
// Accumulate the turn's streamed output so a provider error / disconnect can
// persist the PARTIAL answer the user already saw — the SDK's onError/onAbort
// callbacks don't hand us the in-progress text. `capturedSteps` holds finished
@@ -1305,6 +1478,11 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
const capturedSteps: StepLike[] = [];
let inProgressText = '';
// Per-turn step->parts memo (#490): shared across every flushAssistant call
// this turn so each finished step's (large) output is JSON-stringified ONCE,
// not re-stringified on every subsequent onStepFinish flush (was O(N²)).
const partsCache: StepPartsCache = new WeakMap();
// Token-degeneration guard (#444). When the final-step lockdown is OFF, a
// runaway repetition loop (the 255KB "loadTools." incident) is aborted via
// this internal controller, unioned with the run/socket signal below. The
@@ -1374,7 +1552,10 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
await this.aiChatMessageRepo.update(
assistantId,
workspace.id,
flushAssistant(capturedSteps, '', 'streaming', { pageChanged }),
flushAssistant(capturedSteps, '', 'streaming', {
pageChanged,
partsCache,
}),
{ onlyIfStreaming: true },
);
} catch (err) {
@@ -1607,6 +1788,8 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
// closure scope here). Omitted/0 = no limit.
maxContextTokens: resolved?.chatContextWindow,
pageChanged,
partsCache,
replayTrimmedToTokens,
}),
);
// #184/#487: the RUN is finalized ALWAYS (never gated on the message).
@@ -1631,6 +1814,8 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
// own edits are baked in — and this also SEEDS the snapshot on the first
// turn. Runs once across every terminal path (see snapshotTurnEnd).
await snapshotTurnEnd();
// #490: persist the deferred-tool activation set for the next turn.
await persistActivatedTools();
// Generate the chat title for a freshly created chat AFTER the stream's
// provider call has completed — NOT concurrently with it. The z.ai coding
@@ -1654,7 +1839,16 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
// object, so the actual provider cause is clearly logged. Reuse the
// shared formatter so provider error formatting stays unified.
const e = error as { stack?: string };
const errorText = describeProviderError(error, String(error));
// #490 reactive branch: classify a CONTEXT-OVERFLOW rejection (the
// replayed history exceeded the model window). The overflowing turn had
// no prior usage to trigger preventive trimming, so we record a clear,
// distinguishable cause AND stamp the row so the NEXT turn's budgeter
// trims aggressively — the reactive recovery that un-bricks the chat.
const overflow = isContextOverflowError(error);
const providerError = describeProviderError(error, String(error));
const errorText = overflow
? `${CONTEXT_OVERFLOW_ERROR_PREFIX} (${providerError})`
: providerError;
this.logger.error(`AI chat stream error: ${errorText}`, e?.stack);
// DIAGNOSTIC (Safari stream-drop investigation) — temporary: timing of
// an error-terminated stream.
@@ -1672,6 +1866,9 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
flushAssistant(capturedSteps, inProgressText, 'error', {
error: errorText,
pageChanged,
partsCache,
replayTrimmedToTokens,
replayOverflow: overflow || undefined,
}),
);
// #184: settle the RUN as failed, carrying the provider/transport cause.
@@ -1681,6 +1878,8 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
// committed before the error must be baked into the snapshot, or the
// next turn would mis-report it as a user edit.
await snapshotTurnEnd();
// #490: persist the deferred-tool activation set for the next turn.
await persistActivatedTools();
},
onAbort: async ({ steps }) => {
// #444: distinguish a degeneration abort (our internal controller) from
@@ -1695,6 +1894,7 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
flushAssistant(capturedSteps, truncated, 'error', {
error: OUTPUT_DEGENERATION_ERROR,
pageChanged,
partsCache,
}),
);
if (runId)
@@ -1705,6 +1905,8 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
);
await closeExternalClients();
await snapshotTurnEnd();
// #490: persist the deferred-tool activation set for the next turn.
await persistActivatedTools();
return;
}
const partialChars =
@@ -1729,6 +1931,7 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
await finalizeAssistant(
flushAssistant(capturedSteps, inProgressText, 'aborted', {
pageChanged,
partsCache,
}),
);
// #184: settle the RUN as aborted (an explicit user stop reached the
@@ -1739,6 +1942,8 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
// committed before the client disconnect / stop() must be baked into the
// snapshot, or the next turn would mis-report it as a user edit.
await snapshotTurnEnd();
// #490: persist the deferred-tool activation set for the next turn.
await persistActivatedTools();
},
});
@@ -2049,6 +2254,70 @@ export function chatStreamMetadata(
return undefined;
}
/**
* The provider-reported context size of the most recent assistant turn, read from
* its persisted `metadata.contextTokens` (#490 replay budgeter's PRIMARY signal —
* the provider's own fact, not an estimate). Returns undefined for a chat with no
* assistant turn yet, or one whose last turn recorded no usage (e.g. it errored),
* in which case the budgeter falls back to the char-estimate.
*/
export function lastAssistantContextTokens(
history: ReadonlyArray<AiChatMessage>,
): number | undefined {
for (let i = history.length - 1; i >= 0; i--) {
const row = history[i];
if (row.role !== 'assistant') continue;
const meta = (row.metadata ?? {}) as { contextTokens?: unknown };
const n = meta.contextTokens;
return typeof n === 'number' && Number.isFinite(n) && n > 0 ? n : undefined;
}
return undefined;
}
/**
* Seed the per-turn deferred-tool activation set from a chat's persisted metadata
* (#490), INTERSECTED with the current valid deferred names. Persisting the set
* across turns saves the model re-running loadTools every turn to re-activate the
* same tools; intersecting on load means a changed allowlist / role can never
* resurrect a tool that no longer exists (which would hand prepareAgentStep a
* phantom active name). Tolerant of any stored shape — a non-array is ignored.
*/
export function seedActivatedTools(
metadata: Record<string, unknown> | undefined,
validDeferredNames: ReadonlySet<string>,
): string[] {
const stored = metadata?.activatedTools;
if (!Array.isArray(stored)) return [];
const seen = new Set<string>();
const out: string[] = [];
for (const name of stored) {
if (typeof name === 'string' && validDeferredNames.has(name) && !seen.has(name)) {
seen.add(name);
out.push(name);
}
}
return out;
}
/**
* Whether the most recent assistant turn was rejected for CONTEXT OVERFLOW
* (#490): its row carries `metadata.replayOverflow` (stamped by the stream's
* onError). The next turn's budgeter reads this to trim aggressively — the
* reactive recovery. Only the LAST assistant turn matters (an older overflow was
* already recovered), so we stop at the first assistant row scanning backwards.
*/
export function lastAssistantReplayOverflow(
history: ReadonlyArray<AiChatMessage>,
): boolean {
for (let i = history.length - 1; i >= 0; i--) {
const row = history[i];
if (row.role !== 'assistant') continue;
const meta = (row.metadata ?? {}) as { replayOverflow?: unknown };
return meta.replayOverflow === true;
}
return false;
}
/** The last message with role 'user' from a useChat payload, if any. */
function lastUserMessage(
messages: UIMessage[] | undefined,
@@ -2074,6 +2343,91 @@ function textPart(text: string): Array<{ type: 'text'; text: string }> {
return text ? [{ type: 'text', text }] : [];
}
/**
* Part types accepted on an INCOMING user turn (#489). The client only ever
* sends `sendMessage({ text })` (a single text part); there is no file/attachment
* path. Everything else on a client-supplied user message — most dangerously a
* tool-part in `input-available` state — is untrusted data that would be
* persisted to `metadata.parts` verbatim and replayed through
* `convertToModelMessages` on every later turn, potentially bricking the chat.
*/
const ALLOWED_USER_PART_TYPES: ReadonlySet<string> = new Set(['text']);
/**
* Keep only whitelisted parts on a client-supplied user message; report each
* dropped part's type via `onDrop` (the caller warns). Returns `undefined` when
* nothing survives (no parts / none whitelisted), so the caller persists a null
* metadata rather than an empty-parts object. Never throws.
*/
export function sanitizeUserParts(
parts: UIMessage['parts'] | undefined,
onDrop: (type: string) => void,
): UIMessage['parts'] | undefined {
if (!Array.isArray(parts)) return undefined;
const kept = parts.filter((p) => {
const type =
typeof (p as { type?: unknown })?.type === 'string'
? (p as { type: string }).type
: '';
if (ALLOWED_USER_PART_TYPES.has(type)) return true;
onDrop(type || '(unknown)');
return false;
});
return kept.length > 0 ? (kept as UIMessage['parts']) : undefined;
}
/** Marker for a history row whose tool parts could not be replayed (#489). */
export const TOOL_CONTEXT_OMITTED_MARKER = '[tool context omitted]';
/**
* Synthetic error text for a tool call that neither returned a result nor threw
* a `tool-error` — i.e. it was interrupted mid-step (an abort / server restart).
* Shared by `assistantParts` (the replayed `output-error` part) and
* `serializeSteps` (the `{ kind: 'interrupted' }` trace element) so the replay
* text and the trace stay in lockstep (#490).
*/
export const TOOL_CALL_INCOMPLETE_TEXT = 'Tool call did not complete.';
/**
* Convert persisted UI history to model messages, tolerating a single poisoned
* row (#489). `convertToModelMessages` over the WHOLE array throws if ANY row is
* malformed (e.g. a tool-part left unbalanced / in `input-available` state),
* which would otherwise 500 every turn of the chat forever. On a batch failure we
* fall back to per-row conversion so the bad row is isolated: it is degraded to
* plain text carrying its readable text plus a `[tool context omitted]` marker
* (the model MUST see that its tool context was truncated — silent loss is not
* acceptable), while every healthy row converts normally. Because AI SDK v6
* carries a tool call and its result inside the SAME assistant UIMessage's parts,
* per-row conversion preserves call/result pairing.
*/
export async function convertHistoryResilient(
uiMessages: Array<Omit<UIMessage, 'id'> & { id: string }>,
onDegrade: (index: number, err: unknown) => void,
): Promise<ModelMessage[]> {
try {
return await convertToModelMessages(uiMessages as UIMessage[]);
} catch {
const out: ModelMessage[] = [];
for (let i = 0; i < uiMessages.length; i++) {
const m = uiMessages[i];
try {
out.push(...(await convertToModelMessages([m as UIMessage])));
} catch (err) {
onDegrade(i, err);
const text = uiMessageText(m as UIMessage);
const degraded = text
? `${text}\n\n${TOOL_CONTEXT_OMITTED_MARKER}`
: TOOL_CONTEXT_OMITTED_MARKER;
out.push({
role: m.role === 'assistant' ? 'assistant' : 'user',
content: degraded,
} as ModelMessage);
}
}
return out;
}
}
/**
* Minimal shapes of the AI SDK v6 step objects we read to rebuild UIMessage
* parts (see ai@6.0.134 `StepResult`: `text`, `toolCalls` -> TypedToolCall,
@@ -2241,71 +2595,97 @@ function normalizeToolError(error: unknown): string {
*/
// Exported only so the unit tests can import these pure helpers; exporting
// them does not change runtime behavior.
/**
* Per-turn memo for {@link assistantParts}: a step's rebuilt parts keyed by the
* step OBJECT's identity (#490). A finished step in `capturedSteps` keeps a stable
* reference across every mid-stream flush, and `compactToolOutput` inside it does a
* `JSON.stringify` of the whole (often 50–200 KB) output — so without a memo each
* `onStepFinish` re-stringifies EVERY prior step's output (O(N²) stringify over a
* turn). Keyed by step identity => one stringify per step per turn. WeakMap so a
* turn's steps are GC'd with the turn.
*/
export type StepPartsCache = WeakMap<object, Array<Record<string, unknown>>>;
/** Build the parts for ONE step (text + a part per tool call). Pure. */
function buildStepParts(step: StepLike): Array<Record<string, unknown>> {
const parts: Array<Record<string, unknown>> = [];
if (step.text) {
parts.push({ type: 'text', text: step.text });
}
// Index this step's results by tool call id to pair calls with outputs.
const resultsById = new Map<string, unknown>();
for (const r of step.toolResults ?? []) {
if (r.toolCallId) resultsById.set(r.toolCallId, r.output);
}
// Index this step's THROWN tool failures (ai@6 `tool-error` content parts)
// by tool call id, so a call that failed replays with its real error text.
const errorsById = new Map<string, unknown>();
for (const part of step.content ?? []) {
if (part.type === 'tool-error' && part.toolCallId) {
errorsById.set(part.toolCallId, part.error);
}
}
for (const call of step.toolCalls ?? []) {
if (!call.toolName || !call.toolCallId) continue;
const hasResult = resultsById.has(call.toolCallId);
if (hasResult) {
// output-available: the tool returned; the next turn replays its result.
parts.push({
type: `tool-${call.toolName}`,
toolCallId: call.toolCallId,
state: 'output-available',
input: call.input,
output: compactToolOutput(resultsById.get(call.toolCallId)),
});
} else if (errorsById.has(call.toolCallId)) {
// The tool THREW: replay the REAL error so the model on the next turn
// knows WHY the call failed (and does not blindly repeat it). An
// output-error round-trips through convertToModelMessages as a balanced
// tool-call + tool-result, keeping the rebuilt history valid.
parts.push({
type: `tool-${call.toolName}`,
toolCallId: call.toolCallId,
state: 'output-error',
input: call.input,
errorText: normalizeToolError(errorsById.get(call.toolCallId)),
});
} else {
// No paired result AND no tool-error (e.g. aborted mid-step). Persisting
// a bare tool-call (input-available) would replay as an unpaired call and
// throw MissingToolResultsError on the next turn (convertToModelMessages
// emits no tool-result for it). Emit a SYNTHETIC paired result instead:
// an output-error round-trips through convertToModelMessages as a
// balanced tool-call + tool-result, keeping the rebuilt history valid.
parts.push({
type: `tool-${call.toolName}`,
toolCallId: call.toolCallId,
state: 'output-error',
input: call.input,
errorText: TOOL_CALL_INCOMPLETE_TEXT,
});
}
}
return parts;
}
export function assistantParts(
steps: ReadonlyArray<StepLike> | undefined,
fallbackText: string,
cache?: StepPartsCache,
): UIMessage['parts'] {
const parts: Array<Record<string, unknown>> = [];
let sawText = false;
for (const step of steps ?? []) {
if (step.text) {
parts.push({ type: 'text', text: step.text });
sawText = true;
}
// Index this step's results by tool call id to pair calls with outputs.
const resultsById = new Map<string, unknown>();
for (const r of step.toolResults ?? []) {
if (r.toolCallId) resultsById.set(r.toolCallId, r.output);
}
// Index this step's THROWN tool failures (ai@6 `tool-error` content parts)
// by tool call id, so a call that failed replays with its real error text.
const errorsById = new Map<string, unknown>();
for (const part of step.content ?? []) {
if (part.type === 'tool-error' && part.toolCallId) {
errorsById.set(part.toolCallId, part.error);
}
}
for (const call of step.toolCalls ?? []) {
if (!call.toolName || !call.toolCallId) continue;
const hasResult = resultsById.has(call.toolCallId);
if (hasResult) {
// output-available: the tool returned; the next turn replays its result.
parts.push({
type: `tool-${call.toolName}`,
toolCallId: call.toolCallId,
state: 'output-available',
input: call.input,
output: compactToolOutput(resultsById.get(call.toolCallId)),
});
} else if (errorsById.has(call.toolCallId)) {
// The tool THREW: replay the REAL error so the model on the next turn
// knows WHY the call failed (and does not blindly repeat it). An
// output-error round-trips through convertToModelMessages as a balanced
// tool-call + tool-result, keeping the rebuilt history valid.
parts.push({
type: `tool-${call.toolName}`,
toolCallId: call.toolCallId,
state: 'output-error',
input: call.input,
errorText: normalizeToolError(errorsById.get(call.toolCallId)),
});
} else {
// No paired result AND no tool-error (e.g. aborted mid-step). Persisting
// a bare tool-call (input-available) would replay as an unpaired call and
// throw MissingToolResultsError on the next turn (convertToModelMessages
// emits no tool-result for it). Emit a SYNTHETIC paired result instead:
// an output-error round-trips through convertToModelMessages as a
// balanced tool-call + tool-result, keeping the rebuilt history valid.
parts.push({
type: `tool-${call.toolName}`,
toolCallId: call.toolCallId,
state: 'output-error',
input: call.input,
errorText: 'Tool call did not complete.',
});
}
// Memoize per step object (#490): a finished step is immutable and keeps its
// reference across flushes, so its parts (and the costly output stringify) are
// built exactly once per turn. A cache miss (or no cache) just rebuilds.
let stepParts = cache?.get(step as object);
if (!stepParts) {
stepParts = buildStepParts(step);
cache?.set(step as object, stepParts);
}
parts.push(...stepParts);
}
const sawText = parts.some((p) => p.type === 'text');
if (!sawText && fallbackText) {
// No per-step text (e.g. a single final block): append the final text after
// any tool parts so the natural call -> result -> answer order is preserved.
@@ -2468,6 +2848,16 @@ export function flushAssistant(
maxContextTokens?: number;
error?: string;
pageChanged?: { title: string; diff: string } | null;
// Per-turn step->parts memo (#490): pass the SAME cache on every flush of a
// turn so each finished step's output is stringified once, not once per flush.
partsCache?: StepPartsCache;
// #490 observability: when the replay budgeter trimmed this turn's history,
// the (estimated) token size it trimmed to — the UI can show "replay truncated
// at N tokens". Omitted when nothing was trimmed.
replayTrimmedToTokens?: number;
// #490 reactive branch: set when the provider rejected this turn for context
// overflow. Stamped into metadata so the NEXT turn's budgeter trims aggressively.
replayOverflow?: boolean;
},
): AssistantFlush {
const finished = capturedSteps ?? [];
@@ -2477,13 +2867,20 @@ export function flushAssistant(
// in-progress step's text (the partial answer cut off by an error/abort, or
// simply not yet flushed mid-stream) as the last text part so the persisted
// parts match what streamed to the client.
const parts = assistantParts(finished, '') as unknown as Array<
Record<string, unknown>
>;
const parts = assistantParts(
finished,
'',
extra?.partsCache,
) as unknown as Array<Record<string, unknown>>;
if (trailing) parts.push({ type: 'text', text: trailing });
const metadata: Record<string, unknown> = {
parts: parts as unknown as UIMessage['parts'],
// Era marker for the `tool_calls` trace shape (#490): v2 stores outcome flags
// ({ ok } / { error, kind }) and NO tool output (the output lives once in
// `parts`). Old rows have no marker and the legacy { output } shape; a
// dual-shape query branches on this. Old rows are deliberately NOT migrated.
toolTraceVersion: 2,
};
// finishReason: prefer an explicit one; else derive a sensible value from the
// terminal status (so onError/onAbort records keep their historical reason).
@@ -2499,6 +2896,9 @@ export function flushAssistant(
if (extra?.contextTokens) metadata.contextTokens = extra.contextTokens;
if (extra?.maxContextTokens)
metadata.maxContextTokens = extra.maxContextTokens;
if (extra?.replayTrimmedToTokens)
metadata.replayTrimmedToTokens = extra.replayTrimmedToTokens;
if (extra?.replayOverflow) metadata.replayOverflow = true;
if (extra?.error) metadata.error = extra.error;
// Persist the page-change diff the agent saw this turn (#274 observability),
// so history / the Markdown export can show what the user changed. Only when
@@ -2524,42 +2924,85 @@ export function flushAssistant(
/**
* Reduce SDK step objects to a compact, JSON-serializable trace for the
* `tool_calls` column. Stores only what the UI action-log and history need —
* never raw provider payloads or keys.
* `tool_calls` column — trace format **v2** (#490).
*
* v2 stores, per call, ONLY the metadata a queryable trace needs — never the
* tool OUTPUT. Before #490 each output was persisted TWICE: once here (compacted)
* and once in `metadata.parts` (via `assistantParts`), so a 50-step run with
* 50–200 KB outputs wrote hundreds of MB per turn (each `onStepFinish` rewrote
* the whole row). The parts copy is the one the model replays and the UI/Markdown
* export render, so the trace copy of the output was pure duplication. v2 keeps
* the output ONLY in parts and reduces the trace to outcome flags.
*
* Element shapes (paired per call, in order):
* - `{ toolName, input }` — the call
* - `{ toolName, ok: true }` — it returned a result (success)
* - `{ toolName, error, kind: 'thrown' }` — it threw a `tool-error`
* - `{ toolName, error, kind: 'interrupted' }` — no result and no throw (an
* abort / server restart mid-step). `kind` is MANDATORY: without it a
* synthetic "Tool call did not complete." is indistinguishable from a real
* hard-fail and pollutes any error-rate scan. The distinction is STRUCTURAL
* (an `errorsById` hit vs the synthetic fallback branch), NOT a per-tool
* classifier — soft failures stay OUT of the trace (they live in
* `metadata.parts` outputs; a per-tool mirror would persist its own bugs).
*
* Rows carry `metadata.toolTraceVersion: 2` (set by {@link flushAssistant}) so a
* dual-shape query can branch on the era. Old rows are NOT migrated (rewriting
* giant jsonb is the very WAL churn this removes); see docs/reading-ai-logs.md.
*/
export function serializeSteps(
steps: ReadonlyArray<{
toolCalls?: ReadonlyArray<{ toolName?: string; input?: unknown }>;
toolResults?: ReadonlyArray<{ toolName?: string; output?: unknown }>;
toolCalls?: ReadonlyArray<{
toolCallId?: string;
toolName?: string;
input?: unknown;
}>;
toolResults?: ReadonlyArray<{ toolCallId?: string; toolName?: string }>;
content?: ReadonlyArray<{
type?: string;
toolCallId?: string;
toolName?: string;
error?: unknown;
}>;
}>,
): unknown {
const calls: Array<{
toolName?: string;
input?: unknown;
output?: unknown;
error?: string;
}> = [];
const calls: Array<
| { toolName?: string; input?: unknown }
| { toolName?: string; ok: true }
| { toolName?: string; error: string; kind: 'thrown' | 'interrupted' }
> = [];
for (const step of steps ?? []) {
// Index this step's results + thrown errors by tool call id, so each call is
// paired with its outcome (mirrors assistantParts' pairing exactly).
const resultIds = new Set<string>();
for (const r of step.toolResults ?? []) {
if (r.toolCallId) resultIds.add(r.toolCallId);
}
const errorsById = new Map<string, unknown>();
for (const part of step.content ?? []) {
if (part.type === 'tool-error' && part.toolCallId) {
errorsById.set(part.toolCallId, part.error);
}
}
for (const call of step.toolCalls ?? []) {
calls.push({ toolName: call.toolName, input: call.input });
}
for (const r of step.toolResults ?? []) {
calls.push({ toolName: r.toolName, output: compactToolOutput(r.output) });
}
// ai@6 surfaces a THROWN tool failure as a `tool-error` content part, NOT as
// a `toolResults` entry. Record it as its own paired element (mirroring how a
// successful result is appended) so the failure and its reason survive in the
// trace instead of leaving an orphaned call with no result.
for (const part of step.content ?? []) {
if (part.type === 'tool-error') {
if (call.toolCallId && resultIds.has(call.toolCallId)) {
// Success: the output itself lives in metadata.parts, not here.
calls.push({ toolName: call.toolName, ok: true });
} else if (call.toolCallId && errorsById.has(call.toolCallId)) {
// Hard fail: the tool threw. Persist the real (bounded) reason.
calls.push({
toolName: part.toolName,
error: normalizeToolError(part.error),
toolName: call.toolName,
error: normalizeToolError(errorsById.get(call.toolCallId)),
kind: 'thrown',
});
} else {
// Neither a result nor a throw: interrupted mid-step (abort/restart).
// Marked structurally so it never inflates a thrown-error count.
calls.push({
toolName: call.toolName,
error: TOOL_CALL_INCOMPLETE_TEXT,
kind: 'interrupted',
});
}
}
@@ -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,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();
}
+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
+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,
+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
+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();
@@ -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",
);
});
+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':