From 2f8c5d9a98b2d6607048a9274eda42ccb00fd39d Mon Sep 17 00:00:00 2001 From: agent_coder Date: Sun, 12 Jul 2026 04:35:06 +0300 Subject: [PATCH 1/3] =?UTF-8?q?perf(client):=20=D0=B8=D0=BD=D0=BA=D1=80?= =?UTF-8?q?=D0=B5=D0=BC=D0=B5=D0=BD=D1=82=D0=B0=D0=BB=D1=8C=D0=BD=D1=8B?= =?UTF-8?q?=D0=B9=20=D1=80=D0=B5=D0=BD=D0=B4=D0=B5=D1=80=20=D1=81=D1=82?= =?UTF-8?q?=D1=80=D0=B8=D0=BC=D0=B0=20=D0=BE=D1=82=D0=B2=D0=B5=D1=82=D0=B0?= =?UTF-8?q?=20(#492)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Путь ответа ассистента теперь рендерится инкрементально, по образцу StreamingPlainText из ветки reasoning (#492, волна C эпика #497). Раньше MarkdownPart прогонял ВЕСЬ накопленный ответ через канонический конвейер (markdownToProseMirrorSync → PMNode.fromJSON → DOMSerializer → DOMPurify) на КАЖДОМ throttled-тике (~20 Гц). На синтетическом потоке ~100 КБ это 394 вызова renderChatMarkdown — O(числа тиков), причём каждый вызов заново парсит всю растущую строку. Теперь: - StreamingMarkdownText делит текст на блоки по безопасному срезу (splitPlainChunks — та же append-only-инвариантность, что у reasoning): СТАБИЛИЗИРОВАННЫЕ блоки идут через канонический конвейер и мемоизируются (каждый блок парсится РОВНО ОДИН раз), живой ХВОСТ — дешёвый plain-text (React-escaped, без парсера/санитайзера/innerHTML) до стабилизации. - На финализации (флип state → done или конец хода) — ОДИН полный канонический рендер всего текста: побайтовая визуальная паритетность с прежним выводом (включая

  • -обёртки схемы и scoped-CSS из #498). - Гейт liveness тот же, что у ReasoningBlock: streaming = turnStreaming && part.state === "streaming". Также цикл рендера частей переведён на ИСЧЕРПЫВАЮЩИЙ switch по видам частей с never-проверкой в default (вместо прежнего WARNING-комментария): новый закрытый вид части в UIMessagePart теперь ошибка компиляции. Тесты: - perf-smoke: на ~100 КБ потоке число вызовов renderChatMarkdown ≤ blocks+2 и ≪ ticks (O(блоков), не O(тиков)); мутация (снять memo с MarkdownChunk) краснит ассерт (78631 вызовов). - visual-regression: финальный рендер побайтово равен renderChatMarkdown всего текста (в т.ч.

  • ), инкрементальный вид сходится к нему на финализации; учтён neutralizeInternalLinks. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../components/message-item-memo.test.tsx | 87 +++++++++++ .../components/message-item.render.test.tsx | 112 ++++++++++++++ .../ai-chat/components/message-item.tsx | 145 ++++++++++++------ .../components/streaming-markdown-text.tsx | 96 ++++++++++++ 4 files changed, 395 insertions(+), 45 deletions(-) create mode 100644 apps/client/src/features/ai-chat/components/message-item.render.test.tsx create mode 100644 apps/client/src/features/ai-chat/components/streaming-markdown-text.tsx diff --git a/apps/client/src/features/ai-chat/components/message-item-memo.test.tsx b/apps/client/src/features/ai-chat/components/message-item-memo.test.tsx index 61894ad7..98b88d27 100644 --- a/apps/client/src/features/ai-chat/components/message-item-memo.test.tsx +++ b/apps/client/src/features/ai-chat/components/message-item-memo.test.tsx @@ -27,6 +27,7 @@ vi.mock("@/features/ai-chat/utils/markdown.ts", async () => { import MessageItem from "./message-item"; import { messageSignature } from "@/features/ai-chat/utils/message-signature.ts"; +import { splitPlainChunks } from "./streaming-plain-text"; // matchMedia (read by MantineProvider) is stubbed globally in vitest.setup.ts. @@ -114,3 +115,89 @@ describe("MessageItem markdown memoization", () => { expect(queryByText("streamed answer")).not.toBeNull(); }); }); + +// PERF SMOKE (#492): the whole point of the incremental streaming render is that +// the ANSWER path costs O(number of markdown blocks), NOT O(number of throttled +// ~20Hz ticks). Pre-#492 the finalized MarkdownPart re-parsed the WHOLE growing +// answer on every delta — a synthetic ~100 KB stream measured 394 renderChatMarkdown +// calls (one per tick). With the incremental render each STABILIZED block is parsed +// exactly once (memoized in MarkdownChunk) and the live tail is cheap plain text, so +// the call count collapses to ~= the block count regardless of tick granularity. +describe("MessageItem streaming answer render is O(blocks), not O(ticks)", () => { + // ~100 KB answer. Each section is a heading + a paragraph — TWO blank-line + // delimited markdown blocks — so the safe-cut block count is ~2× the section + // count. The perf claim is about the BLOCK count (the memoization granularity), + // measured directly with splitPlainChunks below, not the section count. + const buildAnswer = () => { + const SECTIONS = 100; + const paragraphs: string[] = []; + for (let i = 0; i < SECTIONS; i++) { + paragraphs.push(`## Section ${i}\n\n` + "lorem ipsum dolor ".repeat(55)); + } + const full = paragraphs.join("\n\n"); + // The number of memoized markdown blocks the incremental render splits into + // (all but the live tail are parsed once each). + return { full, blocks: splitPlainChunks(full).length }; + }; + + const streamMsg = (text: string, state: "streaming" | "done"): UIMessage => + ({ + id: "m1", + role: "assistant", + parts: [{ type: "text", text, state }], + }) as UIMessage; + + it("parses each block ~once over a 100KB stream (≈blocks, ≪ ticks)", () => { + renderChatMarkdownSpy.mockClear(); + const { full, blocks } = buildAnswer(); + const CHUNK = 128; // a realistic ~20Hz throttled delta size + const ticks = Math.ceil(full.length / CHUNK); + + let msg = streamMsg(full.slice(0, CHUNK), "streaming"); + const { rerender } = render( + + + , + ); + for (let end = 2 * CHUNK; end < full.length; end += CHUNK) { + msg = streamMsg(full.slice(0, end), "streaming"); + rerender( + + + , + ); + } + // Finalize: the streaming→done flip renders the whole answer through ONE + // canonical pass (visual parity), so the finished DOM matches the pre-#492 + // output. This is the single extra parse on top of the per-block ones. + const done = streamMsg(full, "done"); + rerender( + + + , + ); + + const calls = renderChatMarkdownSpy.mock.calls.length; + // Sanity: the stream really had far more ticks than blocks (else the test is + // vacuous — the point is that calls scale with blocks, not ticks). + expect(ticks).toBeGreaterThan(blocks * 3); + // O(blocks): each stabilized block parsed once + the single final whole-text + // parse. A small constant absorbs the finalize render and the live-tail block; + // the load-bearing claim is the bound below. + expect(calls).toBeLessThanOrEqual(blocks + 2); + // ≪ ticks — and, non-vacuously, the blocks WERE parsed (not skipped entirely). + expect(calls).toBeLessThan(ticks / 3); + expect(calls).toBeGreaterThan(blocks / 2); + // MUTATION-VERIFY (documented, not run here): dropping the `memo()` wrapper on + // MarkdownChunk (so every stable block re-parses each tick) drives `calls` + // toward `ticks` (~394), reddening both upper-bound assertions above. + }); +}); diff --git a/apps/client/src/features/ai-chat/components/message-item.render.test.tsx b/apps/client/src/features/ai-chat/components/message-item.render.test.tsx new file mode 100644 index 00000000..9e4266b8 --- /dev/null +++ b/apps/client/src/features/ai-chat/components/message-item.render.test.tsx @@ -0,0 +1,112 @@ +import { describe, expect, it, vi } from "vitest"; +import { render } from "@testing-library/react"; +import { MantineProvider } from "@mantine/core"; +import type { UIMessage } from "@ai-sdk/react"; + +// Stub react-i18next (the component reads `useTranslation`). Mirrors the other +// message-item specs. +vi.mock("react-i18next", () => ({ + useTranslation: () => ({ t: (key: string) => key }), +})); + +import MessageItem from "./message-item"; +import { messageSignature } from "@/features/ai-chat/utils/message-signature.ts"; +// The REAL canonical renderer (NOT the spy the memo test installs): this file +// exercises the actual markdown output so the visual-regression assertions below +// compare against genuine HTML (incl. the schema's `

  • ` wrappers). +import { renderChatMarkdown } from "@/features/ai-chat/utils/markdown.ts"; +import classes from "./ai-chat.module.css"; + +const msg = ( + parts: UIMessage["parts"], + extra?: Partial, +): UIMessage => + ({ id: "m1", role: "assistant", parts, ...extra }) as UIMessage; + +const renderRow = (message: UIMessage, turnStreaming = false) => + render( + + + , + ); + +// A rich multi-block answer that exercises headings, a list (the `

  • ` case +// the scoped CSS tightens), inline emphasis, and multiple paragraphs. +const ANSWER = [ + "# Заголовок", + "", + "Первый абзац с **жирным** и `кодом`.", + "", + "- пункт один", + "- пункт два", + "", + "Второй абзац.", +].join("\n"); + +describe("MessageItem final render — visual parity with the canonical pipeline", () => { + it("a finalized text part renders exactly renderChatMarkdown(text)", () => { + const { container } = renderRow( + msg([{ type: "text", text: ANSWER, state: "done" }]), + ); + const block = container.querySelector(`.${classes.markdown}`); + expect(block).not.toBeNull(); + // Byte-for-byte the canonical output (the SAME whole-text pass the pre-#492 + // MarkdownPart produced), including `

  • …

    • ` wrappers. + expect(block!.innerHTML).toBe(renderChatMarkdown(ANSWER, {})); + // The list wrapper is really present (guards against a vacuous empty render). + expect(container.querySelectorAll("li p").length).toBe(2); + }); + + it("the streaming incremental view CONVERGES to the canonical render on finish", () => { + // Mount mid-stream (live tail) — the DOM here is the incremental view. + const { container, rerender } = render( + + + , + ); + // Finish the turn: state flips to done AND the turn is no longer streaming. + const done = msg([{ type: "text", text: ANSWER, state: "done" }]); + rerender( + + + , + ); + // After finish there is exactly ONE canonical markdown container whose HTML is + // the whole-text render — identical to the non-streaming path above. + const blocks = container.querySelectorAll(`.${classes.markdown}`); + expect(blocks.length).toBe(1); + expect(blocks[0].innerHTML).toBe(renderChatMarkdown(ANSWER, {})); + }); + + it("neutralizeInternalLinks is honored on the finalized render", () => { + const linkAnswer = "См. [страницу](/p/abc)."; + const { container } = render( + + + , + ); + const block = container.querySelector(`.${classes.markdown}`); + expect(block!.innerHTML).toBe( + renderChatMarkdown(linkAnswer, { neutralizeInternalLinks: true }), + ); + // The internal link was made inert (no href) by the neutralization flag. + const a = container.querySelector("a"); + expect(a?.hasAttribute("href")).toBe(false); + }); +}); diff --git a/apps/client/src/features/ai-chat/components/message-item.tsx b/apps/client/src/features/ai-chat/components/message-item.tsx index 878477b4..0cf7b7a4 100644 --- a/apps/client/src/features/ai-chat/components/message-item.tsx +++ b/apps/client/src/features/ai-chat/components/message-item.tsx @@ -4,6 +4,7 @@ import { useTranslation } from "react-i18next"; import type { UIMessage } from "@ai-sdk/react"; import ToolCallCard from "@/features/ai-chat/components/tool-call-card.tsx"; import ReasoningBlock from "@/features/ai-chat/components/reasoning-block.tsx"; +import { StreamingMarkdownText } from "@/features/ai-chat/components/streaming-markdown-text.tsx"; import ChatErrorAlert from "@/features/ai-chat/components/chat-error-alert.tsx"; import ChatStoppedNotice from "@/features/ai-chat/components/chat-stopped-notice.tsx"; import { ToolUiPart, isToolPart } from "@/features/ai-chat/utils/tool-parts.tsx"; @@ -86,17 +87,39 @@ interface MessageItemProps { * One assistant text part rendered as sanitized markdown. Memoized on its inputs * so a finalized text part is NOT re-parsed on every streamed delta: during a * turn only the actively-growing tail part changes its `text`, so every earlier - * part hits the memo and skips the expensive marked + DOMPurify pass. Props are - * primitives, so React.memo's default shallow compare is exactly right (the - * `text` string is compared by value). + * part hits the memo and skips the expensive canonical parse + DOMPurify pass. + * Props are primitives, so React.memo's default shallow compare is exactly right + * (the `text` string is compared by value). + * + * Streaming gate (#492) — mirrors ReasoningBlock: + * - `streaming` (this is the live, actively-growing tail part of an in-flight + * turn): render incrementally via StreamingMarkdownText — the stabilized blocks + * go through the canonical pipeline (each parsed ONCE, memoized) and only the + * live tail is cheap plain text. This makes the per-tick cost O(new blocks), + * not the pre-#492 O(ticks) whole-answer re-parse on every ~20Hz delta. + * - finalized (the common case, and the turn-end flip): render the WHOLE text + * through ONE canonical pass — byte-identical to the pre-#492 output (visual + * parity). The row re-renders on the streaming→done flip because + * `messageSignature` tracks each part's `state` (and `turnStreaming` flips at + * turn end), so the incremental view always converges to this single render. */ const MarkdownPart = memo(function MarkdownPart({ text, neutralizeInternalLinks, + streaming, }: { text: string; neutralizeInternalLinks: boolean; + streaming: boolean; }) { + if (streaming) { + return ( + + ); + } const html = renderChatMarkdown(text, { neutralizeInternalLinks }); if (html) { return ( @@ -179,47 +202,10 @@ function MessageItem({ {resolveAssistantName(assistantName) ?? t("AI agent")} {message.parts.map((part, index) => { - if (part.type === "reasoning") { - // Reasoning ("thinking") -> a collapsible block with its own token - // count. Empty/whitespace reasoning with no authoritative count carries - // nothing to show, so skip it (avoids an empty 0-token block). - const text = (part as { text?: string }).text ?? ""; - if (!text.trim() && !(reasoningTokens && reasoningTokens > 0)) - return null; - // Absent state (persisted rows) and "done" both mean finalized. - // `messageSignature` already includes each part's `state`, so the - // streaming→done flip changes the row signature and re-renders this - // row — which is what lets ReasoningBlock switch from chunked plain - // text to its one-time markdown parse (see reasoning-block.tsx). - // ALSO require the turn to be live: a part stranded at - // `state:"streaming"` after the turn ended (no `reasoning-end` — see - // the `turnStreaming` prop doc) must still finalize and parse. - const streaming = - turnStreaming && (part as { state?: string }).state === "streaming"; - return ( - - ); - } - - if (part.type === "text") { - // Skip empty/whitespace-only text parts (a streaming message often - // starts with an empty text part before the first token arrives); the - // typing indicator covers that gap until real content streams in. - if (!part.text.trim()) return null; - return ( - - ); - } - + // Tool parts (`tool-*` / `dynamic-tool`) are template-literal kinds, so + // they cannot be a `switch` case; the runtime guard handles them, and the + // switch below covers every CLOSED (literal-typed) part kind with a + // compile-time exhaustiveness check in its default. if (isToolPart(part.type)) { return ( a collapsible block with its own token + // count. Empty/whitespace reasoning with no authoritative count + // carries nothing to show, so skip it (avoids an empty 0-token block). + const text = part.text ?? ""; + if (!text.trim() && !(reasoningTokens && reasoningTokens > 0)) + return null; + // Absent state (persisted rows) and "done" both mean finalized. + // `messageSignature` already includes each part's `state`, so the + // streaming→done flip changes the row signature and re-renders this + // row — which is what lets ReasoningBlock switch from chunked plain + // text to its one-time markdown parse (see reasoning-block.tsx). + // ALSO require the turn to be live: a part stranded at + // `state:"streaming"` after the turn ended (no `reasoning-end` — see + // the `turnStreaming` prop doc) must still finalize and parse. + const streaming = turnStreaming && part.state === "streaming"; + return ( + + ); + } + + case "text": { + // Skip empty/whitespace-only text parts (a streaming message often + // starts with an empty text part before the first token arrives); the + // typing indicator covers that gap until real content streams in. + if (!part.text.trim()) return null; + // The live, actively-growing tail part of the in-flight turn renders + // incrementally (see MarkdownPart); a finalized part (persisted, or + // the turn-end flip) renders the whole text through one canonical + // pass. Same liveness rule as the reasoning branch above. + const streaming = turnStreaming && part.state === "streaming"; + return ( + + ); + } + + case "source-url": + case "source-document": + case "file": + case "step-start": + // Not surfaced in the chat bubble (v1) — same as the pre-#492 default. + return null; + + default: { + // Compile-time exhaustiveness over the CLOSED union members: every + // literal-typed part kind is handled above, so the only kinds that + // can reach here are the OPEN template-literal ones (`tool-*` — caught + // by the guard at runtime — and `data-*`) plus `dynamic-tool`. Adding + // a NEW closed part kind to UIMessagePart makes this assignment fail + // to compile, forcing it to be handled instead of silently ignored + // (this replaces the pre-#492 fall-through `return null` + WARNING). + const _exhaustive: + | `tool-${string}` + | "dynamic-tool" + | `data-${string}` = part.type; + void _exhaustive; + return null; + } + } })} {/* A persisted turn error (server stored it in metadata.error). Rendered here so it survives a thread remount and shows in reopened history. */} diff --git a/apps/client/src/features/ai-chat/components/streaming-markdown-text.tsx b/apps/client/src/features/ai-chat/components/streaming-markdown-text.tsx new file mode 100644 index 00000000..7a6fdca6 --- /dev/null +++ b/apps/client/src/features/ai-chat/components/streaming-markdown-text.tsx @@ -0,0 +1,96 @@ +import { memo, useMemo } from "react"; +import { splitPlainChunks } from "@/features/ai-chat/components/streaming-plain-text.tsx"; +import { renderChatMarkdown } from "@/features/ai-chat/utils/markdown.ts"; +import classes from "@/features/ai-chat/components/ai-chat.module.css"; + +/** + * One STABILIZED markdown block, rendered through the canonical pipeline and + * memoized on its string prop. During streaming only the TAIL chunk grows (the + * `splitPlainChunks` append-only invariant guarantees every earlier chunk is + * byte-identical across deltas), so React skips every stable block and each one + * is parsed by `renderChatMarkdown` EXACTLY ONCE — turning the pre-#492 + * "re-parse the whole accumulated answer on every ~20Hz tick" (O(ticks)) into + * O(number of blocks). The markup is DOMPurify-sanitized inside renderChatMarkdown + * before it reaches `dangerouslySetInnerHTML`. + * + * NOTE (transient streaming-only artifact): a safe cut is a blank-line boundary, + * so a construct that legitimately contains a blank line (e.g. a fenced code block + * with an empty line) can be split across chunks and render oddly WHILE it is still + * streaming. This is cosmetic and self-heals: the moment the part finalizes, + * MarkdownPart renders the WHOLE text through one canonical pass (visual parity + * with the pre-#492 output). The reasoning path makes the same trade (plain text + * while streaming, one markdown parse at the end). + */ +const MarkdownChunk = memo(function MarkdownChunk({ + text, + neutralizeInternalLinks, +}: { + text: string; + neutralizeInternalLinks: boolean; +}) { + const html = renderChatMarkdown(text, { neutralizeInternalLinks }); + if (html) { + return ( +
    + ); + } + // Malformed/unsupported markdown could not render synchronously: raw text. + return ( +
    + {text} +
    + ); +}); + +/** + * The cheap streaming-time stand-in for the finalized answer's one-time markdown + * parse (see MarkdownPart in message-item.tsx). Mirrors StreamingPlainText's + * chunked-memo pattern but renders the STABILIZED prefix as real markdown (each + * block parsed once, memoized) and only the LIVE tail as flat plain text — so the + * user sees formatted output for everything up to the last safe cut, and the not- + * yet-stable tail (which markdown-parsing every tick would make O(ticks)) stays a + * single cheap escaped text node until it stabilizes into a new block. + * + * `splitPlainChunks` yields chunks where, under append-only growth, every chunk + * except the LAST is immutable; the last chunk is the live tail. Index keys are + * therefore stable (a given index never changes to a different chunk's content). + */ +export function StreamingMarkdownText({ + text, + neutralizeInternalLinks, +}: { + text: string; + neutralizeInternalLinks: boolean; +}) { + const chunks = useMemo(() => splitPlainChunks(text), [text]); + return ( + <> + {chunks.map((chunk, index) => + index < chunks.length - 1 ? ( + + ) : ( + // The live tail: flat, React-escaped plain text (no markdown parse, no + // sanitizer, no innerHTML). `pre-wrap` preserves its newlines; trailing + // separator newlines are dropped at display time so the block gap comes + // from the markdown margins, not a doubled empty line (mirrors + // PlainChunk in streaming-plain-text.tsx). +
    + {chunk.replace(/\n+$/, "")} +
    + ), + )} + + ); +} From c96fafc4ad2c37abc8f3ef765a3c80012cd11230 Mon Sep 17 00:00:00 2001 From: agent_coder Date: Sun, 12 Jul 2026 05:02:25 +0300 Subject: [PATCH 2/3] =?UTF-8?q?perf(ai-chat):=20append-=D0=BF=D0=B5=D1=80?= =?UTF-8?q?=D1=81=D0=B8=D1=81=D1=82=20=D1=88=D0=B0=D0=B3=D0=BE=D0=B2=20?= =?UTF-8?q?=E2=80=94=20per-step=20INSERT=20=D0=B2=D0=BC=D0=B5=D1=81=D1=82?= =?UTF-8?q?=D0=BE=20=D0=BF=D0=B5=D1=80=D0=B5=D0=BF=D0=B8=D1=81=D0=B8=20?= =?UTF-8?q?=D1=81=D1=82=D1=80=D0=BE=D0=BA=D0=B8=20(#492)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Раньше каждый onStepFinish переписывал ВСЮ строку ассистента (растущий metadata.parts jsonb со всеми выводами инструментов) → O(n²) объёма записи на прогон: под MVCC/TOAST апдейт jsonb переписывает всю версию строки, так что шаг k пишет ~k×вывод. Прогон из 50 шагов по ~100 КБ = сотни МБ WAL и мёртвых кортежей за ход, что молотит autovacuum. (#490 убрал только ВТОРУЮ копию в tool_calls; сам metadata.parts всё ещё рос и переписывался.) Теперь каждый завершённый шаг ДОПИСЫВАЕТСЯ отдельной строкой в лёгкую таблицу ai_chat_run_steps (только парты этого шага), а строка сообщения получает дешёвый маркер (stepsPersisted + toolTraceVersion, без растущего блоба parts). Полный metadata.parts собирается ОДИН раз на финализации. НЕ jsonb-append (||): апдейт всё равно переписывает всю TOAST-версию — экономится только сетевой payload, а WAL/мёртвые кортежи остаются; поэтому именно ОТДЕЛЬНАЯ таблица + INSERT. Три обязательные интеграции: - reconstructRunParts(row, stepRows) → { parts, stepsPersisted }: единый шов переключения бэкенда. Читает парты из СТРОКИ, если она уже несёт inline-parts (старые записи + ЛЮБАЯ финализированная), иначе из ТАБЛИЦЫ ШАГОВ (mid-run запись #492). Дискриминатор — наличие непустого metadata.parts (флаг схемы не нужен). Потребители (attach-seed, delta-poll, export, reconnect) прогоняют строки через hydrateAssistantParts на границе чтения — их контракт/вывод не меняется, старые и новые записи восстанавливаются идентично. - сигнал ротации кольца реестра #491 (confirmPersistedStep) теперь стреляет на подтверждённый INSERT шага, под тем же контрактом (updateStreaming возвращает stepsPersisted / null). - era-marker toolTraceVersion (#490) больше не ставится полной переписью — ставится в маркере шага и на финализации (flushAssistant), остаётся консистентным. Полная обратная совместимость: прогон, записанный по-старому (полная строка, без строк шагов), восстанавливается/attach/export идентично. При отсутствии репозитория шагов (позиционные тест-конструкции) — фолбэк на прежний полнострочный flush (без регрессии, только без выигрыша WAL). Тесты (реальный pg, int-lane): - WAL-дельта (pg_current_wal_lsn) на прогоне 40×100КБ: new=4.3МБ vs old=90.3МБ (20.8x) — O(Σ шагов) против O(n²); старый путь в тесте И есть ревертнутое поведение (мутация-проверка). - reconstruct-контракт: new-style (таблица шагов) и old-style (inline) прогоны восстанавливаются в идентичные parts; hydrate заполняет строку. - миграция up/down roundtrip. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../src/core/ai-chat/ai-chat.controller.ts | 61 +++++- .../src/core/ai-chat/ai-chat.service.ts | 183 +++++++++++++++++- apps/server/src/database/database.module.ts | 3 + .../20260708T120000-ai-chat-run-steps.ts | 70 +++++++ .../repos/ai-chat/ai-chat-run-step.repo.ts | 111 +++++++++++ apps/server/src/database/types/db.d.ts | 17 ++ .../server/src/database/types/entity.types.ts | 7 + .../ai-chat-append-persist-wal.int-spec.ts | 173 +++++++++++++++++ .../ai-chat-run-step-repo.int-spec.ts | 164 ++++++++++++++++ .../ai-chat-run-steps-migration.int-spec.ts | 63 ++++++ 10 files changed, 837 insertions(+), 15 deletions(-) create mode 100644 apps/server/src/database/migrations/20260708T120000-ai-chat-run-steps.ts create mode 100644 apps/server/src/database/repos/ai-chat/ai-chat-run-step.repo.ts create mode 100644 apps/server/test/integration/ai-chat-append-persist-wal.int-spec.ts create mode 100644 apps/server/test/integration/ai-chat-run-step-repo.int-spec.ts create mode 100644 apps/server/test/integration/ai-chat-run-steps-migration.int-spec.ts diff --git a/apps/server/src/core/ai-chat/ai-chat.controller.ts b/apps/server/src/core/ai-chat/ai-chat.controller.ts index 2b8aa6ab..66822728 100644 --- a/apps/server/src/core/ai-chat/ai-chat.controller.ts +++ b/apps/server/src/core/ai-chat/ai-chat.controller.ts @@ -35,6 +35,7 @@ import { import { PaginationOptions } from '@docmost/db/pagination/pagination-options'; import { AiChatRepo } from '@docmost/db/repos/ai-chat/ai-chat.repo'; import { AiChatMessageRepo } from '@docmost/db/repos/ai-chat/ai-chat-message.repo'; +import { AiChatRunStepRepo } from '@docmost/db/repos/ai-chat/ai-chat-run-step.repo'; import { PageRepo } from '@docmost/db/repos/page/page.repo'; import { UserThrottlerGuard } from '../../integrations/throttle/user-throttler.guard'; import { AI_CHAT_THROTTLER } from '../../integrations/throttle/throttler-names'; @@ -43,6 +44,8 @@ import { AiChatRunHooks, AiChatService, AiChatStreamBody, + rowHasInlineParts, + hydrateAssistantParts, } from './ai-chat.service'; import { AiChatRunService } from './ai-chat-run.service'; import { AiTranscriptionService } from './ai-transcription.service'; @@ -129,8 +132,39 @@ export class AiChatController { // production. Only touched on the resumable-stream (flag-on) path. private readonly streamRegistry?: AiChatStreamRegistryService, private readonly environment?: EnvironmentService, + // #492: reconstruct a #492 mid-run record's parts from the steps table before + // returning rows to the client / export. OPTIONAL so positional controller + // specs compile unchanged; when absent, hydration is skipped (old-era rows + // already carry inline parts, so nothing to reconstruct). + private readonly aiChatRunStepRepo?: AiChatRunStepRepo, ) {} + /** + * Reconstruct parts for any assistant rows that don't carry them INLINE — a + * #492 mid-run record whose per-step parts live in `ai_chat_run_steps` (the + * append-persist backend). Every FINISHED row (old-era + #492) and every old-era + * streaming snapshot already has inline `metadata.parts`, so the common path + * fetches NOTHING and returns the rows untouched; only an actively-streaming + * new-style row triggers the batch step fetch. Consumers (seed/poll/export) read + * `metadata.parts` off the returned rows exactly as before — the era switch is + * invisible to them (reconstructRunParts contract). + */ + private async withReconstructedParts( + rows: AiChatMessage[], + workspaceId: string, + ): Promise { + if (!this.aiChatRunStepRepo) return rows; + const needy = rows.filter( + (r) => r.role === 'assistant' && !rowHasInlineParts(r), + ); + if (needy.length === 0) return rows; + const stepsByMessage = await this.aiChatRunStepRepo.findByMessageIds( + needy.map((r) => r.id), + workspaceId, + ); + return hydrateAssistantParts(rows, stepsByMessage); + } + /** List the requesting user's chats in this workspace (paginated). */ @HttpCode(HttpStatus.OK) @Post('chats') @@ -184,11 +218,17 @@ export class AiChatController { @AuthWorkspace() workspace: Workspace, ) { await this.assertOwnedChat(dto.chatId, user, workspace); - return this.aiChatMessageRepo.findByChat( + const page = await this.aiChatMessageRepo.findByChat( dto.chatId, workspace.id, pagination, ); + // #492: reconstruct parts for any active new-style row so the client seed sees + // `metadata.parts` unchanged (a no-op for the finished rows that fill a page). + return { + ...page, + items: await this.withReconstructedParts(page.items, workspace.id), + }; } /** @@ -225,7 +265,10 @@ export class AiChatController { workspace.id, ); return { - rows, + // #492: the delta of an actively-streaming new-style row carries its parts + // reconstructed from the steps table, so the degraded poll shows persisted + // progress exactly as the pre-#492 full-row snapshot did. + rows: await this.withReconstructedParts(rows, workspace.id), cursor, run: run ? { id: run.id, status: run.status } : null, }; @@ -247,8 +290,10 @@ export class AiChatController { @AuthWorkspace() workspace: Workspace, ): Promise<{ markdown: string }> { const chat = await this.assertOwnedChat(dto.chatId, user, workspace); - const rows = await this.aiChatMessageRepo.findAllByChat( - dto.chatId, + const rows = await this.withReconstructedParts( + await this.aiChatMessageRepo.findAllByChat(dto.chatId, workspace.id), + // #492: an interrupted-but-still-active turn exports its persisted steps + // (reconstructed from the steps table) just like the pre-#492 full row did. workspace.id, ); const markdown = buildChatMarkdown({ @@ -288,7 +333,13 @@ export class AiChatController { workspace.id, ) : undefined; - return { run, message: message ?? null }; + // #492: reconnect to an IN-FLIGHT run reconstructs the projection row's parts + // from the steps table (the row itself carries only the step marker mid-run); + // a finished run's row already has inline parts, so this is a no-op. + const [hydrated] = message + ? await this.withReconstructedParts([message], workspace.id) + : [undefined]; + return { run, message: hydrated ?? null }; } /** 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 9922ba5a..201c4cf8 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.ts @@ -22,6 +22,7 @@ import { AiSettingsService } from '../../integrations/ai/ai-settings.service'; import { describeProviderError } from '../../integrations/ai/ai-error.util'; import { AiChatRepo } from '@docmost/db/repos/ai-chat/ai-chat.repo'; import { AiChatMessageRepo } from '@docmost/db/repos/ai-chat/ai-chat-message.repo'; +import { AiChatRunStepRepo } from '@docmost/db/repos/ai-chat/ai-chat-run-step.repo'; import { AiChatPageSnapshotRepo } from '@docmost/db/repos/ai-chat/ai-chat-page-snapshot.repo'; import { AiAgentRoleRepo } from '@docmost/db/repos/ai-agent-roles/ai-agent-roles.repo'; import { PageRepo } from '@docmost/db/repos/page/page.repo'; @@ -517,6 +518,12 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy { // constructions compile unchanged; Nest always injects the real singleton, so // reconcile sees the SAME in-memory active/zombie maps the runner mutates. private readonly aiChatRunService?: AiChatRunService, + // #492 append-persist: per-step INSERT into the lightweight steps table (the + // O(Σ steps) replacement for the O(n²) full-row `metadata.parts` rewrite). + // OPTIONAL so existing positional constructions (int-specs) compile unchanged; + // Nest injects the real singleton. When ABSENT the per-step path falls back to + // the pre-#492 full-row flush (no regression, only no WAL win). + private readonly aiChatRunStepRepo?: AiChatRunStepRepo, ) {} // #487: periodic reconcile timer (single-process phase 1). Started in @@ -1557,17 +1564,57 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy { // connection when finalize runs, so the SQL `WHERE status='streaming'` // (not this flag) is what prevents it clobbering the terminal row. if (finalized) return null; - // Build the flush ONCE so the returned count is EXACTLY the persisted - // `stepsPersisted` (both derive from capturedSteps.length at this instant). - const flushed = flushAssistant(capturedSteps, '', 'streaming', { - pageChanged, - partsCache, - }); - const stepsPersisted = flushed.metadata.stepsPersisted as number; + // The count derives from capturedSteps.length at THIS instant, so the + // returned value is EXACTLY the persisted `stepsPersisted` the ring rotates + // on (whether we take the append-persist path or the legacy fallback). + const stepsPersisted = capturedSteps.length; try { - await this.aiChatMessageRepo.update(assistantId, workspace.id, flushed, { - onlyIfStreaming: true, - }); + if (this.aiChatRunStepRepo) { + // #492 APPEND-PERSIST: write only THIS finished step's parts to the + // steps table (O(step) WAL), then bump the row's CHEAP step marker — + // NO growing `metadata.parts` blob (that O(n²) full-row rewrite is + // exactly what this removes). The full `metadata.parts` is assembled + // once at finalize; a mid-run resume seed is reconstructed from the + // step rows (reconstructRunParts). The INSERT is idempotent + // (ON CONFLICT DO NOTHING), so a re-fired step never doubles the parts. + const index = stepsPersisted - 1; + if (index >= 0) { + const stepParts = assistantParts( + [capturedSteps[index]], + '', + partsCache, + ); + await this.aiChatRunStepRepo.insertStep( + assistantId, + workspace.id, + index, + stepParts, + ); + } + // Marker UPDATE: advance stepsPersisted + keep the toolTrace era marker + // (bumps updatedAt so the delta poll observes the step, and carries the + // frontier a resuming client attaches from). Scoped onlyIfStreaming so a + // late marker never clobbers the terminal finalize. + await this.aiChatMessageRepo.update( + assistantId, + workspace.id, + { metadata: stepMarkerMetadata(stepsPersisted) }, + { onlyIfStreaming: true }, + ); + } else { + // Legacy fallback (no steps table wired — positional test builds): the + // pre-#492 full-row flush, so parts still land inline on the row. + const flushed = flushAssistant(capturedSteps, '', 'streaming', { + pageChanged, + partsCache, + }); + await this.aiChatMessageRepo.update( + assistantId, + workspace.id, + flushed, + { onlyIfStreaming: true }, + ); + } return stepsPersisted; } catch (err) { this.logger.warn( @@ -2747,6 +2794,122 @@ export function rowToUiMessage(row: AiChatMessage): Omit & { return { id: row.id, role, parts: parts as UIMessage['parts'] }; } +/** + * Cheap step-marker metadata for the #492 per-step UPDATE. Advances + * `stepsPersisted` (the resume attach frontier) and keeps the `toolTraceVersion` + * era marker, WITHOUT the growing `parts` blob (those live in the steps table + * now; the full `metadata.parts` is assembled once at finalize by flushAssistant). + * `parts: []` is kept for shape stability — it reads as an empty inline-parts row, + * which is exactly the discriminator that routes reconstruction to the steps table. + */ +export function stepMarkerMetadata( + stepsPersisted: number, +): Record { + return { parts: [], toolTraceVersion: 2, stepsPersisted }; +} + +/** + * Whether an assistant row already carries its full UI parts INLINE on the row + * (`metadata.parts`). TRUE for every FINISHED row — old-era rows AND #492 rows, + * whose full parts are assembled once at finalize — and for old-era streaming + * snapshots (the pre-#492 per-step full-row flush). FALSE for a #492 MID-RUN + * record, whose per-step parts live in the `ai_chat_run_steps` table. This is the + * era discriminator the reconstruct seam branches on — no schema flag needed. + */ +export function rowHasInlineParts(row: { metadata?: unknown }): boolean { + const meta = (row.metadata ?? {}) as { parts?: unknown }; + return Array.isArray(meta.parts) && meta.parts.length > 0; +} + +/** + * Concatenate persisted per-step parts (in `stepIndex` order) into the turn's UI + * parts (#492). Reproduces EXACTLY what flushAssistant → assistantParts would have + * written to `metadata.parts` for those finished steps, since each step row stored + * `assistantParts([step])` at persist time. + */ +export function assembleStepParts( + stepRows: ReadonlyArray<{ stepIndex: number; parts: unknown }>, +): UIMessage['parts'] { + const parts: Array> = []; + for (const step of [...stepRows].sort((a, b) => a.stepIndex - b.stepIndex)) { + if (Array.isArray(step.parts)) { + parts.push(...(step.parts as Array>)); + } + } + return parts as UIMessage['parts']; +} + +/** + * reconstructRunParts (#492) — the single backend-switch seam. Given an assistant + * ROW and its persisted step rows, return the turn's UI `parts` + the persisted + * step count, reading from the ROW when it already carries inline parts (old-era + * records AND every finished record) and from the STEPS TABLE otherwise (a #492 + * mid-run record). The higher-level consumers (attach seed, delta poll, export) + * route their row→parts through this / {@link hydrateAssistantParts}, so old and + * new records reconstruct identically WITHOUT the consumers branching on the era. + */ +export function reconstructRunParts( + row: { metadata?: unknown; content?: string | null }, + stepRows: ReadonlyArray<{ stepIndex: number; parts: unknown }>, +): { parts: UIMessage['parts']; stepsPersisted: number } { + if (rowHasInlineParts(row)) { + const meta = row.metadata as { + parts: UIMessage['parts']; + stepsPersisted?: number; + }; + return { + parts: meta.parts, + stepsPersisted: + typeof meta.stepsPersisted === 'number' + ? meta.stepsPersisted + : stepRows.length, + }; + } + if (stepRows.length > 0) { + return { + parts: assembleStepParts(stepRows), + stepsPersisted: stepRows.length, + }; + } + // No inline parts and no step rows: an old-era seed / empty streaming row. Fall + // back to a single text part from `content` (mirrors rowToUiMessage). + return { + parts: textPart(row.content ?? '') as UIMessage['parts'], + stepsPersisted: 0, + }; +} + +/** + * Fill each assistant row's `metadata.parts` from its step rows when the row does + * not already carry them inline (a #492 mid-run record), so a consumer that reads + * `metadata.parts` off the RAW row (the client seed/poll, the Markdown export) + * sees the reconstructed parts with NO change to itself. Rows that already have + * inline parts (old-era + finished) and non-assistant rows pass through untouched. + * Pure: returns new row objects, never mutates the inputs. + */ +export function hydrateAssistantParts< + T extends { id: string; role?: string; metadata?: unknown }, +>( + rows: ReadonlyArray, + stepsByMessage: Map< + string, + ReadonlyArray<{ stepIndex: number; parts: unknown }> + >, +): T[] { + return rows.map((row) => { + if (row.role !== 'assistant' || rowHasInlineParts(row)) return row; + const steps = stepsByMessage.get(row.id); + if (!steps || steps.length === 0) return row; + return { + ...row, + metadata: { + ...((row.metadata ?? {}) as Record), + parts: assembleStepParts(steps), + }, + }; + }); +} + /** * The persisted-row patch shape produced by {@link flushAssistant}. It is the * SAME shape the assistant repo insert/update consume (content + toolCalls + diff --git a/apps/server/src/database/database.module.ts b/apps/server/src/database/database.module.ts index 3e9218cc..71f87d14 100644 --- a/apps/server/src/database/database.module.ts +++ b/apps/server/src/database/database.module.ts @@ -32,6 +32,7 @@ import { TemplateRepo } from '@docmost/db/repos/template/template.repo'; import { AiChatRepo } from '@docmost/db/repos/ai-chat/ai-chat.repo'; import { AiChatMessageRepo } from '@docmost/db/repos/ai-chat/ai-chat-message.repo'; import { AiChatRunRepo } from '@docmost/db/repos/ai-chat/ai-chat-run.repo'; +import { AiChatRunStepRepo } from '@docmost/db/repos/ai-chat/ai-chat-run-step.repo'; import { AiChatPageSnapshotRepo } from '@docmost/db/repos/ai-chat/ai-chat-page-snapshot.repo'; import { AiProviderCredentialsRepo } from '@docmost/db/repos/ai-chat/ai-provider-credentials.repo'; import { AiMcpServerRepo } from '@docmost/db/repos/ai-chat/ai-mcp-server.repo'; @@ -124,6 +125,7 @@ import { firstSqlToken } from '../integrations/metrics/metrics.constants'; AiChatRepo, AiChatMessageRepo, AiChatRunRepo, + AiChatRunStepRepo, AiChatPageSnapshotRepo, AiProviderCredentialsRepo, AiMcpServerRepo, @@ -159,6 +161,7 @@ import { firstSqlToken } from '../integrations/metrics/metrics.constants'; AiChatRepo, AiChatMessageRepo, AiChatRunRepo, + AiChatRunStepRepo, AiChatPageSnapshotRepo, AiProviderCredentialsRepo, AiMcpServerRepo, diff --git a/apps/server/src/database/migrations/20260708T120000-ai-chat-run-steps.ts b/apps/server/src/database/migrations/20260708T120000-ai-chat-run-steps.ts new file mode 100644 index 00000000..e576b09e --- /dev/null +++ b/apps/server/src/database/migrations/20260708T120000-ai-chat-run-steps.ts @@ -0,0 +1,70 @@ +import { type Kysely, sql } from 'kysely'; + +/** + * `ai_chat_run_steps` — append-only per-step persistence for an assistant turn + * (#492 wave C). Each finished agent step's UI `parts` (its text part + a part + * per tool call, WITH the tool output) is INSERTed as its own lightweight row the + * moment the step ends, instead of REWRITING the whole assistant row's growing + * `metadata.parts` jsonb on every `onStepFinish`. + * + * WHY a separate table + INSERT (not a jsonb `||` append on the message row): a + * Postgres jsonb UPDATE rewrites the ENTIRE TOASTed row version under MVCC, so + * re-persisting a growing `metadata.parts` on every step is O(n²) write volume + * (a 50-step run with ~100 KB tool outputs wrote hundreds of MB of WAL / dead + * tuples per turn, hammering autovacuum). `||` would only shave the network + * payload — the WAL/TOAST rewrite harm remains. An INSERT into a per-step table + * writes ONLY that step's bytes, so the per-turn write volume is O(Σ steps). + * + * The full `metadata.parts` on the message row is assembled ONCE at finalize (the + * terminal completed/error/aborted write). Mid-run, a resuming client's seed is + * reconstructed by concatenating these step rows in `step_index` order — which + * reproduces exactly what the old per-step full-row rewrite persisted. Records + * written the OLD way (full `metadata.parts` on the row, no step rows) still + * reconstruct from the row unchanged; the two eras are distinguished by whether + * the row already carries non-empty `metadata.parts` (see reconstructRunParts / + * assembleStepParts in ai-chat.service.ts). + * + * ON DELETE CASCADE on `message_id`: the step rows are a derived projection of the + * assistant message; they must vanish with it (or with its workspace). + */ +export async function up(db: Kysely): Promise { + await db.schema + .createTable('ai_chat_run_steps') + .ifNotExists() + .addColumn('id', 'uuid', (col) => + col.primaryKey().defaultTo(sql`gen_uuid_v7()`), + ) + // The assistant message row this step belongs to (the #183 projection). The + // step rows are a derived, per-step slice of that message, so they cascade. + .addColumn('message_id', 'uuid', (col) => + col.references('ai_chat_messages.id').onDelete('cascade').notNull(), + ) + .addColumn('workspace_id', 'uuid', (col) => + col.references('workspaces.id').onDelete('cascade').notNull(), + ) + // 0-based index of the finished step within the turn. Ordering key for + // reconstruction; unique per message (idempotent step re-persist). + .addColumn('step_index', 'integer', (col) => col.notNull()) + // The step's UI parts (text part + a `tool-*` part per call, WITH output). + // Concatenated in step order to rebuild the turn's `metadata.parts`. + .addColumn('parts', 'jsonb', (col) => col.notNull()) + .addColumn('created_at', 'timestamptz', (col) => + col.notNull().defaultTo(sql`now()`), + ) + .execute(); + + // Idempotent per-step persist: a retried INSERT of the same (message, step) + // is a no-op (the service uses ON CONFLICT DO NOTHING). This also serves the + // reconstruction read (WHERE message_id ORDER BY step_index). + await db.schema + .createIndex('ai_chat_run_steps_message_step_uidx') + .ifNotExists() + .on('ai_chat_run_steps') + .columns(['message_id', 'step_index']) + .unique() + .execute(); +} + +export async function down(db: Kysely): Promise { + await db.schema.dropTable('ai_chat_run_steps').ifExists().execute(); +} diff --git a/apps/server/src/database/repos/ai-chat/ai-chat-run-step.repo.ts b/apps/server/src/database/repos/ai-chat/ai-chat-run-step.repo.ts new file mode 100644 index 00000000..656186c6 --- /dev/null +++ b/apps/server/src/database/repos/ai-chat/ai-chat-run-step.repo.ts @@ -0,0 +1,111 @@ +import { Injectable } from '@nestjs/common'; +import { InjectKysely } from 'nestjs-kysely'; +import { sql } from 'kysely'; +import { KyselyDB, KyselyTransaction } from '../../types/kysely.types'; +import { dbOrTx } from '../../utils'; +import { AiChatRunStep } from '@docmost/db/types/entity.types'; + +/** + * Append-only per-step persistence for an assistant turn (#492). Each finished + * agent step's UI `parts` (its text part + a `tool-*` part per call, WITH the + * tool output) is INSERTed as its own lightweight row the moment the step ends — + * instead of REWRITING the assistant row's growing `metadata.parts` jsonb on every + * `onStepFinish` (a Postgres jsonb UPDATE rewrites the whole TOASTed row version + * under MVCC, so that was O(n²) WAL/dead-tuple churn per turn). + * + * The full `metadata.parts` on the message row is assembled ONCE at finalize; + * mid-run, a resuming client's seed is rebuilt from these rows in `stepIndex` + * order (see `assembleStepParts` / the reconstruct seam in ai-chat.service.ts). + * Every method is workspace-scoped as defense-in-depth. + */ +@Injectable() +export class AiChatRunStepRepo { + constructor(@InjectKysely() private readonly db: KyselyDB) {} + + /** + * Append one finished step's parts. Idempotent: a retried persist of the SAME + * (message, stepIndex) is a no-op via ON CONFLICT DO NOTHING — the per-step + * writes are fired fire-and-forget + serialized, and a duplicate must never + * throw into the stream or double the parts. Returns whether a NEW row landed + * (false = the step was already persisted). + */ + async insertStep( + messageId: string, + workspaceId: string, + stepIndex: number, + parts: unknown, + trx?: KyselyTransaction, + ): Promise { + const db = dbOrTx(this.db, trx); + const inserted = await db + .insertInto('aiChatRunSteps') + .values({ + messageId, + workspaceId, + stepIndex, + // jsonb column: cast through never (same pattern as the message repo). + parts: parts as never, + }) + .onConflict((oc) => oc.columns(['messageId', 'stepIndex']).doNothing()) + .returning('id') + .executeTakeFirst(); + return inserted !== undefined; + } + + /** All persisted steps for ONE assistant message, in step order. */ + async findByMessage( + messageId: string, + workspaceId: string, + ): Promise { + return this.db + .selectFrom('aiChatRunSteps') + .selectAll('aiChatRunSteps') + .where('messageId', '=', messageId) + .where('workspaceId', '=', workspaceId) + .orderBy('stepIndex', 'asc') + .execute(); + } + + /** + * All persisted steps for a SET of assistant messages, grouped by messageId + * (each group in step order). One query for the batch — the hydration seam + * (getMessages / delta / export) calls this only for the rows that actually + * need reconstruction (an active new-style row whose `metadata.parts` is still + * empty), which is usually none, so this is skipped on the common path. + */ + async findByMessageIds( + messageIds: string[], + workspaceId: string, + ): Promise> { + const byMessage = new Map(); + if (messageIds.length === 0) return byMessage; + const rows = await this.db + .selectFrom('aiChatRunSteps') + .selectAll('aiChatRunSteps') + .where('messageId', 'in', messageIds) + .where('workspaceId', '=', workspaceId) + .orderBy('stepIndex', 'asc') + .execute(); + for (const row of rows) { + const list = byMessage.get(row.messageId); + if (list) list.push(row); + else byMessage.set(row.messageId, [row]); + } + return byMessage; + } + + /** + * How many steps are persisted for a message (its step-marker floor). Exposed + * for the reconstruct contract (`reconstructRunParts → { parts, stepsPersisted }`) + * so a caller can align a resume attach without materializing every step's parts. + */ + async countByMessage(messageId: string, workspaceId: string): Promise { + const row = await this.db + .selectFrom('aiChatRunSteps') + .select(sql`count(*)::int`.as('n')) + .where('messageId', '=', messageId) + .where('workspaceId', '=', workspaceId) + .executeTakeFirst(); + return row?.n ?? 0; + } +} diff --git a/apps/server/src/database/types/db.d.ts b/apps/server/src/database/types/db.d.ts index a17d628c..e1495e74 100644 --- a/apps/server/src/database/types/db.d.ts +++ b/apps/server/src/database/types/db.d.ts @@ -691,6 +691,22 @@ export interface AiChatRuns { updatedAt: Generated; } +// Append-only per-step persistence for an assistant turn (#492). Mirrors +// migration 20260708T120000-ai-chat-run-steps.ts. Each finished agent step's UI +// `parts` are INSERTed as their own row (instead of rewriting the message row's +// growing `metadata.parts` jsonb every step — an O(n²) WAL/TOAST churn). The full +// `metadata.parts` is assembled once at finalize; mid-run a resuming client's seed +// is rebuilt by concatenating these rows in `stepIndex` order. Cascades with the +// assistant message row it projects. +export interface AiChatRunSteps { + id: Generated; + messageId: string; + workspaceId: string; + stepIndex: number; + parts: Json; + createdAt: Generated; +} + // Per-(chat,page) snapshot of the open page's Markdown at the END of the agent's // previous turn (#274). Mirrors migration 20260702T120000-ai-chat-page-snapshot.ts. // The next turn diffs the CURRENT Markdown against `contentMd` to surface edits a @@ -728,6 +744,7 @@ export interface DB { aiChats: AiChats; aiChatMessages: AiChatMessages; aiChatRuns: AiChatRuns; + aiChatRunSteps: AiChatRunSteps; aiChatPageSnapshots: AiChatPageSnapshots; apiKeys: ApiKeys; attachments: Attachments; diff --git a/apps/server/src/database/types/entity.types.ts b/apps/server/src/database/types/entity.types.ts index 7ea1a1e1..7230162b 100644 --- a/apps/server/src/database/types/entity.types.ts +++ b/apps/server/src/database/types/entity.types.ts @@ -4,6 +4,7 @@ import { AiChats, AiChatMessages, AiChatRuns, + AiChatRunSteps, AiChatPageSnapshots, Attachments, Comments, @@ -64,6 +65,12 @@ export type InsertableAiChatMessage = Omit, 'tsv'>; export type AiChatRun = Selectable; export type InsertableAiChatRun = Insertable; +// AI Chat Run Step (#492): append-only per-step parts persistence. Each finished +// agent step's UI parts are stored as their own row; the full turn's parts are +// assembled from these (in stepIndex order) for a mid-run resume seed. +export type AiChatRunStep = Selectable; +export type InsertableAiChatRunStep = Insertable; + // AI Chat Page Snapshot (#274): per-(chat,page) Markdown snapshot taken at the // end of the agent's previous turn, diffed against the current page next turn to // detect human edits made between turns. diff --git a/apps/server/test/integration/ai-chat-append-persist-wal.int-spec.ts b/apps/server/test/integration/ai-chat-append-persist-wal.int-spec.ts new file mode 100644 index 00000000..faa2cc2c --- /dev/null +++ b/apps/server/test/integration/ai-chat-append-persist-wal.int-spec.ts @@ -0,0 +1,173 @@ +import { randomBytes } from 'crypto'; +import { Kysely, sql } from 'kysely'; +import { AiChatRunStepRepo } from '@docmost/db/repos/ai-chat/ai-chat-run-step.repo'; +import { AiChatMessageRepo } from '@docmost/db/repos/ai-chat/ai-chat-message.repo'; +import { + assistantParts, + flushAssistant, + stepMarkerMetadata, +} from '../../src/core/ai-chat/ai-chat.service'; +import { + getTestDb, + destroyTestDb, + createWorkspace, + createUser, + createChat, +} from './db'; + +/** + * #492 append-persist — WRITE-VOLUME regression on a LIVE Postgres, measured via + * the `pg_current_wal_lsn()` delta around a realistic multi-step run driven through + * the REAL repos (not a mock — a mock cannot observe MVCC/TOAST rewrite volume, the + * whole point). Proves the core claim: + * + * NEW (per-step INSERT into ai_chat_run_steps + a CHEAP step-marker UPDATE on the + * message row) writes O(Σ steps) of WAL — each step writes only its own bytes. + * + * OLD (the pre-#492 full-row rewrite: re-persist the GROWING metadata.parts on + * every onStepFinish) writes O(n²) — step k rewrites the whole TOASTed jsonb of + * all k prior outputs. + * + * The OLD path here IS the reverted behavior, so this doubles as the mutation + * check: swapping the new path back to `flushAssistant` full-row UPDATEs reddens + * the assertion (OLD is many times larger). + */ +type Step = { + text: string; + toolCalls: Array<{ toolCallId: string; toolName: string; input: unknown }>; + toolResults: Array<{ toolCallId: string; toolName: string; output: unknown }>; +}; + +// ~100 KB INCOMPRESSIBLE output per step (a page read). Random base64 so TOAST +// cannot compress it away and hide the real write volume. +function makeStep(i: number, outputBytes = 100_000): Step { + const body = randomBytes(Math.ceil(outputBytes * 0.75)).toString('base64'); + return { + text: `step ${i} reasoning`, + toolCalls: [ + { toolCallId: `c${i}`, toolName: 'getPage', input: { id: `p${i}` } }, + ], + toolResults: [ + { + toolCallId: `c${i}`, + toolName: 'getPage', + output: { id: `p${i}`, title: `Page ${i}`, body }, + }, + ], + }; +} + +async function walDelta( + db: Kysely, + fn: () => Promise, +): Promise { + const before = ( + await sql<{ l: string }>`select pg_current_wal_lsn() as l`.execute(db) + ).rows[0].l; + await fn(); + // NOTE: no pg_switch_wal() — a segment switch pads the LSN to the next 16 MB + // boundary and would swamp the delta. The raw LSN advances by the WAL bytes. + const after = ( + await sql<{ l: string }>`select pg_current_wal_lsn() as l`.execute(db) + ).rows[0].l; + return Number( + ( + await sql<{ + d: string; + }>`select pg_wal_lsn_diff(${after}::pg_lsn, ${before}::pg_lsn) as d`.execute( + db, + ) + ).rows[0].d, + ); +} + +describe('#492 append-persist write volume (pg_current_wal_lsn delta) [integration]', () => { + let db: Kysely; + let stepRepo: AiChatRunStepRepo; + let msgRepo: AiChatMessageRepo; + let workspaceId: string; + let userId: string; + let chatId: string; + + beforeAll(async () => { + db = getTestDb(); + stepRepo = new AiChatRunStepRepo(db as any); + msgRepo = new AiChatMessageRepo(db as any); + workspaceId = (await createWorkspace(db)).id; + userId = (await createUser(db, workspaceId)).id; + chatId = (await createChat(db, { workspaceId, creatorId: userId })).id; + }); + + afterAll(async () => { + await destroyTestDb(); + }); + + const seedRow = () => + msgRepo.insert({ + chatId, + workspaceId, + userId, + role: 'assistant', + content: '', + status: 'streaming', + metadata: stepMarkerMetadata(0) as never, + }); + + const STEPS = 40; + + it('NEW per-step INSERT is O(Σ steps); OLD full-row rewrite is O(n²)', async () => { + const steps: Step[] = []; + for (let i = 0; i < STEPS; i++) steps.push(makeStep(i)); + + // NEW: per-step INSERT of THIS step's parts + a cheap marker UPDATE. + const newRow = await seedRow(); + const newWal = await walDelta(db, async () => { + for (let i = 0; i < STEPS; i++) { + await stepRepo.insertStep( + newRow.id, + workspaceId, + i, + assistantParts([steps[i]], ''), + ); + await msgRepo.update( + newRow.id, + workspaceId, + { metadata: stepMarkerMetadata(i + 1) }, + { onlyIfStreaming: true }, + ); + } + }); + + // OLD (the pre-#492 revert): re-persist the GROWING metadata.parts on the + // message row on every step. + const oldRow = await seedRow(); + const oldWal = await walDelta(db, async () => { + const acc: Step[] = []; + for (let i = 0; i < STEPS; i++) { + acc.push(steps[i]); + await msgRepo.update( + oldRow.id, + workspaceId, + flushAssistant(acc as never, '', 'streaming'), + { onlyIfStreaming: true }, + ); + } + }); + + // eslint-disable-next-line no-console + console.log( + `[#492 WAL] ${STEPS} steps ×100KB: new=${(newWal / 1e6).toFixed(1)}MB ` + + `old=${(oldWal / 1e6).toFixed(1)}MB (${(oldWal / newWal).toFixed( + 1, + )}x smaller)`, + ); + + // O(Σ steps): ~STEPS × (100KB output + marker) of WAL. 40 × ~100KB parts plus + // 40 tiny markers is a few tens of MB at most — bounded, linear in step count. + expect(newWal).toBeLessThan(30_000_000); + // O(n²): step k rewrites ~k × 100KB. Σ over 40 steps ≈ 80+ MB — far larger. + expect(oldWal).toBeGreaterThan(30_000_000); + // The load-bearing claim: the new path writes a small FRACTION of the old. + expect(newWal).toBeLessThan(oldWal * 0.35); + }, 120_000); +}); diff --git a/apps/server/test/integration/ai-chat-run-step-repo.int-spec.ts b/apps/server/test/integration/ai-chat-run-step-repo.int-spec.ts new file mode 100644 index 00000000..96b71a3a --- /dev/null +++ b/apps/server/test/integration/ai-chat-run-step-repo.int-spec.ts @@ -0,0 +1,164 @@ +import { randomBytes } from 'crypto'; +import { Kysely } from 'kysely'; +import { AiChatRunStepRepo } from '@docmost/db/repos/ai-chat/ai-chat-run-step.repo'; +import { AiChatMessageRepo } from '@docmost/db/repos/ai-chat/ai-chat-message.repo'; +import { + assistantParts, + reconstructRunParts, + hydrateAssistantParts, + stepMarkerMetadata, + rowHasInlineParts, +} from '../../src/core/ai-chat/ai-chat.service'; +import { + getTestDb, + destroyTestDb, + createWorkspace, + createUser, + createChat, +} from './db'; + +/** + * #492 append-persist — the reconstruct CONTRACT on a live Postgres. Proves that a + * turn persisted the NEW way (per-step rows in `ai_chat_run_steps`, only a step + * marker on the message row) reconstructs to the SAME UI parts as a turn persisted + * the OLD way (full `metadata.parts` inline on the row, no step rows) — so the + * era-switch is invisible to attach / delta-poll / export. Real repos + real jsonb + * roundtrip, not a mock (a mock cannot prove the parts survive the jsonb column + * byte-identical). + */ +type Step = { + text: string; + toolCalls: Array<{ toolCallId: string; toolName: string; input: unknown }>; + toolResults: Array<{ toolCallId: string; toolName: string; output: unknown }>; +}; + +// A realistic step: some text + a getPage tool call whose ~100 KB body is +// INCOMPRESSIBLE random base64 (a 'x'.repeat filler would TOAST away and hide the +// real bytes). Under MAX_TOOL_OUTPUT_BYTES (200 KB) it is stored uncompacted. +function makeStep(i: number, outputBytes = 4_000): Step { + const body = randomBytes(Math.ceil(outputBytes * 0.75)).toString('base64'); + return { + text: `step ${i} text`, + toolCalls: [ + { toolCallId: `c${i}`, toolName: 'getPage', input: { id: `p${i}` } }, + ], + toolResults: [ + { + toolCallId: `c${i}`, + toolName: 'getPage', + output: { id: `p${i}`, title: `Page ${i}`, body }, + }, + ], + }; +} + +describe('AiChatRunStepRepo + reconstruct contract [integration]', () => { + let db: Kysely; + let stepRepo: AiChatRunStepRepo; + let msgRepo: AiChatMessageRepo; + let workspaceId: string; + let userId: string; + let chatId: string; + + beforeAll(async () => { + db = getTestDb(); + stepRepo = new AiChatRunStepRepo(db as any); + msgRepo = new AiChatMessageRepo(db as any); + workspaceId = (await createWorkspace(db)).id; + userId = (await createUser(db, workspaceId)).id; + chatId = (await createChat(db, { workspaceId, creatorId: userId })).id; + }); + + afterAll(async () => { + await destroyTestDb(); + }); + + const seedRow = (metadata: unknown, status: string) => + msgRepo.insert({ + chatId, + workspaceId, + userId, + role: 'assistant', + content: '', + status, + metadata: metadata as never, + }); + + it('insertStep is idempotent per (message, stepIndex) and reads back in order', async () => { + const row = await seedRow(stepMarkerMetadata(0), 'streaming'); + const parts0 = assistantParts([makeStep(0)], ''); + const parts1 = assistantParts([makeStep(1)], ''); + + expect(await stepRepo.insertStep(row.id, workspaceId, 0, parts0)).toBe(true); + expect(await stepRepo.insertStep(row.id, workspaceId, 1, parts1)).toBe(true); + // A retried persist of the SAME step is a no-op (ON CONFLICT DO NOTHING). + expect(await stepRepo.insertStep(row.id, workspaceId, 0, parts0)).toBe( + false, + ); + + const steps = await stepRepo.findByMessage(row.id, workspaceId); + expect(steps.map((s) => s.stepIndex)).toEqual([0, 1]); + expect(await stepRepo.countByMessage(row.id, workspaceId)).toBe(2); + + // Batch fetch groups by message id in step order. + const map = await stepRepo.findByMessageIds([row.id], workspaceId); + expect(map.get(row.id)!.map((s) => s.stepIndex)).toEqual([0, 1]); + }); + + it('a NEW-style (step-table) run reconstructs identically to an OLD-style (inline) run', async () => { + const steps = [makeStep(10), makeStep(11)]; + // The inline parts the OLD full-row flush would have written. + const fullParts = assistantParts(steps, ''); + + // OLD-style record: full parts inline on the row, NO step rows. + const oldRow = await seedRow( + { parts: fullParts, toolTraceVersion: 2, stepsPersisted: 2 }, + 'completed', + ); + + // NEW-style record: only a step marker on the row + per-step rows. + const newRow = await seedRow(stepMarkerMetadata(2), 'streaming'); + for (let i = 0; i < steps.length; i++) { + await stepRepo.insertStep( + newRow.id, + workspaceId, + i, + assistantParts([steps[i]], ''), + ); + } + + // Re-read both from the DB (proves the jsonb roundtrip). + const oldFetched = await msgRepo.findById(oldRow.id, workspaceId); + const newFetched = await msgRepo.findById(newRow.id, workspaceId); + const oldSteps = await stepRepo.findByMessage(oldRow.id, workspaceId); + const newSteps = await stepRepo.findByMessage(newRow.id, workspaceId); + + // The discriminator: the old row carries inline parts, the new one does not. + expect(rowHasInlineParts(oldFetched!)).toBe(true); + expect(rowHasInlineParts(newFetched!)).toBe(false); + expect(oldSteps).toHaveLength(0); + expect(newSteps).toHaveLength(2); + + const oldRecon = reconstructRunParts(oldFetched!, oldSteps); + const newRecon = reconstructRunParts(newFetched!, newSteps); + + // Both reconstruct to the SAME parts + step count — the era is invisible. + expect(newRecon.parts).toEqual(fullParts); + expect(oldRecon.parts).toEqual(fullParts); + expect(newRecon.parts).toEqual(oldRecon.parts); + expect(newRecon.stepsPersisted).toBe(2); + expect(oldRecon.stepsPersisted).toBe(2); + + // hydrateAssistantParts fills the new row's metadata.parts to match the old + // row's inline parts — so a consumer reading `metadata.parts` off the raw row + // (the client seed/poll, export) is unchanged across the era. + const map = await stepRepo.findByMessageIds([newRow.id], workspaceId); + const [hydrated] = hydrateAssistantParts([newFetched!], map); + expect((hydrated.metadata as { parts: unknown }).parts).toEqual(fullParts); + // A row that already has inline parts passes through untouched (same ref-shape). + const [oldPassThrough] = hydrateAssistantParts([oldFetched!], map); + expect((oldPassThrough.metadata as { parts: unknown }).parts).toEqual( + fullParts, + ); + }); +}); diff --git a/apps/server/test/integration/ai-chat-run-steps-migration.int-spec.ts b/apps/server/test/integration/ai-chat-run-steps-migration.int-spec.ts new file mode 100644 index 00000000..4c6e88a7 --- /dev/null +++ b/apps/server/test/integration/ai-chat-run-steps-migration.int-spec.ts @@ -0,0 +1,63 @@ +import { Kysely, sql } from 'kysely'; +import { + up, + down, +} from '../../src/database/migrations/20260708T120000-ai-chat-run-steps'; +import { getTestDb, destroyTestDb } from './db'; + +/** + * #492 migration up/down roundtrip on a LIVE Postgres. global-setup already + * migrated docmost_test to latest (so the table exists at start); this drives the + * migration's own down()/up() and asserts the table presence toggles, then leaves + * it PRESENT (up) so the shared test DB is intact for any spec that runs after. + */ +async function tableExists(db: Kysely): Promise { + const row = ( + await sql<{ t: string | null }>`select to_regclass('ai_chat_run_steps') as t`.execute( + db, + ) + ).rows[0]; + return row.t !== null; +} + +async function uniqueIndexExists(db: Kysely): Promise { + const row = ( + await sql<{ + t: string | null; + }>`select to_regclass('ai_chat_run_steps_message_step_uidx') as t`.execute(db) + ).rows[0]; + return row.t !== null; +} + +describe('20260708 ai_chat_run_steps migration roundtrip [integration]', () => { + let db: Kysely; + + beforeAll(() => { + db = getTestDb(); + }); + + afterAll(async () => { + // Belt-and-suspenders: guarantee the table is present for later specs even if + // an assertion threw mid-roundtrip. + if (!(await tableExists(db))) await up(db); + await destroyTestDb(); + }); + + it('down() drops the table+index and up() recreates them (idempotent)', async () => { + // Starts applied (global-setup migrated to latest). + expect(await tableExists(db)).toBe(true); + expect(await uniqueIndexExists(db)).toBe(true); + + await down(db); + expect(await tableExists(db)).toBe(false); + expect(await uniqueIndexExists(db)).toBe(false); + + await up(db); + expect(await tableExists(db)).toBe(true); + expect(await uniqueIndexExists(db)).toBe(true); + + // up() is idempotent (ifNotExists) — a second run is a harmless no-op. + await expect(up(db)).resolves.not.toThrow(); + expect(await tableExists(db)).toBe(true); + }); +}); From 8503ff1f3d70eb2df3cafe29dbc981e8b424f302 Mon Sep 17 00:00:00 2001 From: agent_coder Date: Sun, 12 Jul 2026 07:50:50 +0300 Subject: [PATCH 3/3] fix(ai-chat): hydrate crashed mid-run steps for model replay + cover write/read seams (#492) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit F1: the server model-replay loaded history via findAllByChat().map(rowToUiMessage) WITHOUT hydrating parts from ai_chat_run_steps. A HARD crash mid-run (SIGKILL/OOM) fires no terminal callback, so the assistant row stays parts:[] and its partial tool-calls/results/text (durable in the steps table) dropped out of the model's next-turn context. Hydrate needy assistant rows (role==='assistant' && !rowHasInlineParts) via findByMessageIds + hydrateAssistantParts before the replay map — mirroring the controller's withReconstructedParts exactly — guarded on the optional repo. Fix the now-false interrupt-resume comment. F2: add a service int-spec that drives the REAL onStep append-persist WRITE branch through AiChatService.stream with a real AiChatRunStepRepo injected, asserting the per-step rows' stepIndex + parts slice and the step-marker metadata match a single-row flush (catches an stepsPersisted-1 off-by-one). F3: add a controller int-spec that drives withReconstructedParts through getMessages WITH the repo present (a mid-run marker-only row + its step rows), asserting the reconstructed metadata.parts and workspace-scoping. F4: remove the dead countByMessage (zero prod callers; reconstructRunParts derives stepsPersisted inline) + its now-unused sql import and the redundant test assertion. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../src/core/ai-chat/ai-chat.service.ts | 32 +- .../repos/ai-chat/ai-chat-run-step.repo.ts | 16 - ...ai-chat-append-persist-service.int-spec.ts | 412 ++++++++++++++++++ .../ai-chat-hydrate-controller.int-spec.ts | 169 +++++++ .../ai-chat-run-step-repo.int-spec.ts | 1 - 5 files changed, 611 insertions(+), 19 deletions(-) create mode 100644 apps/server/test/integration/ai-chat-append-persist-service.int-spec.ts create mode 100644 apps/server/test/integration/ai-chat-hydrate-controller.int-spec.ts 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 0cd86cde..b080804f 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.ts @@ -1121,8 +1121,34 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy { chatId, workspace.id, ); + // #492: HYDRATE needy assistant rows from the steps table BEFORE the replay + // map. A #492 mid-run assistant row carries only a step marker + // (metadata.parts:[]); its real per-step parts live in `ai_chat_run_steps`. + // The graceful terminal callbacks (onFinish/onError/onAbort -> flushAssistant) + // assemble the full inline parts, so a normally-ended turn already has them. + // But a HARD crash mid-run (SIGKILL/OOM) fires NO terminal callback, so the + // row stays parts:[]; without this, rowToUiMessage falls back to an empty + // text part and the partial tool-calls/results/text — durable in the steps + // table — would DROP OUT of the model's replay context (regressing #183 + // step-granular durability for the model consumer). Mirrors the controller's + // withReconstructedParts EXACTLY (same needy predicate + hydration helper). + // Guarded on the optional repo: absent (positional test builds) degrades to + // the current behavior rather than crashing. + let replayHistory = oldHistory; + if (this.aiChatRunStepRepo) { + const needy = oldHistory.filter( + (r) => r.role === 'assistant' && !rowHasInlineParts(r), + ); + if (needy.length > 0) { + const stepsByMessage = await this.aiChatRunStepRepo.findByMessageIds( + needy.map((r) => r.id), + workspace.id, + ); + replayHistory = hydrateAssistantParts(oldHistory, stepsByMessage); + } + } const uiMessages: Array & { id: string }> = [ - ...oldHistory.map(rowToUiMessage), + ...replayHistory.map(rowToUiMessage), { id: 'pending-user', role: 'user', @@ -1161,7 +1187,9 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy { // hint — confirm it against the persisted history (the preceding assistant // turn must really be aborted/streaming) so a spoofed flag cannot inject the // interrupt note onto an ordinary turn. The partial output the model needs is - // already in `messages` (the aborted assistant row replays via findRecent). + // already in `messages`: a #492 mid-run row's per-step parts live only in the + // `ai_chat_run_steps` table and were hydrated into the replay history above, + // so the aborted assistant turn replays WITH its partial parts intact. // Append the new user turn (shape-only) so index -2 is the prior assistant. const interrupted = isInterruptResume( [...oldHistory, { role: 'user', status: null, metadata: null }], diff --git a/apps/server/src/database/repos/ai-chat/ai-chat-run-step.repo.ts b/apps/server/src/database/repos/ai-chat/ai-chat-run-step.repo.ts index 656186c6..8075c7ab 100644 --- a/apps/server/src/database/repos/ai-chat/ai-chat-run-step.repo.ts +++ b/apps/server/src/database/repos/ai-chat/ai-chat-run-step.repo.ts @@ -1,6 +1,5 @@ import { Injectable } from '@nestjs/common'; import { InjectKysely } from 'nestjs-kysely'; -import { sql } from 'kysely'; import { KyselyDB, KyselyTransaction } from '../../types/kysely.types'; import { dbOrTx } from '../../utils'; import { AiChatRunStep } from '@docmost/db/types/entity.types'; @@ -93,19 +92,4 @@ export class AiChatRunStepRepo { } return byMessage; } - - /** - * How many steps are persisted for a message (its step-marker floor). Exposed - * for the reconstruct contract (`reconstructRunParts → { parts, stepsPersisted }`) - * so a caller can align a resume attach without materializing every step's parts. - */ - async countByMessage(messageId: string, workspaceId: string): Promise { - const row = await this.db - .selectFrom('aiChatRunSteps') - .select(sql`count(*)::int`.as('n')) - .where('messageId', '=', messageId) - .where('workspaceId', '=', workspaceId) - .executeTakeFirst(); - return row?.n ?? 0; - } } diff --git a/apps/server/test/integration/ai-chat-append-persist-service.int-spec.ts b/apps/server/test/integration/ai-chat-append-persist-service.int-spec.ts new file mode 100644 index 00000000..0f007529 --- /dev/null +++ b/apps/server/test/integration/ai-chat-append-persist-service.int-spec.ts @@ -0,0 +1,412 @@ +import * as http from 'node:http'; +import { Kysely } from 'kysely'; +import { tool } from 'ai'; +import { z } from 'zod'; +import { MockLanguageModelV3, convertArrayToReadableStream } from 'ai/test'; +import { AiChatRepo } from '@docmost/db/repos/ai-chat/ai-chat.repo'; +import { AiChatMessageRepo } from '@docmost/db/repos/ai-chat/ai-chat-message.repo'; +import { AiChatRunStepRepo } from '@docmost/db/repos/ai-chat/ai-chat-run-step.repo'; +import { + AiChatService, + assembleStepParts, + assistantParts, + rowHasInlineParts, + stepMarkerMetadata, +} from 'src/core/ai-chat/ai-chat.service'; +import { + getTestDb, + destroyTestDb, + createWorkspace, + createUser, + createChat, + createMessage, +} from './db'; + +/** + * #492 append-persist — the REAL onStep WRITE path (F2) and the model-REPLAY + * hydration path (F1), driven through `AiChatService.stream` against a LIVE + * Postgres with a REAL `AiChatRunStepRepo` INJECTED. The existing append-persist + * int-specs hand-roll the insert+marker cycle via the repos directly and build + * the service with `aiChatRunStepRepo: undefined` (only the legacy-fallback branch + * is covered), so an off-by-one on `stepsPersisted-1`, a wrong `capturedSteps` + * slice, or a broken marker payload would pass all of them. These tests exercise + * the actual `updateStreaming` append-persist branch end to end. + * + * The seam is the injected `model` (a seeded `MockLanguageModelV3` from `ai/test`) + * plus a REAL Node `ServerResponse` as the hijacked socket — mirrors + * ai-chat-stream.int-spec.ts. + */ + +const sleep = (ms: number) => new Promise((r) => setTimeout(r, ms)); + +async function waitFor( + cond: () => Promise | boolean, + { timeoutMs = 15_000, stepMs = 25 } = {}, +): Promise { + const start = Date.now(); + while (Date.now() - start < timeoutMs) { + if (await cond()) return; + await sleep(stepMs); + } + throw new Error('waitFor: condition not met within timeout'); +} + +// A real Node ServerResponse wired to a live socket (identical helper to the +// stream int-spec) so the SDK's pipe/heartbeat writes behave as in prod. +function makeRealResponse(): Promise<{ + res: http.ServerResponse; + cleanup: () => Promise; +}> { + return new Promise((resolve) => { + const server = http.createServer((_req, res) => { + resolve({ + res, + cleanup: () => + new Promise((done) => { + try { + if (!res.writableEnded) res.end(); + } catch { + /* socket already gone */ + } + server.close(() => done()); + }), + }); + }); + server.listen(0, () => { + const port = (server.address() as any).port; + const creq = http.request({ port, method: 'GET' }, (cres) => { + cres.resume(); + }); + creq.on('error', () => undefined); + creq.end(); + }); + }); +} + +// Stream parts for a normal, successful single-step turn. +function successStream() { + return convertArrayToReadableStream([ + { type: 'stream-start', warnings: [] }, + { type: 'text-start', id: 't1' }, + { type: 'text-delta', id: 't1', delta: 'Hello' }, + { type: 'text-delta', id: 't1', delta: ' there' }, + { type: 'text-end', id: 't1' }, + { + type: 'finish', + finishReason: 'stop', + usage: { inputTokens: 10, outputTokens: 5, totalTokens: 15 }, + }, + ] as any); +} + +// A THREE-step turn: steps 0 and 1 each emit text + an `echo` tool call (the SDK +// runs the tool and continues); step 2 answers and stops. Three steps is +// deliberate: the LAST finished step's append-persist write races the terminal +// finalize (which writes the full inline parts anyway, so a lost last-step row is +// by design), but the NON-final steps 0 and 1 always drain to the steps table +// before finalize — so those are what the test asserts on deterministically. +function threeStepModel(): MockLanguageModelV3 { + let step = 0; + const toolStep = (i: number) => ({ + stream: convertArrayToReadableStream([ + { type: 'stream-start', warnings: [] }, + { type: 'text-start', id: `s${i}` }, + { type: 'text-delta', id: `s${i}`, delta: `step ${i} ` }, + { type: 'text-end', id: `s${i}` }, + { + type: 'tool-call', + toolCallId: `c${i}`, + toolName: 'echo', + input: JSON.stringify({ msg: `m${i}` }), + }, + { + type: 'finish', + finishReason: 'tool-calls', + usage: { inputTokens: 5, outputTokens: 3, totalTokens: 8 }, + }, + ] as any), + }); + return new MockLanguageModelV3({ + doStream: async () => { + const n = step++; + // Realistic inter-step latency. A real model spends seconds per step, so the + // fire-and-forget per-step write chain drains to the steps table BETWEEN + // steps; the mock otherwise collapses all steps into microseconds and the + // terminal finalize wins the race before any but the first step persists. + if (n > 0) await sleep(200); + if (n < 2) return toolStep(n); + return { + stream: convertArrayToReadableStream([ + { type: 'stream-start', warnings: [] }, + { type: 'text-start', id: 's2' }, + { type: 'text-delta', id: 's2', delta: 'final answer' }, + { type: 'text-end', id: 's2' }, + { + type: 'finish', + finishReason: 'stop', + usage: { inputTokens: 6, outputTokens: 4, totalTokens: 10 }, + }, + ] as any), + }; + }, + } as any); +} + +describe('#492 append-persist service paths [integration]', () => { + let db: Kysely; + let aiChatRepo: AiChatRepo; + let msgRepo: AiChatMessageRepo; + let stepRepo: AiChatRunStepRepo; + let workspaceId: string; + let userId: string; + + let closeCalls: number; + const mcpClients = { + toolsFor: async () => ({ + tools: {}, + clients: [ + { + close: async () => { + closeCalls += 1; + }, + }, + ], + outcomes: [], + instructions: [], + }), + }; + + // Build the service WITH a REAL AiChatRunStepRepo injected (the property under + // test) — unlike the legacy-fallback harness that passes it as undefined. + const echoTool = tool({ + description: 'echo the message back', + inputSchema: z.object({ msg: z.string() }), + execute: async ({ msg }) => ({ echoed: msg }), + }); + + function buildService(): AiChatService { + return new AiChatService( + { getChatModel: async () => null } as any, + aiChatRepo, + msgRepo, + {} as any, // aiChatPageSnapshotRepo + { resolve: async () => null } as any, // aiSettings + { forUser: async () => ({ echo: echoTool }) } as any, // tools + mcpClients as any, + {} as any, // aiAgentRoleRepo + {} as any, // pageRepo + {} as any, // pageAccess + { + isAiChatDeferredToolsEnabled: () => false, + isAiChatFinalStepLockdownEnabled: () => false, + } as any, // environment (deferred OFF -> all tools active every step) + undefined, // streamRegistry + undefined, // aiChatRunService + stepRepo, // #492 aiChatRunStepRepo — the append-persist backend + ); + } + + function userUiMessage(text: string) { + return { + id: `u-${Math.random()}`, + role: 'user', + parts: [{ type: 'text', text }], + }; + } + + async function runStream(opts: { + model: MockLanguageModelV3; + chatId: string; + body: any; + }): Promise { + closeCalls = 0; + const service = buildService(); + const { res, cleanup } = await makeRealResponse(); + try { + await service.stream({ + user: { id: userId, workspaceId } as any, + workspace: { id: workspaceId, name: 'WS' } as any, + sessionId: 'sess-1', + body: opts.body, + res: { raw: res } as any, + signal: new AbortController().signal, + model: opts.model as any, + role: null, + } as any); + await waitFor(async () => { + const rows = await msgRepo.findAllByChat(opts.chatId, workspaceId); + return rows.some( + (r) => + r.role === 'assistant' && + ['completed', 'error', 'aborted'].includes(r.status as string), + ); + }); + await waitFor(() => closeCalls > 0, { timeoutMs: 5_000 }); + } finally { + await cleanup(); + } + } + + beforeAll(async () => { + db = getTestDb(); + aiChatRepo = new AiChatRepo(db as any); + msgRepo = new AiChatMessageRepo(db as any); + stepRepo = new AiChatRunStepRepo(db as any); + workspaceId = (await createWorkspace(db)).id; + userId = (await createUser(db, workspaceId)).id; + }); + + afterAll(async () => { + await destroyTestDb(); + }); + + // --- F2: the real onStep append-persist WRITE branch ----------------------- + it('drives steps through the real onStep path: per-step rows + marker match a single-row flush', async () => { + const chatId = (await createChat(db, { workspaceId, creatorId: userId })).id; + const model = threeStepModel(); + + // Capture the mid-run step-marker UPDATEs the append-persist branch writes on + // the assistant row (a { parts: [], toolTraceVersion, stepsPersisted } patch). + const updateSpy = jest.spyOn(msgRepo, 'update'); + try { + await runStream({ + model, + chatId, + body: { chatId, messages: [userUiMessage('call the tool then answer')] }, + }); + + const rows = await msgRepo.findAllByChat(chatId, workspaceId); + const assistant = rows.find((r) => r.role === 'assistant')!; + expect(assistant).toBeDefined(); + expect(assistant.status).toBe('completed'); + // The turn finalizes with the FULL inline parts assembled by a single-row + // flush (assistantParts over every step) — the baseline the per-step slices + // must reproduce. + expect(rowHasInlineParts(assistant)).toBe(true); + const finalParts = (assistant.metadata as { parts: any[] }).parts; + + // The two NON-final finished steps each landed their own row, in stepIndex + // order. (The fire-and-forget write chain drains before the next step, so + // poll until both are on disk; the LAST step's write may lose the finalize + // race, which is by design — its parts are already in `finalParts`.) + await waitFor(async () => { + const s = await stepRepo.findByMessage(assistant.id, workspaceId); + return s.length >= 2; + }); + const steps = await stepRepo.findByMessage(assistant.id, workspaceId); + expect(steps[0].stepIndex).toBe(0); + expect(steps[1].stepIndex).toBe(1); + + // Each per-step row carries a NON-trivial slice: this step's text part + its + // paired tool part (guards a mutation that persists empty/whole-turn parts). + const s0 = steps[0].parts as any[]; + expect(s0).toContainEqual({ type: 'text', text: 'step 0 ' }); + expect(s0.some((p) => p.type === 'tool-echo')).toBe(true); + + // The per-step slices are EXACTLY the corresponding prefix of the single-row + // flush: assembleStepParts([step0, step1]) === finalParts[0 .. len0+len1]. + // This is what an off-by-one on `stepsPersisted-1` (a wrong `capturedSteps` + // slice) or a shifted stepIndex breaks — the prefix no longer aligns. + const prefixLen = + (steps[0].parts as any[]).length + (steps[1].parts as any[]).length; + expect(assembleStepParts([steps[0], steps[1]] as any)).toEqual( + finalParts.slice(0, prefixLen), + ); + + // The mid-run step markers advanced 1 -> 2 -> ... (the resume frontier), each + // a shape-stable empty-parts marker equal to a single-row flush's marker. + const markerCounts = updateSpy.mock.calls + .map((c) => (c[2] as any)?.metadata) + .filter( + (m) => + m && + Array.isArray(m.parts) && + m.parts.length === 0 && + typeof m.stepsPersisted === 'number', + ) + .map((m) => m.stepsPersisted); + // Monotonic from 1, covering at least the two non-final steps. + expect(markerCounts.slice(0, 2)).toEqual([1, 2]); + expect( + updateSpy.mock.calls + .map((c) => (c[2] as any)?.metadata) + .find((m) => m && m.stepsPersisted === 2), + ).toEqual(stepMarkerMetadata(2)); + } finally { + updateSpy.mockRestore(); + } + }, 60_000); + + // --- F1: model-REPLAY hydrates a hard-crashed mid-run turn from the steps table + it('replays a hard-crashed mid-run turn WITH its partial steps hydrated from the steps table', async () => { + const chatId = (await createChat(db, { workspaceId, creatorId: userId })).id; + + // Prior turn: a genuine user question... + await createMessage(db, { + workspaceId, + chatId, + userId, + role: 'user', + content: 'What is in the design doc?', + createdAt: new Date(Date.now() - 3000), + }); + // ...and an assistant row that a HARD crash (SIGKILL/OOM) left mid-run: only a + // step marker on the row (metadata.parts:[] , content:''), NO terminal + // callback ever fired, so its real parts live ONLY in ai_chat_run_steps. + const crashed = await createMessage(db, { + workspaceId, + chatId, + role: 'assistant', + content: '', + status: 'aborted', + metadata: stepMarkerMetadata(1), + createdAt: new Date(Date.now() - 2000), + }); + // The durable partial step: some reasoning text + a completed getPage tool + // call (input + output), exactly what #183 step-granular durability preserves. + await stepRepo.insertStep( + crashed.id, + workspaceId, + 0, + assistantParts( + [ + { + text: 'HYDRATED_PARTIAL_STEP the doc says', + toolCalls: [ + { toolCallId: 'g1', toolName: 'getPage', input: { id: 'p1' } }, + ], + toolResults: [ + { + toolCallId: 'g1', + toolName: 'getPage', + output: { id: 'p1', body: 'PARTIAL_TOOL_OUTPUT budget section' }, + }, + ], + } as any, + ], + '', + ), + ); + + // The NEXT turn: the model just answers. The service must REPLAY the crashed + // assistant turn with its partial parts hydrated from the steps table. + const model = new MockLanguageModelV3({ + doStream: async () => ({ stream: successStream() }), + } as any); + await runStream({ + model, + chatId, + body: { chatId, messages: [userUiMessage('Continue please')] }, + }); + + expect(model.doStreamCalls.length).toBeGreaterThan(0); + const prompt = JSON.stringify(model.doStreamCalls[0].prompt); + // The partial step's TEXT reached the model context (it would be an empty text + // part without hydration — rowToUiMessage falls back to `content:''`). + expect(prompt).toContain('HYDRATED_PARTIAL_STEP'); + // The partial TOOL RESULT survived too (durable in the steps table, replayed). + expect(prompt).toContain('PARTIAL_TOOL_OUTPUT'); + // The genuine prior user turn is present as well (sanity: real history replay). + expect(prompt).toContain('What is in the design doc?'); + }, 60_000); +}); diff --git a/apps/server/test/integration/ai-chat-hydrate-controller.int-spec.ts b/apps/server/test/integration/ai-chat-hydrate-controller.int-spec.ts new file mode 100644 index 00000000..d48d9d1c --- /dev/null +++ b/apps/server/test/integration/ai-chat-hydrate-controller.int-spec.ts @@ -0,0 +1,169 @@ +import { Kysely } from 'kysely'; +import { AiChatController } from 'src/core/ai-chat/ai-chat.controller'; +import { + assembleStepParts, + assistantParts, + stepMarkerMetadata, +} from 'src/core/ai-chat/ai-chat.service'; +import { AiChatRepo } from '@docmost/db/repos/ai-chat/ai-chat.repo'; +import { AiChatMessageRepo } from '@docmost/db/repos/ai-chat/ai-chat-message.repo'; +import { AiChatRunStepRepo } from '@docmost/db/repos/ai-chat/ai-chat-run-step.repo'; +import type { User, Workspace } from '@docmost/db/types/entity.types'; +import { + getTestDb, + destroyTestDb, + createWorkspace, + createUser, + createChat, + createMessage, +} from './db'; + +/** + * #492 controller hydration (crash-before-finalize RESUME) on a LIVE Postgres. + * `AiChatController.withReconstructedParts` is wired into getMessages/delta/export/ + * run, but `aiChatRunStepRepo` is OPTIONAL and every controller unit spec passes it + * as `undefined`, so the hydration branch early-returns and NEVER executes in those + * tests. This drives the real read path — a mid-run streaming row (marker only, + * empty inline parts) PLUS its `ai_chat_run_steps` rows — through getMessages WITH + * the repo present, exercising the `role==='assistant' && !rowHasInlineParts` + * needy predicate, the workspace-scoped batch step fetch, and the endpoint binding. + */ +describe('#492 controller hydration read path [integration]', () => { + let db: Kysely; + let aiChatRepo: AiChatRepo; + let msgRepo: AiChatMessageRepo; + let stepRepo: AiChatRunStepRepo; + let workspaceId: string; + let otherWorkspaceId: string; + let userId: string; + + // Build the controller WITH a real AiChatRunStepRepo injected (position 9), the + // seam the unit specs leave undefined. Only the read-path deps are real. + function buildController(): AiChatController { + return new AiChatController( + {} as any, // aiChatService + {} as any, // aiChatRunService + aiChatRepo, + msgRepo, + {} as any, // aiTranscription + {} as any, // pageRepo + undefined, // streamRegistry + undefined, // environment + stepRepo, // #492 aiChatRunStepRepo + ); + } + + beforeAll(async () => { + db = getTestDb(); + aiChatRepo = new AiChatRepo(db as any); + msgRepo = new AiChatMessageRepo(db as any); + stepRepo = new AiChatRunStepRepo(db as any); + workspaceId = (await createWorkspace(db)).id; + otherWorkspaceId = (await createWorkspace(db)).id; + userId = (await createUser(db, workspaceId)).id; + }); + + afterAll(async () => { + await destroyTestDb(); + }); + + it('getMessages reconstructs a mid-run row from the steps table (finished rows untouched)', async () => { + const chatId = ( + await createChat(db, { workspaceId, creatorId: userId }) + ).id; + const user = { id: userId } as User; + const workspace = { id: workspaceId } as Workspace; + + // A prior FINISHED assistant row that already carries inline parts — the needy + // predicate must SKIP it (no step fetch), returned untouched. + const finishedParts = assistantParts( + [{ text: 'done earlier', toolCalls: [], toolResults: [] } as any], + '', + ); + await createMessage(db, { + workspaceId, + chatId, + role: 'assistant', + content: 'done earlier', + status: 'completed', + metadata: { parts: finishedParts, toolTraceVersion: 2, stepsPersisted: 1 }, + createdAt: new Date(Date.now() - 3000), + }); + + // The mid-run row a crash-before-finalize left behind: a step marker only + // (parts:[] , content:''), status 'streaming'. Its real parts live ONLY in the + // steps table. + const midRun = await createMessage(db, { + workspaceId, + chatId, + role: 'assistant', + content: '', + status: 'streaming', + metadata: stepMarkerMetadata(2), + createdAt: new Date(Date.now() - 1000), + }); + + const step0 = assistantParts( + [ + { + text: 'reasoning about the page', + toolCalls: [ + { toolCallId: 'g1', toolName: 'getPage', input: { id: 'p1' } }, + ], + toolResults: [ + { toolCallId: 'g1', toolName: 'getPage', output: { id: 'p1', body: 'B' } }, + ], + } as any, + ], + '', + ); + const step1 = assistantParts( + [{ text: 'partial synthesis so far', toolCalls: [], toolResults: [] } as any], + '', + ); + await stepRepo.insertStep(midRun.id, workspaceId, 0, step0); + await stepRepo.insertStep(midRun.id, workspaceId, 1, step1); + + // Workspace-scoping guard: a step row for the SAME message id under a DIFFERENT + // workspace must NEVER leak into this workspace's reconstruction. + await stepRepo.insertStep(midRun.id, otherWorkspaceId, 99, [ + { type: 'text', text: 'FOREIGN_WORKSPACE_LEAK' }, + ]); + + const res = await buildController().getMessages( + { chatId } as any, + { limit: 50 } as any, + user, + workspace, + ); + + const items = res.items as any[]; + const finished = items.find((r) => r.status === 'completed'); + const reconstructed = items.find((r) => r.id === midRun.id); + + // The finished row passed through with its inline parts unchanged. + expect(finished.metadata.parts).toEqual(finishedParts); + + // The mid-run row's parts were reconstructed from the two step rows, in order, + // exactly as assembleStepParts concatenates them — the client seed sees the + // persisted progress with no change to itself. + const expected = assembleStepParts([ + { stepIndex: 0, parts: step0 }, + { stepIndex: 1, parts: step1 }, + ] as any); + expect(reconstructed.metadata.parts).toEqual(expected); + // The foreign-workspace step row did NOT leak in. + expect(JSON.stringify(reconstructed.metadata.parts)).not.toContain( + 'FOREIGN_WORKSPACE_LEAK', + ); + // Sanity: reconstruction produced real content (text + the paired tool part + + // the second step's text), not an empty fallback. + expect(reconstructed.metadata.parts).toContainEqual({ + type: 'text', + text: 'reasoning about the page', + }); + expect( + (reconstructed.metadata.parts as any[]).some((p) => p.type === 'tool-getPage'), + ).toBe(true); + }, 60_000); +}); diff --git a/apps/server/test/integration/ai-chat-run-step-repo.int-spec.ts b/apps/server/test/integration/ai-chat-run-step-repo.int-spec.ts index 96b71a3a..78211f4a 100644 --- a/apps/server/test/integration/ai-chat-run-step-repo.int-spec.ts +++ b/apps/server/test/integration/ai-chat-run-step-repo.int-spec.ts @@ -98,7 +98,6 @@ describe('AiChatRunStepRepo + reconstruct contract [integration]', () => { const steps = await stepRepo.findByMessage(row.id, workspaceId); expect(steps.map((s) => s.stepIndex)).toEqual([0, 1]); - expect(await stepRepo.countByMessage(row.id, workspaceId)).toBe(2); // Batch fetch groups by message id in step order. const map = await stepRepo.findByMessageIds([row.id], workspaceId);