` 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+$/, "")}
+
+ ),
+ )}
+
+ );
+}
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 7ec93545..b080804f 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';
@@ -518,6 +519,12 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
// constructions compile unchanged; Nest always injects the real singleton, so
// reconcile sees the SAME in-memory active/zombie maps the runner mutates.
private readonly aiChatRunService?: AiChatRunService,
+ // #492 append-persist: per-step INSERT into the lightweight steps table (the
+ // O(Σ steps) replacement for the O(n²) full-row `metadata.parts` rewrite).
+ // OPTIONAL so existing positional constructions (int-specs) compile unchanged;
+ // Nest injects the real singleton. When ABSENT the per-step path falls back to
+ // the pre-#492 full-row flush (no regression, only no WAL win).
+ private readonly aiChatRunStepRepo?: AiChatRunStepRepo,
) {}
// #487: periodic reconcile timer (single-process phase 1). Started in
@@ -1114,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',
@@ -1154,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 }],
@@ -1559,17 +1594,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(
@@ -2749,6 +2824,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 8d16e5f9..22b4962c 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';
@@ -125,6 +126,7 @@ import { firstSqlToken } from '../integrations/metrics/metrics.constants';
AiChatRepo,
AiChatMessageRepo,
AiChatRunRepo,
+ AiChatRunStepRepo,
AiChatPageSnapshotRepo,
AiProviderCredentialsRepo,
AiMcpServerRepo,
@@ -161,6 +163,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..8075c7ab
--- /dev/null
+++ b/apps/server/src/database/repos/ai-chat/ai-chat-run-step.repo.ts
@@ -0,0 +1,95 @@
+import { Injectable } from '@nestjs/common';
+import { InjectKysely } from 'nestjs-kysely';
+import { KyselyDB, KyselyTransaction } from '../../types/kysely.types';
+import { dbOrTx } from '../../utils';
+import { AiChatRunStep } from '@docmost/db/types/entity.types';
+
+/**
+ * Append-only per-step persistence for an assistant turn (#492). Each finished
+ * agent step's UI `parts` (its text part + a `tool-*` part per call, WITH the
+ * tool output) is INSERTed as its own lightweight row the moment the step ends —
+ * instead of REWRITING the assistant row's growing `metadata.parts` jsonb on every
+ * `onStepFinish` (a Postgres jsonb UPDATE rewrites the whole TOASTed row version
+ * under MVCC, so that was O(n²) WAL/dead-tuple churn per turn).
+ *
+ * The full `metadata.parts` on the message row is assembled ONCE at finalize;
+ * mid-run, a resuming client's seed is rebuilt from these rows in `stepIndex`
+ * order (see `assembleStepParts` / the reconstruct seam in ai-chat.service.ts).
+ * Every method is workspace-scoped as defense-in-depth.
+ */
+@Injectable()
+export class AiChatRunStepRepo {
+ constructor(@InjectKysely() private readonly db: KyselyDB) {}
+
+ /**
+ * Append one finished step's parts. Idempotent: a retried persist of the SAME
+ * (message, stepIndex) is a no-op via ON CONFLICT DO NOTHING — the per-step
+ * writes are fired fire-and-forget + serialized, and a duplicate must never
+ * throw into the stream or double the parts. Returns whether a NEW row landed
+ * (false = the step was already persisted).
+ */
+ async insertStep(
+ messageId: string,
+ workspaceId: string,
+ stepIndex: number,
+ parts: unknown,
+ trx?: KyselyTransaction,
+ ): Promise {
+ 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;
+ }
+}
diff --git a/apps/server/src/database/types/db.d.ts b/apps/server/src/database/types/db.d.ts
index 889adcf5..5530db56 100644
--- a/apps/server/src/database/types/db.d.ts
+++ b/apps/server/src/database/types/db.d.ts
@@ -692,6 +692,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
@@ -729,6 +745,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-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-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-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
new file mode 100644
index 00000000..78211f4a
--- /dev/null
+++ b/apps/server/test/integration/ai-chat-run-step-repo.int-spec.ts
@@ -0,0 +1,163 @@
+import { randomBytes } from 'crypto';
+import { Kysely } from 'kysely';
+import { AiChatRunStepRepo } from '@docmost/db/repos/ai-chat/ai-chat-run-step.repo';
+import { AiChatMessageRepo } from '@docmost/db/repos/ai-chat/ai-chat-message.repo';
+import {
+ assistantParts,
+ reconstructRunParts,
+ hydrateAssistantParts,
+ stepMarkerMetadata,
+ rowHasInlineParts,
+} from '../../src/core/ai-chat/ai-chat.service';
+import {
+ getTestDb,
+ destroyTestDb,
+ createWorkspace,
+ createUser,
+ createChat,
+} from './db';
+
+/**
+ * #492 append-persist — the reconstruct CONTRACT on a live Postgres. Proves that a
+ * turn persisted the NEW way (per-step rows in `ai_chat_run_steps`, only a step
+ * marker on the message row) reconstructs to the SAME UI parts as a turn persisted
+ * the OLD way (full `metadata.parts` inline on the row, no step rows) — so the
+ * era-switch is invisible to attach / delta-poll / export. Real repos + real jsonb
+ * roundtrip, not a mock (a mock cannot prove the parts survive the jsonb column
+ * byte-identical).
+ */
+type Step = {
+ text: string;
+ toolCalls: Array<{ toolCallId: string; toolName: string; input: unknown }>;
+ toolResults: Array<{ toolCallId: string; toolName: string; output: unknown }>;
+};
+
+// A realistic step: some text + a getPage tool call whose ~100 KB body is
+// INCOMPRESSIBLE random base64 (a 'x'.repeat filler would TOAST away and hide the
+// real bytes). Under MAX_TOOL_OUTPUT_BYTES (200 KB) it is stored uncompacted.
+function makeStep(i: number, outputBytes = 4_000): Step {
+ const body = randomBytes(Math.ceil(outputBytes * 0.75)).toString('base64');
+ return {
+ text: `step ${i} text`,
+ toolCalls: [
+ { toolCallId: `c${i}`, toolName: 'getPage', input: { id: `p${i}` } },
+ ],
+ toolResults: [
+ {
+ toolCallId: `c${i}`,
+ toolName: 'getPage',
+ output: { id: `p${i}`, title: `Page ${i}`, body },
+ },
+ ],
+ };
+}
+
+describe('AiChatRunStepRepo + reconstruct contract [integration]', () => {
+ let db: Kysely;
+ let stepRepo: AiChatRunStepRepo;
+ let msgRepo: AiChatMessageRepo;
+ let workspaceId: string;
+ let userId: string;
+ let chatId: string;
+
+ beforeAll(async () => {
+ db = getTestDb();
+ stepRepo = new AiChatRunStepRepo(db as any);
+ msgRepo = new AiChatMessageRepo(db as any);
+ workspaceId = (await createWorkspace(db)).id;
+ userId = (await createUser(db, workspaceId)).id;
+ chatId = (await createChat(db, { workspaceId, creatorId: userId })).id;
+ });
+
+ afterAll(async () => {
+ await destroyTestDb();
+ });
+
+ const seedRow = (metadata: unknown, status: string) =>
+ msgRepo.insert({
+ chatId,
+ workspaceId,
+ userId,
+ role: 'assistant',
+ content: '',
+ status,
+ metadata: metadata as never,
+ });
+
+ it('insertStep is idempotent per (message, stepIndex) and reads back in order', async () => {
+ const row = await seedRow(stepMarkerMetadata(0), 'streaming');
+ const parts0 = assistantParts([makeStep(0)], '');
+ const parts1 = assistantParts([makeStep(1)], '');
+
+ expect(await stepRepo.insertStep(row.id, workspaceId, 0, parts0)).toBe(true);
+ expect(await stepRepo.insertStep(row.id, workspaceId, 1, parts1)).toBe(true);
+ // A retried persist of the SAME step is a no-op (ON CONFLICT DO NOTHING).
+ expect(await stepRepo.insertStep(row.id, workspaceId, 0, parts0)).toBe(
+ false,
+ );
+
+ const steps = await stepRepo.findByMessage(row.id, workspaceId);
+ expect(steps.map((s) => s.stepIndex)).toEqual([0, 1]);
+
+ // Batch fetch groups by message id in step order.
+ const map = await stepRepo.findByMessageIds([row.id], workspaceId);
+ expect(map.get(row.id)!.map((s) => s.stepIndex)).toEqual([0, 1]);
+ });
+
+ it('a NEW-style (step-table) run reconstructs identically to an OLD-style (inline) run', async () => {
+ const steps = [makeStep(10), makeStep(11)];
+ // The inline parts the OLD full-row flush would have written.
+ const fullParts = assistantParts(steps, '');
+
+ // OLD-style record: full parts inline on the row, NO step rows.
+ const oldRow = await seedRow(
+ { parts: fullParts, toolTraceVersion: 2, stepsPersisted: 2 },
+ 'completed',
+ );
+
+ // NEW-style record: only a step marker on the row + per-step rows.
+ const newRow = await seedRow(stepMarkerMetadata(2), 'streaming');
+ for (let i = 0; i < steps.length; i++) {
+ await stepRepo.insertStep(
+ newRow.id,
+ workspaceId,
+ i,
+ assistantParts([steps[i]], ''),
+ );
+ }
+
+ // Re-read both from the DB (proves the jsonb roundtrip).
+ const oldFetched = await msgRepo.findById(oldRow.id, workspaceId);
+ const newFetched = await msgRepo.findById(newRow.id, workspaceId);
+ const oldSteps = await stepRepo.findByMessage(oldRow.id, workspaceId);
+ const newSteps = await stepRepo.findByMessage(newRow.id, workspaceId);
+
+ // The discriminator: the old row carries inline parts, the new one does not.
+ expect(rowHasInlineParts(oldFetched!)).toBe(true);
+ expect(rowHasInlineParts(newFetched!)).toBe(false);
+ expect(oldSteps).toHaveLength(0);
+ expect(newSteps).toHaveLength(2);
+
+ const oldRecon = reconstructRunParts(oldFetched!, oldSteps);
+ const newRecon = reconstructRunParts(newFetched!, newSteps);
+
+ // Both reconstruct to the SAME parts + step count — the era is invisible.
+ expect(newRecon.parts).toEqual(fullParts);
+ expect(oldRecon.parts).toEqual(fullParts);
+ expect(newRecon.parts).toEqual(oldRecon.parts);
+ expect(newRecon.stepsPersisted).toBe(2);
+ expect(oldRecon.stepsPersisted).toBe(2);
+
+ // hydrateAssistantParts fills the new row's metadata.parts to match the old
+ // row's inline parts — so a consumer reading `metadata.parts` off the raw row
+ // (the client seed/poll, export) is unchanged across the era.
+ const map = await stepRepo.findByMessageIds([newRow.id], workspaceId);
+ const [hydrated] = hydrateAssistantParts([newFetched!], map);
+ expect((hydrated.metadata as { parts: unknown }).parts).toEqual(fullParts);
+ // A row that already has inline parts passes through untouched (same ref-shape).
+ const [oldPassThrough] = hydrateAssistantParts([oldFetched!], map);
+ expect((oldPassThrough.metadata as { parts: unknown }).parts).toEqual(
+ fullParts,
+ );
+ });
+});
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);
+ });
+});