From 3f217f483573360c5c943e99241065b18c7429e8 Mon Sep 17 00:00:00 2001 From: agent_coder Date: Sat, 11 Jul 2026 10:39:42 +0300 Subject: [PATCH 1/6] =?UTF-8?q?perf(ai-chat):=20=D1=84=D0=BE=D1=80=D0=BC?= =?UTF-8?q?=D0=B0=D1=82=20=D1=82=D1=80=D0=B5=D0=B9=D1=81=D0=B0=20tool=5Fca?= =?UTF-8?q?lls=20v2=20=E2=80=94=20outputs=20=D1=82=D0=BE=D0=BB=D1=8C=D0=BA?= =?UTF-8?q?=D0=BE=20=D0=B2=20parts=20(#490)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Каждый 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) --- .../src/core/ai-chat/ai-chat.service.spec.ts | 102 ++++++- .../src/core/ai-chat/ai-chat.service.ts | 103 +++++-- docs/reading-ai-logs.md | 252 ++++++++++++------ 3 files changed, 343 insertions(+), 114 deletions(-) diff --git a/apps/server/src/core/ai-chat/ai-chat.service.spec.ts b/apps/server/src/core/ai-chat/ai-chat.service.spec.ts index 8dd32123..cd32331a 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.spec.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.spec.ts @@ -231,61 +231,137 @@ 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>; 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>; - // 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>; + 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>; 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>; + // 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); + }); }); describe('rowToUiMessage', () => { diff --git a/apps/server/src/core/ai-chat/ai-chat.service.ts b/apps/server/src/core/ai-chat/ai-chat.service.ts index 66d21473..3c2bd6ec 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.ts @@ -2152,6 +2152,15 @@ export function sanitizeUserParts( /** 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 @@ -2419,7 +2428,7 @@ export function assistantParts( toolCallId: call.toolCallId, state: 'output-error', input: call.input, - errorText: 'Tool call did not complete.', + errorText: TOOL_CALL_INCOMPLETE_TEXT, }); } } @@ -2602,6 +2611,11 @@ export function flushAssistant( const metadata: Record = { 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). @@ -2642,42 +2656,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(); + for (const r of step.toolResults ?? []) { + if (r.toolCallId) resultIds.add(r.toolCallId); + } + const errorsById = new Map(); + 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', }); } } diff --git a/docs/reading-ai-logs.md b/docs/reading-ai-logs.md index fbf2e211..135085dc 100644 --- a/docs/reading-ai-logs.md +++ b/docs/reading-ai-logs.md @@ -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: ""}` 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: -- 2.52.0 From 71517552b7ba56f4ea8b309a5fd9a514af4d1888 Mon Sep 17 00:00:00 2001 From: agent_coder Date: Sat, 11 Jul 2026 10:43:35 +0300 Subject: [PATCH 2/6] =?UTF-8?q?perf(ai-chat):=20=D0=BA=D1=8D=D1=88=20compa?= =?UTF-8?q?ctToolOutput=20=D0=BF=D0=BE=20identity=20=D1=88=D0=B0=D0=B3?= =?UTF-8?q?=D0=B0=20(#490)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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), flushAssistant пробрасывает его, stream() заводит один WeakMap на ход и передаёт во все flush'и. Промах кэша (или его отсутствие в тестах/легаси-вызовах) просто пересобирает — байтового расхождения нет. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../src/core/ai-chat/ai-chat.service.spec.ts | 49 ++++++ .../src/core/ai-chat/ai-chat.service.ts | 165 +++++++++++------- 2 files changed, 153 insertions(+), 61 deletions(-) diff --git a/apps/server/src/core/ai-chat/ai-chat.service.spec.ts b/apps/server/src/core/ai-chat/ai-chat.service.spec.ts index cd32331a..bee3139f 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.spec.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.spec.ts @@ -13,6 +13,7 @@ import { compactToolOutput, assistantParts, serializeSteps, + type StepPartsCache, rowToUiMessage, prepareAgentStep, stepBudgetWarning, @@ -114,6 +115,54 @@ describe('compactToolOutput', () => { describe('assistantParts', () => { type AnyPart = Record; + // #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 = [ { diff --git a/apps/server/src/core/ai-chat/ai-chat.service.ts b/apps/server/src/core/ai-chat/ai-chat.service.ts index 3c2bd6ec..4f8cb2dc 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.ts @@ -1347,6 +1347,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 @@ -1416,7 +1421,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) { @@ -1649,6 +1657,7 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy { // closure scope here). Omitted/0 = no limit. maxContextTokens: resolved?.chatContextWindow, pageChanged, + partsCache, }), ); // #184/#487: the RUN is finalized ALWAYS (never gated on the message). @@ -1714,6 +1723,7 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy { flushAssistant(capturedSteps, inProgressText, 'error', { error: errorText, pageChanged, + partsCache, }), ); // #184: settle the RUN as failed, carrying the provider/transport cause. @@ -1737,6 +1747,7 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy { flushAssistant(capturedSteps, truncated, 'error', { error: OUTPUT_DEGENERATION_ERROR, pageChanged, + partsCache, }), ); if (runId) @@ -1771,6 +1782,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 @@ -2368,71 +2380,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>>; + +/** Build the parts for ONE step (text + a part per tool call). Pure. */ +function buildStepParts(step: StepLike): Array> { + const parts: Array> = []; + 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(); + 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(); + 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 | undefined, fallbackText: string, + cache?: StepPartsCache, ): UIMessage['parts'] { const parts: Array> = []; - 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(); - 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(); - 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, - }); - } + // 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. @@ -2595,6 +2633,9 @@ 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; }, ): AssistantFlush { const finished = capturedSteps ?? []; @@ -2604,9 +2645,11 @@ 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 - >; + const parts = assistantParts( + finished, + '', + extra?.partsCache, + ) as unknown as Array>; if (trailing) parts.push({ type: 'text', text: trailing }); const metadata: Record = { -- 2.52.0 From c144abcb8368402557e472c5c914d2823da9c563 Mon Sep 17 00:00:00 2001 From: agent_coder Date: Sat, 11 Jul 2026 11:21:27 +0300 Subject: [PATCH 3/6] =?UTF-8?q?feat(ai-chat):=20=D1=82=D0=BE=D0=BA=D0=B5?= =?UTF-8?q?=D0=BD-=D0=B1=D1=8E=D0=B4=D0=B6=D0=B5=D1=82=20=D1=80=D0=B5?= =?UTF-8?q?=D0=BF=D0=BB=D0=B5=D1=8F=20=D0=B8=D1=81=D1=82=D0=BE=D1=80=D0=B8?= =?UTF-8?q?=D0=B8=20+=20=D1=80=D0=B5=D0=B0=D0=BA=D1=82=D0=B8=D0=B2=D0=BD?= =?UTF-8?q?=D0=B0=D1=8F=20=D0=B2=D0=B5=D1=82=D0=BA=D0=B0=20(#490)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Вся персистентная история реплеится провайдеру КАЖДЫЙ ход, поэтому длинный чат рано или поздно упирается в контекстное окно и получает провайдерский 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) --- .gitignore | 4 + apps/client/package.json | 1 + .../ai-chat/utils/count-stream-tokens.test.ts | 11 +- .../ai-chat/utils/count-stream-tokens.ts | 20 +- apps/server/package.json | 4 +- .../ai-chat/ai-chat.service.run-race.spec.ts | 51 ++- .../src/core/ai-chat/ai-chat.service.spec.ts | 68 ++++ .../src/core/ai-chat/ai-chat.service.ts | 132 ++++++- .../core/ai-chat/ai-chat.write-volume.spec.ts | 209 ++++++++++ .../src/core/ai-chat/history-budget.spec.ts | 266 +++++++++++++ .../server/src/core/ai-chat/history-budget.ts | 356 ++++++++++++++++++ .../integrations/ai/ai-settings.service.ts | 3 + apps/server/src/integrations/ai/ai.types.ts | 4 + packages/token-estimate/package.json | 19 + packages/token-estimate/src/index.test.ts | 31 ++ packages/token-estimate/src/index.ts | 35 ++ packages/token-estimate/tsconfig.json | 16 + pnpm-lock.yaml | 15 + 18 files changed, 1221 insertions(+), 24 deletions(-) create mode 100644 apps/server/src/core/ai-chat/ai-chat.write-volume.spec.ts create mode 100644 apps/server/src/core/ai-chat/history-budget.spec.ts create mode 100644 apps/server/src/core/ai-chat/history-budget.ts create mode 100644 packages/token-estimate/package.json create mode 100644 packages/token-estimate/src/index.test.ts create mode 100644 packages/token-estimate/src/index.ts create mode 100644 packages/token-estimate/tsconfig.json diff --git a/.gitignore b/.gitignore index 3eb7e75b..f3ce4f3e 100644 --- a/.gitignore +++ b/.gitignore @@ -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 diff --git a/apps/client/package.json b/apps/client/package.json index a9fb48f7..00180fd4 100644 --- a/apps/client/package.json +++ b/apps/client/package.json @@ -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", diff --git a/apps/client/src/features/ai-chat/utils/count-stream-tokens.test.ts b/apps/client/src/features/ai-chat/utils/count-stream-tokens.test.ts index 6b00fbc4..819c057e 100644 --- a/apps/client/src/features/ai-chat/utils/count-stream-tokens.test.ts +++ b/apps/client/src/features/ai-chat/utils/count-stream-tokens.test.ts @@ -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 }); }); diff --git a/apps/client/src/features/ai-chat/utils/count-stream-tokens.ts b/apps/client/src/features/ai-chat/utils/count-stream-tokens.ts index aaf99599..951bb992 100644 --- a/apps/client/src/features/ai-chat/utils/count-stream-tokens.ts +++ b/apps/client/src/features/ai-chat/utils/count-stream-tokens.ts @@ -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"; diff --git a/apps/server/package.json b/apps/server/package.json index a6791d80..7b97752a 100644 --- a/apps/server/package.json +++ b/apps/server/package.json @@ -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/(.*)$": "/database/$1", "^@docmost/transactional/(.*)$": "/integrations/transactional/$1", "^@docmost/ee/(.*)$": "/ee/$1", + "^@docmost/token-estimate$": "/../../../packages/token-estimate/src/index.ts", "^src/(.*)$": "/$1", "^@tiptap/react$": "/../test/stubs/tiptap-react.js" } diff --git a/apps/server/src/core/ai-chat/ai-chat.service.run-race.spec.ts b/apps/server/src/core/ai-chat/ai-chat.service.run-race.spec.ts index 8b170c11..b15c2ed1 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.run-race.spec.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.run-race.spec.ts @@ -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; + }; + 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; + }; + expect('replayOverflow' in patch.metadata).toBe(false); + }); }); /** diff --git a/apps/server/src/core/ai-chat/ai-chat.service.spec.ts b/apps/server/src/core/ai-chat/ai-chat.service.spec.ts index bee3139f..70293eb0 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.spec.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.spec.ts @@ -29,6 +29,8 @@ import { FINAL_STEP_NUDGE, STEP_LIMIT_NO_ANSWER_MARKER, OUTPUT_DEGENERATION_ERROR, + lastAssistantContextTokens, + lastAssistantReplayOverflow, } from './ai-chat.service'; import type { AiChatMessage, Workspace } from '@docmost/db/types/entity.types'; import { buildSystemPrompt } from './ai-chat.prompt'; @@ -413,6 +415,55 @@ describe('toolTraceVersion era marker (#490)', () => { }); }); +// #490 replay-budget signal helpers over persisted history. +describe('lastAssistantContextTokens', () => { + const row = ( + role: string, + metadata: Record | 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(); + }); +}); + +describe('lastAssistantReplayOverflow', () => { + const row = ( + role: string, + metadata: Record | 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', () => { it('prefers metadata.parts over content', () => { const row = { @@ -743,6 +794,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. diff --git a/apps/server/src/core/ai-chat/ai-chat.service.ts b/apps/server/src/core/ai-chat/ai-chat.service.ts index 4f8cb2dc..b4c65333 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.ts @@ -55,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, @@ -127,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 @@ -1086,7 +1101,7 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy { // 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). - const messages = await convertHistoryResilient(uiMessages, (index, err) => + 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' @@ -1135,6 +1150,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 @@ -1658,6 +1725,7 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy { maxContextTokens: resolved?.chatContextWindow, pageChanged, partsCache, + replayTrimmedToTokens, }), ); // #184/#487: the RUN is finalized ALWAYS (never gated on the message). @@ -1705,7 +1773,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. @@ -1724,6 +1801,8 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy { error: errorText, pageChanged, partsCache, + replayTrimmedToTokens, + replayOverflow: overflow || undefined, }), ); // #184: settle the RUN as failed, carrying the provider/transport cause. @@ -2103,6 +2182,45 @@ 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, +): 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; +} + +/** + * 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, +): 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, @@ -2636,6 +2754,13 @@ export function flushAssistant( // 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 ?? []; @@ -2674,6 +2799,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 diff --git a/apps/server/src/core/ai-chat/ai-chat.write-volume.spec.ts b/apps/server/src/core/ai-chat/ai-chat.write-volume.spec.ts new file mode 100644 index 00000000..cbd12ad2 --- /dev/null +++ b/apps/server/src/core/ai-chat/ai-chat.write-volume.spec.ts @@ -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[]): 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, +): Promise { + 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[] = []; + + // v1: each step re-persists ALL prior outputs into the trace (the O(N²) churn). + const v1 = await walDelta(c, async () => { + const acc: ReturnType[] = []; + 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[] = []; + 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[] = []; + 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[] = []; + 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); +}); diff --git a/apps/server/src/core/ai-chat/history-budget.spec.ts b/apps/server/src/core/ai-chat/history-budget.spec.ts new file mode 100644 index 00000000..16f6fd76 --- /dev/null +++ b/apps/server/src/core/ai-chat/history-budget.spec.ts @@ -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); + }); +}); diff --git a/apps/server/src/core/ai-chat/history-budget.ts b/apps/server/src/core/ai-chat/history-budget.ts new file mode 100644 index 00000000..fb45702f --- /dev/null +++ b/apps/server/src/core/ai-chat/history-budget.ts @@ -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; + 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; + 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, +): 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, + 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; +} diff --git a/apps/server/src/integrations/ai/ai-settings.service.ts b/apps/server/src/integrations/ai/ai-settings.service.ts index 6be0b5be..ad98c256 100644 --- a/apps/server/src/integrations/ai/ai-settings.service.ts +++ b/apps/server/src/integrations/ai/ai-settings.service.ts @@ -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 diff --git a/apps/server/src/integrations/ai/ai.types.ts b/apps/server/src/integrations/ai/ai.types.ts index 06bf83e3..c43d2dfc 100644 --- a/apps/server/src/integrations/ai/ai.types.ts +++ b/apps/server/src/integrations/ai/ai.types.ts @@ -105,6 +105,10 @@ export interface ResolvedAiConfig extends Partial { // 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 diff --git a/packages/token-estimate/package.json b/packages/token-estimate/package.json new file mode 100644 index 00000000..eef030a6 --- /dev/null +++ b/packages/token-estimate/package.json @@ -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" + } +} diff --git a/packages/token-estimate/src/index.test.ts b/packages/token-estimate/src/index.test.ts new file mode 100644 index 00000000..2854efd0 --- /dev/null +++ b/packages/token-estimate/src/index.test.ts @@ -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)); + }); +}); diff --git a/packages/token-estimate/src/index.ts b/packages/token-estimate/src/index.ts new file mode 100644 index 00000000..ec8d2c04 --- /dev/null +++ b/packages/token-estimate/src/index.ts @@ -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); +} diff --git a/packages/token-estimate/tsconfig.json b/packages/token-estimate/tsconfig.json new file mode 100644 index 00000000..dec9218d --- /dev/null +++ b/packages/token-estimate/tsconfig.json @@ -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"] +} diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index bdf58b2b..c876808e 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -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': -- 2.52.0 From c0e39a395c8de141b474d73e7ff558d575adea11 Mon Sep 17 00:00:00 2001 From: agent_coder Date: Sat, 11 Jul 2026 11:27:15 +0300 Subject: [PATCH 4/6] =?UTF-8?q?perf(ai-chat):=20deferred-=D0=B0=D0=BA?= =?UTF-8?q?=D1=82=D0=B8=D0=B2=D0=B0=D1=86=D0=B8=D1=8F=20=D1=82=D1=83=D0=BB?= =?UTF-8?q?=D0=BE=D0=B2=20=D0=B2=20metadata=20=D1=87=D0=B0=D1=82=D0=B0=20(?= =?UTF-8?q?#490)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Активированный set сбрасывался каждый ход → модель заново гоняла loadTools, чтобы переактивировать те же тулы (лишний round-trip на каждом ходу). Теперь набор персистится в metadata чата и сидируется на следующем ходу. - Миграция: jsonb-колонка metadata на ai_chats (default '{}'); db.d.ts дополнен вручную (AiChats.metadata: Generated). - seedActivatedTools(metadata, validDeferredNames): читает сохранённый набор, ПЕРЕСЕКАЯ с актуальными validDeferredNames — смена allowlist/ролей не воскресит несуществующий тул (иначе prepareAgentStep получил бы фантомное активное имя). Сид только при deferredEnabled. - Персист на завершении хода (once-guard, во всех терминальных ветках рядом со snapshotTurnEnd): детерминированно отсортированный набор, merge в существующий bag (другие ключи сохраняются), запись пропускается если ничего нового не активировано (обычный ход не даёт лишней записи). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../src/core/ai-chat/ai-chat.service.spec.ts | 40 +++++++++ .../src/core/ai-chat/ai-chat.service.ts | 84 ++++++++++++++++++- .../20260707T120000-ai-chat-metadata.ts | 24 ++++++ apps/server/src/database/types/db.d.ts | 3 + 4 files changed, 150 insertions(+), 1 deletion(-) create mode 100644 apps/server/src/database/migrations/20260707T120000-ai-chat-metadata.ts diff --git a/apps/server/src/core/ai-chat/ai-chat.service.spec.ts b/apps/server/src/core/ai-chat/ai-chat.service.spec.ts index 70293eb0..095bb34a 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.spec.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.spec.ts @@ -31,6 +31,7 @@ import { 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'; @@ -439,6 +440,45 @@ describe('lastAssistantContextTokens', () => { }); }); +// #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, diff --git a/apps/server/src/core/ai-chat/ai-chat.service.ts b/apps/server/src/core/ai-chat/ai-chat.service.ts index b4c65333..80078827 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.ts @@ -936,10 +936,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 | undefined; if (chatId) { const existing = await this.aiChatRepo.findById(chatId, workspace.id); if (!existing) { chatId = undefined; + } else { + chatMetadata = (existing.metadata ?? undefined) as + | Record + | undefined; } } // The open page the client sent is attacker-controllable — BOTH its id and @@ -1393,10 +1400,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(); const validDeferredNames = new Set( 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( + 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 @@ -1406,6 +1422,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 => { + 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 @@ -1750,6 +1799,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 @@ -1812,6 +1863,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 @@ -1837,6 +1890,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 = @@ -1872,6 +1927,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(); }, }); @@ -2202,6 +2259,31 @@ export function lastAssistantContextTokens( 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 | undefined, + validDeferredNames: ReadonlySet, +): string[] { + const stored = metadata?.activatedTools; + if (!Array.isArray(stored)) return []; + const seen = new Set(); + 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 diff --git a/apps/server/src/database/migrations/20260707T120000-ai-chat-metadata.ts b/apps/server/src/database/migrations/20260707T120000-ai-chat-metadata.ts new file mode 100644 index 00000000..429a6bd9 --- /dev/null +++ b/apps/server/src/database/migrations/20260707T120000-ai-chat-metadata.ts @@ -0,0 +1,24 @@ +import { type Kysely, sql } from 'kysely'; + +export async function up(db: Kysely): Promise { + // 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): Promise { + await db.schema.alterTable('ai_chats').dropColumn('metadata').execute(); +} diff --git a/apps/server/src/database/types/db.d.ts b/apps/server/src/database/types/db.d.ts index bf5a9663..a17d628c 100644 --- a/apps/server/src/database/types/db.d.ts +++ b/apps/server/src/database/types/db.d.ts @@ -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; createdAt: Generated; updatedAt: Generated; deletedAt: Timestamp | null; -- 2.52.0 From d35956b9e97fa8816c8c14d20d189d22cb504fee Mon Sep 17 00:00:00 2001 From: agent_coder Date: Sat, 11 Jul 2026 11:30:46 +0300 Subject: [PATCH 5/6] perf(ai-chat): snapshotOpenPage fast-path (#490) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- .../src/core/ai-chat/ai-chat.service.spec.ts | 74 +++++++++++++++++++ .../src/core/ai-chat/ai-chat.service.ts | 15 ++++ 2 files changed, 89 insertions(+) diff --git a/apps/server/src/core/ai-chat/ai-chat.service.spec.ts b/apps/server/src/core/ai-chat/ai-chat.service.spec.ts index 095bb34a..7c0e4273 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.spec.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.spec.ts @@ -440,6 +440,80 @@ describe('lastAssistantContextTokens', () => { }); }); +// #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 }) + .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 }) + .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 }) + .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']); diff --git a/apps/server/src/core/ai-chat/ai-chat.service.ts b/apps/server/src/core/ai-chat/ai-chat.service.ts index 80078827..8ab879ff 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.ts @@ -897,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, -- 2.52.0 From 9f0e880e0bac05bddf48efcfa80ecc2a0b20022a Mon Sep 17 00:00:00 2001 From: agent_coder Date: Sat, 11 Jul 2026 11:35:58 +0300 Subject: [PATCH 6/6] =?UTF-8?q?perf(mcp):=20checkNewComments=20=E2=80=94?= =?UTF-8?q?=20=D0=BF=D0=B0=D1=80=D0=B0=D0=BB=D0=BB=D0=B5=D0=BB=D0=B8=D0=B7?= =?UTF-8?q?=D0=BC=20=D1=81=20=D0=BA=D0=B0=D0=BF=D0=BE=D0=BC=20(#490)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- packages/mcp/src/client/comments.ts | 78 ++++++++++++++----- .../mcp/test/mock/pagination-cursor.test.mjs | 68 ++++++++++++++++ 2 files changed, 126 insertions(+), 20 deletions(-) diff --git a/packages/mcp/src/client/comments.ts b/packages/mcp/src/client/comments.ts index c1fd557d..df5368e3 100644 --- a/packages/mcp/src/client/comments.ts +++ b/packages/mcp/src/client/comments.ts @@ -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( + items: readonly T[], + limit: number, + fn: (item: T, index: number) => Promise, +): Promise { + const results = new Array(items.length); + let cursor = 0; + const worker = async (): Promise => { + 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>( 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, diff --git a/packages/mcp/test/mock/pagination-cursor.test.mjs b/packages/mcp/test/mock/pagination-cursor.test.mjs index fd6ab47f..efc997cd 100644 --- a/packages/mcp/test/mock/pagination-cursor.test.mjs +++ b/packages/mcp/test/mock/pagination-cursor.test.mjs @@ -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", + ); +}); -- 2.52.0