From 800812a31d1b3d299352caaae07f3e3a61fbfd75 Mon Sep 17 00:00:00 2001 From: Yevanchen Date: Thu, 30 Jul 2026 15:26:12 +0800 Subject: [PATCH] fix(runtime): persist final assistant log snapshots --- .../driver-instance/event-persistence.ts | 49 +++++++++++++- .../runtime-final-output-ingestion.test.ts | 67 +++++++++++++++++++ 2 files changed, 114 insertions(+), 2 deletions(-) diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/event-persistence.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/event-persistence.ts index 1bb446e4..78e37990 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/event-persistence.ts +++ b/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); diff --git a/apps/api/tests/runtime-final-output-ingestion.test.ts b/apps/api/tests/runtime-final-output-ingestion.test.ts index 62bff293..08bb99a6 100644 --- a/apps/api/tests/runtime-final-output-ingestion.test.ts +++ b/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(); + const events = [ + ...fragmentTexts.flatMap((text, index) => + messageEvents({ + messageId: createPlatformId(), + 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(); + 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);