From c96fafc4ad2c37abc8f3ef765a3c80012cd11230 Mon Sep 17 00:00:00 2001 From: agent_coder Date: Sun, 12 Jul 2026 05:02:25 +0300 Subject: [PATCH] =?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); + }); +});