fix(ai-chat): hydrate crashed mid-run steps for model replay + cover write/read seams (#492)

F1: the server model-replay loaded history via findAllByChat().map(rowToUiMessage)
WITHOUT hydrating parts from ai_chat_run_steps. A HARD crash mid-run (SIGKILL/OOM)
fires no terminal callback, so the assistant row stays parts:[] and its partial
tool-calls/results/text (durable in the steps table) dropped out of the model's
next-turn context. Hydrate needy assistant rows (role==='assistant' &&
!rowHasInlineParts) via findByMessageIds + hydrateAssistantParts before the replay
map — mirroring the controller's withReconstructedParts exactly — guarded on the
optional repo. Fix the now-false interrupt-resume comment.

F2: add a service int-spec that drives the REAL onStep append-persist WRITE branch
through AiChatService.stream with a real AiChatRunStepRepo injected, asserting the
per-step rows' stepIndex + parts slice and the step-marker metadata match a
single-row flush (catches an stepsPersisted-1 off-by-one).

F3: add a controller int-spec that drives withReconstructedParts through getMessages
WITH the repo present (a mid-run marker-only row + its step rows), asserting the
reconstructed metadata.parts and workspace-scoping.

F4: remove the dead countByMessage (zero prod callers; reconstructRunParts derives
stepsPersisted inline) + its now-unused sql import and the redundant test assertion.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

apps/server/src/core/ai-chat/ai-chat.service.ts

@@ -1121,8 +1121,34 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
         chatId,
         workspace.id,
       );
+      // #492: HYDRATE needy assistant rows from the steps table BEFORE the replay
+      // map. A #492 mid-run assistant row carries only a step marker
+      // (metadata.parts:[]); its real per-step parts live in `ai_chat_run_steps`.
+      // The graceful terminal callbacks (onFinish/onError/onAbort -> flushAssistant)
+      // assemble the full inline parts, so a normally-ended turn already has them.
+      // But a HARD crash mid-run (SIGKILL/OOM) fires NO terminal callback, so the
+      // row stays parts:[]; without this, rowToUiMessage falls back to an empty
+      // text part and the partial tool-calls/results/text — durable in the steps
+      // table — would DROP OUT of the model's replay context (regressing #183
+      // step-granular durability for the model consumer). Mirrors the controller's
+      // withReconstructedParts EXACTLY (same needy predicate + hydration helper).
+      // Guarded on the optional repo: absent (positional test builds) degrades to
+      // the current behavior rather than crashing.
+      let replayHistory = oldHistory;
+      if (this.aiChatRunStepRepo) {
+        const needy = oldHistory.filter(
+          (r) => r.role === 'assistant' && !rowHasInlineParts(r),
+        );
+        if (needy.length > 0) {
+          const stepsByMessage = await this.aiChatRunStepRepo.findByMessageIds(
+            needy.map((r) => r.id),
+            workspace.id,
+          );
+          replayHistory = hydrateAssistantParts(oldHistory, stepsByMessage);
+        }
+      }
       const uiMessages: Array<Omit<UIMessage, 'id'> & { id: string }> = [
-        ...oldHistory.map(rowToUiMessage),
+        ...replayHistory.map(rowToUiMessage),
         {
           id: 'pending-user',
           role: 'user',
@@ -1161,7 +1187,9 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
       // hint — confirm it against the persisted history (the preceding assistant
       // turn must really be aborted/streaming) so a spoofed flag cannot inject the
       // interrupt note onto an ordinary turn. The partial output the model needs is
-      // already in `messages` (the aborted assistant row replays via findRecent).
+      // already in `messages`: a #492 mid-run row's per-step parts live only in the
+      // `ai_chat_run_steps` table and were hydrated into the replay history above,
+      // so the aborted assistant turn replays WITH its partial parts intact.
       // Append the new user turn (shape-only) so index -2 is the prior assistant.
       const interrupted = isInterruptResume(
         [...oldHistory, { role: 'user', status: null, metadata: null }],

apps/server/src/database/repos/ai-chat/ai-chat-run-step.repo.ts

@@ -1,6 +1,5 @@
 import { Injectable } from '@nestjs/common';
 import { InjectKysely } from 'nestjs-kysely';
-import { sql } from 'kysely';
 import { KyselyDB, KyselyTransaction } from '../../types/kysely.types';
 import { dbOrTx } from '../../utils';
 import { AiChatRunStep } from '@docmost/db/types/entity.types';
@@ -93,19 +92,4 @@ export class AiChatRunStepRepo {
     }
     return byMessage;
   }
-
-  /**
-   * How many steps are persisted for a message (its step-marker floor). Exposed
-   * for the reconstruct contract (`reconstructRunParts → { parts, stepsPersisted }`)
-   * so a caller can align a resume attach without materializing every step's parts.
-   */
-  async countByMessage(messageId: string, workspaceId: string): Promise<number> {
-    const row = await this.db
-      .selectFrom('aiChatRunSteps')
-      .select(sql<number>`count(*)::int`.as('n'))
-      .where('messageId', '=', messageId)
-      .where('workspaceId', '=', workspaceId)
-      .executeTakeFirst();
-    return row?.n ?? 0;
-  }
 }