fix(runtime): persist final assistant log snapshots

apps/api/src/modules/runtime/infrastructure/driver-instance/event-persistence.ts

@@ -1,4 +1,4 @@
-import { sessionsTable } from "@mosoo/db";
+import { sessionEventsTable, sessionsTable } from "@mosoo/db";
 import type { DriverInstanceId } from "@mosoo/id";
 import { and, eq, isNull } from "drizzle-orm";
 
@@ -9,6 +9,7 @@ import {
 import type { ApiBindings } from "../../../../platform/cloudflare/worker-types";
 import { getAppDatabase } from "../../../../platform/db/drizzle";
 import { currentTimestampMs } from "../../../../time";
+import { createSessionRuntimeEvent } from "../../../sessions/application/session-event-write.service";
 import { upsertSessionModelCallUsage } from "../../../sessions/application/session-model-call.service";
 import { persistSessionRuntimeEvents } from "../../../sessions/infrastructure/session-runtime-event-store.repository";
 import { setSessionRunStatus } from "../session-runs/session-run-store.repository";
@@ -219,6 +220,7 @@ export async function persistProjectedRuntimeDriverEvents(
     completedTransition === undefined
       ? projection.runtimeEvents
       : projection.runtimeEvents.filter((record) => record.event.kind === "run.completed");
+  const finalAssistantRuntimeEvents: typeof terminalRuntimeEvents = [];
   const persistedSourceEventIds: string[] = [];
 
   if (preCompletionRuntimeEvents.length > 0) {
@@ -280,10 +282,53 @@ export async function persistProjectedRuntimeDriverEvents(
       sessionRunId: link.sessionRunId,
       state: nextLiveState,
     });
+
+    const terminalRuntimeEvent = terminalRuntimeEvents[0];
+    const finalSnapshotAlreadyPersisted =
+      (await getAppDatabase(database)
+        .select({ id: sessionEventsTable.id })
+        .from(sessionEventsTable)
+        .where(
+          and(
+            eq(sessionEventsTable.sessionId, link.sessionId),
+            eq(sessionEventsTable.runId, link.sessionRunId),
+            eq(sessionEventsTable.eventType, "message.added"),
+            eq(sessionEventsTable.processType, "agent.message.delta"),
+            eq(sessionEventsTable.contentText, projection.finalAssistantMessage.text),
+          ),
+        )
+        .limit(1)
+        .get()) !== undefined;
+
+    if (terminalRuntimeEvent !== undefined && !finalSnapshotAlreadyPersisted) {
+      const sourceEventId = `session-run:${link.sessionRunId}:final-assistant`;
+      finalAssistantRuntimeEvents.push({
+        event: createSessionRuntimeEvent({
+          actor: terminalRuntimeEvent.event.actor,
+          kind: "message.added",
+          ...(terminalRuntimeEvent.occurredAt === null
+            ? {}
+            : { occurredAtMs: terminalRuntimeEvent.occurredAt }),
+          origin: terminalRuntimeEvent.event.origin,
+          payload: {
+            content: projection.finalAssistantMessage.text,
+            messageId: projection.finalAssistantMessage.id,
+            role: "agent",
+          },
+          runId: link.sessionRunId,
+          sessionId: link.sessionId,
+          sourceEventId,
+          traceId: terminalRuntimeEvent.event.traceId ?? link.traceId,
+          visibility: terminalRuntimeEvent.event.visibility,
+        }),
+        occurredAt: terminalRuntimeEvent.occurredAt,
+        sourceEventId,
+      });
+    }
   }
 
   const persistedTerminalEvents = await persistSessionRuntimeEvents(database, {
-    records: terminalRuntimeEvents,
+    records: [...finalAssistantRuntimeEvents, ...terminalRuntimeEvents],
     sessionId: link.sessionId,
   });
   persistedSourceEventIds.push(...persistedTerminalEvents.persistedSourceEventIds);

apps/api/tests/runtime-final-output-ingestion.test.ts

@@ -25,6 +25,8 @@ import { DriverInstanceRpcEventIngestionController } from "../src/modules/runtim
 import { RuntimeSessionViewCache } from "../src/modules/runtime/infrastructure/driver-instance/runtime-session-view-cache";
 import { recordDriverInstanceCompletion } from "../src/modules/runtime/infrastructure/driver-instance/terminal-driver-events";
 import { loadSessionViewerState } from "../src/modules/sessions/application/session-live-state.service";
+import { createSessionProcessEventsFromSessionEventRows } from "../src/modules/sessions/application/session-process-events.service";
+import type { SessionEventProcessRow } from "../src/modules/sessions/application/session-process-events.service";
 import type { ApiBindings } from "../src/platform/cloudflare/worker-types";
 import {
   createPublicHttpContractDatabase,
@@ -296,6 +298,71 @@ async function pushFreshController(
 }
 
 describe("runtime final output ingestion", () => {
+  test.each([
+    ["omits", false],
+    ["provides", true],
+  ] as const)(
+    "persists one final assistant snapshot when the driver %s it",
+    async (_driverBehavior, driverProvidesSnapshot) => {
+      const database = await createPublicHttpContractDatabase();
+      await insertRuntimeFixture(database);
+      const bindings = createPublicHttpTestBindings(database) as ApiBindings;
+      const finalText = "The final answer.";
+      const fragmentTexts = ["The ", "final ", "answer."];
+      const finalMessageId = createPlatformId<SessionMessageId>();
+      const events = [
+        ...fragmentTexts.flatMap((text, index) =>
+          messageEvents({
+            messageId: createPlatformId<SessionMessageId>(),
+            sourcePrefix: `fractured:${index + 1}`,
+            text,
+          }),
+        ),
+        ...(driverProvidesSnapshot
+          ? [
+              runtimeEvent({
+                kind: "message.added",
+                payload: { content: finalText, messageId: finalMessageId, role: "agent" },
+                sourceEventId: "fractured:final-snapshot",
+              }),
+            ]
+          : []),
+        runtimeEvent({
+          kind: "run.completed",
+          payload: {
+            finalMessageId,
+            finalMessageText: finalText,
+            stopReason: "end_turn",
+          },
+          sourceEventId: "fractured:run-completed",
+        }),
+      ];
+
+      await pushFreshController(bindings, events);
+
+      const rows = await database
+        .prepare(
+          `SELECT content_text, ended_at, event_type, id, occurred_at, process_status,
+                process_type, run_id, seq, tokens
+         FROM session_event
+         WHERE session_id = ? AND run_id = ?
+         ORDER BY seq`,
+        )
+        .bind(SESSION_ID, RUN_ID)
+        .all<SessionEventProcessRow>();
+      const assistantMessages = createSessionProcessEventsFromSessionEventRows(rows.results).filter(
+        (event) => event.type === "agent.message.delta",
+      );
+
+      expect(
+        rows.results
+          .filter((row) => row.event_type === "message.added")
+          .map((row) => row.content_text),
+      ).toEqual([finalText]);
+      expect(assistantMessages.map((event) => event.content)).toEqual([finalText]);
+    },
+  );
+
   test("preserves a long final snapshot across hibernation, terminal failure, and replay", async () => {
     const database = await createPublicHttpContractDatabase();
     await insertRuntimeFixture(database);