feat: persist agent messages to transcript

packages/core/src/__tests__/agent-session.test.ts

@@ -3,68 +3,121 @@ import { mkdtemp, mkdir, writeFile, rm } from "node:fs/promises";
 import { join } from "node:path";
 import { tmpdir } from "node:os";
 
-const { agentInstances } = vi.hoisted(() => ({
+const EMPTY_USAGE = {
+  input: 0,
+  output: 0,
+  cacheRead: 0,
+  cacheWrite: 0,
+  totalTokens: 0,
+  cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+};
+
+const { agentInstances, streamCalls } = vi.hoisted(() => ({
   agentInstances: [] as any[],
+  streamCalls: [] as Array<{ model: any; context: any }>,
 }));
 
 vi.mock("@mariozechner/pi-agent-core", async () => {
   const actual = await vi.importActual<any>("@mariozechner/pi-agent-core");
-  class FakeAgent {
-    state: any;
-    transformContext: any;
-    streamFn: any;
-    getApiKey: any;
+  class SpyAgent extends actual.Agent {
     constructor(options: any) {
-      this.state = {
-        model: options.initialState?.model,
-        systemPrompt: options.initialState?.systemPrompt,
-        tools: options.initialState?.tools ?? [],
-        messages: [],
-      };
-      this.transformContext = options.transformContext;
-      this.streamFn = options.streamFn;
-      this.getApiKey = options.getApiKey;
+      super(options);
       agentInstances.push(this);
     }
-    subscribe() {
-      return () => {};
-    }
-    async prompt(userMessage: string) {
-      const now = Date.now();
-      this.state.messages.push({ role: "user", content: userMessage, timestamp: now });
-      this.state.messages.push({
-        role: "assistant",
-        content: [{ type: "text", text: "ok" }],
-        api: "anthropic-messages",
-        provider: "anthropic",
-        model: "fake",
-        usage: {
-          input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0,
-          cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
-        },
-        stopReason: "stop",
-        timestamp: now + 1,
-      });
-    }
   }
-  return { ...actual, Agent: FakeAgent };
+  return { ...actual, Agent: SpyAgent };
 });
 
 vi.mock("@mariozechner/pi-ai", async () => {
   const actual = await vi.importActual<any>("@mariozechner/pi-ai");
+
+  function clone(value: unknown): unknown {
+    return JSON.parse(JSON.stringify(value));
+  }
+
+  function textFromContent(content: unknown): string {
+    if (typeof content === "string") return content;
+    if (!Array.isArray(content)) return "";
+    return content
+      .filter((block: any) => block?.type === "text" && typeof block.text === "string")
+      .map((block: any) => block.text)
+      .join("");
+  }
+
+  function lastVisibleUserText(messages: any[]): string {
+    for (let index = messages.length - 1; index >= 0; index -= 1) {
+      const message = messages[index];
+      if (message?.role === "user") return textFromContent(message.content);
+    }
+    return "";
+  }
+
+  function assistant(content: any[], timestamp = Date.now()) {
+    return {
+      role: "assistant",
+      content,
+      api: "anthropic-messages",
+      provider: "anthropic",
+      model: "fake",
+      usage: EMPTY_USAGE,
+      stopReason: content.some((block) => block.type === "toolCall") ? "toolUse" : "stop",
+      timestamp,
+    };
+  }
+
+  const streamSimple = vi.fn((model: any, context: any) => {
+    streamCalls.push({ model: clone(model), context: clone(context) });
+    const stream = actual.createAssistantMessageEventStream();
+    const last = context.messages.at(-1);
+    const prompt = lastVisibleUserText(context.messages);
+    const timestamp = Date.now();
+    const message = last?.role === "toolResult"
+      ? assistant([{ type: "text", text: "ok" }], timestamp)
+      : prompt === "think"
+        ? assistant([
+            { type: "thinking", thinking: "raw thought", thinkingSignature: "sig-1" },
+            { type: "text", text: "ok" },
+          ], timestamp)
+        : prompt === "use tool"
+          ? assistant([
+              {
+                type: "toolCall",
+                id: "tool-1",
+                name: "read",
+                arguments: { path: "book-a/story/story_bible.md" },
+              },
+            ], timestamp)
+          : assistant([{ type: "text", text: "ok" }], timestamp);
+
+    stream.push({
+      type: "done",
+      reason: message.stopReason === "toolUse" ? "toolUse" : "stop",
+      message,
+    });
+    return stream;
+  });
+
   return {
     ...actual,
-    streamSimple: vi.fn(),
+    streamSimple,
     getEnvApiKey: vi.fn(() => "fake-key"),
     getModel: vi.fn((provider: string, id: string) => ({
       provider,
       id,
+      name: id,
       api: "anthropic-messages",
+      baseUrl: "",
+      reasoning: false,
+      input: ["text"],
+      cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
+      contextWindow: 200_000,
+      maxTokens: 4096,
     })),
   };
 });
 
 import { runAgentSession, evictAgentCache } from "../agent/agent-session.js";
+import { readTranscriptEvents } from "../interaction/session-transcript.js";
 
 describe("runAgentSession cache — bookId switch", () => {
   let projectRoot: string;
@@ -82,6 +135,7 @@ describe("runAgentSession cache — bookId switch", () => {
       "书B 的真相",
     );
     agentInstances.length = 0;
+    streamCalls.length = 0;
   });
 
   afterEach(async () => {
@@ -106,17 +160,10 @@ describe("runAgentSession cache — bookId switch", () => {
 
     expect(agentInstances).toHaveLength(2);
 
-    const injected = await agentInstances[1].transformContext([]);
-    const body = JSON.stringify(injected);
+    const body = JSON.stringify(streamCalls.at(-1)?.context.messages);
     expect(body).toContain("书B 的真相");
     expect(body).not.toContain("书A 的真相");
-
-    // Prior conversation must be replayed into the rebuilt Agent via initialMessages,
-    // not silently dropped — this is the whole reason eviction preserves messages.
-    const preservedUser = agentInstances[1].state.messages.find(
-      (m: any) => m.role === "user" && typeof m.content === "string" && m.content.includes("earlier question about book A"),
-    );
-    expect(preservedUser).toBeDefined();
+    expect(body).toContain("earlier question about book A");
   });
 
   it("rebuilds Agent when bookId goes from null to a real book", async () => {
@@ -135,8 +182,7 @@ describe("runAgentSession cache — bookId switch", () => {
     );
 
     expect(agentInstances).toHaveLength(2);
-    const injected = await agentInstances[1].transformContext([]);
-    expect(JSON.stringify(injected)).toContain("书A 的真相");
+    expect(JSON.stringify(streamCalls.at(-1)?.context.messages)).toContain("书A 的真相");
   });
 
   it("treats undefined bookId as null (no spurious rebuild)", async () => {
@@ -149,8 +195,6 @@ describe("runAgentSession cache — bookId switch", () => {
     );
     expect(agentInstances).toHaveLength(1);
 
-    // A caller passing `undefined` (e.g., `activeBookId` not set and no `?? null`
-    // guard) must not cause an eviction when the cached agent already holds null.
     await runAgentSession(
       { sessionId: "s1", bookId: undefined as any, language: "zh", pipeline, projectRoot, model },
       "hi",
@@ -174,4 +218,64 @@ describe("runAgentSession cache — bookId switch", () => {
 
     expect(agentInstances).toHaveLength(1);
   });
+
+  it("把真实 Agent 的 message_end 写入 JSONL，并在 cache 失效后恢复 raw AgentMessage", async () => {
+    const model = { provider: "x", id: "y", api: "anthropic-messages" } as any;
+    const pipeline = {} as any;
+
+    await runAgentSession(
+      { sessionId: "s1", bookId: "book-a", language: "zh", pipeline, projectRoot, model },
+      "think",
+    );
+
+    const events = await readTranscriptEvents(projectRoot, "s1");
+    expect(events.map((event) => event.type)).toContain("request_committed");
+
+    evictAgentCache("s1");
+
+    await runAgentSession(
+      { sessionId: "s1", bookId: "book-a", language: "zh", pipeline, projectRoot, model },
+      "again",
+    );
+
+    expect(agentInstances).toHaveLength(2);
+    expect(JSON.stringify(streamCalls.at(-1)?.context.messages)).toContain("raw thought");
+    expect(streamCalls.at(-1)?.context.messages).toEqual(
+      expect.arrayContaining([
+        expect.objectContaining({ role: "assistant" }),
+      ]),
+    );
+  });
+
+  it("恢复 transcript 中的 toolResult message", async () => {
+    const model = { provider: "x", id: "y", api: "anthropic-messages" } as any;
+    const pipeline = {} as any;
+
+    await runAgentSession(
+      { sessionId: "s1", bookId: "book-a", language: "zh", pipeline, projectRoot, model },
+      "use tool",
+    );
+
+    evictAgentCache("s1");
+
+    await runAgentSession(
+      { sessionId: "s1", bookId: "book-a", language: "zh", pipeline, projectRoot, model },
+      "again",
+    );
+
+    expect(agentInstances).toHaveLength(2);
+    expect(streamCalls.at(-1)?.context.messages.some(
+      (message: any) => message.role === "toolResult" && message.toolCallId === "tool-1",
+    )).toBe(true);
+
+    const messageEvents = (await readTranscriptEvents(projectRoot, "s1"))
+      .filter((event) => event.type === "message");
+    const toolAssistant = messageEvents.find(
+      (event: any) => event.toolCallId === "tool-1" && event.role === "assistant",
+    ) as any;
+    const toolResult = messageEvents.find(
+      (event: any) => event.toolCallId === "tool-1" && event.role === "toolResult",
+    ) as any;
+    expect(toolResult.sourceToolAssistantUuid).toBe(toolAssistant.uuid);
+  });
 });

packages/core/src/agent/agent-session.ts

@@ -1,3 +1,4 @@
+import { randomUUID } from "node:crypto";
 import { Agent } from "@mariozechner/pi-agent-core";
 import type { AgentEvent, AgentMessage } from "@mariozechner/pi-agent-core";
 import { streamSimple, getModel, getEnvApiKey } from "@mariozechner/pi-ai";
@@ -16,6 +17,13 @@ import {
   createWriteTruthFileTool,
 } from "./agent-tools.js";
 import { createBookContextTransform } from "./context-transform.js";
+import {
+  appendTranscriptEvent,
+  nextTranscriptSeq,
+  readTranscriptEvents,
+} from "../interaction/session-transcript.js";
+import { restoreAgentMessagesFromTranscript } from "../interaction/session-transcript-restore.js";
+import type { TranscriptRole } from "../interaction/session-transcript-schema.js";
 
 // ---------------------------------------------------------------------------
 // Types
@@ -43,8 +51,8 @@ export interface AgentSessionConfig {
 export interface AgentSessionResult {
   /** Extracted text from the final assistant message. */
   responseText: string;
-  /** Full conversation history for persistence. */
-  messages: Array<{ role: string; content: string; thinking?: string }>;
+  /** Full raw Agent conversation history. */
+  messages: AgentMessage[];
 }
 
 // ---------------------------------------------------------------------------
@@ -59,6 +67,8 @@ export interface AgentSessionResult {
 interface CachedAgent {
   agent: Agent;
   bookId: string | null;
+  modelId: string | null;
+  lastCommittedSeq: number;
   lastActive: number;
 }
 
@@ -110,6 +120,84 @@ function resolveModel(spec: AgentSessionConfig["model"]): Model<Api> {
   return getModel(provider as any, modelId as any);
 }
 
+function modelIdFromSpec(spec: AgentSessionConfig["model"]): string | null {
+  if (!spec || typeof spec !== "object") return null;
+  if ("id" in spec && typeof spec.id === "string") return spec.id;
+  if ("modelId" in spec && typeof spec.modelId === "string") return spec.modelId;
+  return null;
+}
+
+async function latestCommittedSeq(projectRoot: string, sessionId: string): Promise<number> {
+  const events = await readTranscriptEvents(projectRoot, sessionId);
+  return events
+    .filter((event) => event.type === "request_committed")
+    .reduce((max, event) => Math.max(max, event.seq), 0);
+}
+
+function transcriptRoleForMessage(message: AgentMessage): TranscriptRole | null {
+  if (!message || typeof message !== "object" || !("role" in message)) return null;
+  const role = (message as { role?: unknown }).role;
+  return role === "user" || role === "assistant" || role === "toolResult" || role === "system"
+    ? role
+    : null;
+}
+
+function firstToolCallId(message: AgentMessage): string | undefined {
+  if (!message || typeof message !== "object" || !("content" in message)) return undefined;
+  const content = (message as { content?: unknown }).content;
+  if (!Array.isArray(content)) return undefined;
+  const block = content.find(
+    (item): item is { type: "toolCall"; id: string } =>
+      !!item &&
+      typeof item === "object" &&
+      (item as { type?: unknown }).type === "toolCall" &&
+      typeof (item as { id?: unknown }).id === "string",
+  );
+  return block?.id;
+}
+
+function toolCallIdForMessage(message: AgentMessage): string | undefined {
+  if (!message || typeof message !== "object") return undefined;
+  if ((message as { role?: unknown }).role === "toolResult") {
+    const toolCallId = (message as { toolCallId?: unknown }).toolCallId;
+    return typeof toolCallId === "string" && toolCallId.length > 0 ? toolCallId : undefined;
+  }
+  return firstToolCallId(message);
+}
+
+function messageTimestamp(message: AgentMessage): number {
+  if (message && typeof message === "object") {
+    const timestamp = (message as { timestamp?: unknown }).timestamp;
+    if (typeof timestamp === "number" && Number.isFinite(timestamp) && timestamp >= 0) {
+      return Math.floor(timestamp);
+    }
+  }
+  return Date.now();
+}
+
+async function ensureSessionCreatedEvent(
+  projectRoot: string,
+  sessionId: string,
+  bookId: string | null,
+): Promise<number> {
+  const events = await readTranscriptEvents(projectRoot, sessionId);
+  if (events.length > 0) return events.reduce((max, event) => Math.max(max, event.seq), 0) + 1;
+
+  const now = Date.now();
+  await appendTranscriptEvent(projectRoot, {
+    type: "session_created",
+    version: 1,
+    sessionId,
+    seq: 1,
+    timestamp: now,
+    bookId,
+    title: null,
+    createdAt: now,
+    updatedAt: now,
+  });
+  return 2;
+}
+
 /**
  * Extract readable text from an AssistantMessage's content array.
  * Filters out tool-call blocks; concatenates text blocks.
@@ -218,6 +306,7 @@ export async function runAgentSession(
   // skipped) and we don't want that to (a) throw in path.join or (b) trigger
   // a spurious cache eviction because `null !== undefined`.
   const bookId: string | null = config.bookId ?? null;
+  const requestedModelId = modelIdFromSpec(config.model);
 
   // ----- Resolve or create Agent -----
   let cached = agentCache.get(sessionId);
@@ -228,27 +317,27 @@ export async function runAgentSession(
     // in systemPrompt / tools / transformContext), so a mismatch means the
     // cached Agent would keep using stale context — including reading truth
     // files from the wrong book's story/ directory.
-    const currentModelId = (cached.agent.state.model as any)?.id;
-    const newModelId = typeof config.model === 'object' && 'id' in config.model
-      ? (config.model as any).id
-      : undefined;
-    const modelChanged = !!(currentModelId && newModelId && currentModelId !== newModelId);
+    const modelChanged = !!(
+      cached.modelId &&
+      requestedModelId &&
+      cached.modelId !== requestedModelId
+    );
     const bookChanged = cached.bookId !== bookId;
 
     if (modelChanged || bookChanged) {
-      // Preserve conversation messages for re-injection
-      const preservedMessages = agentMessagesToPlain(cached.agent.state.messages);
       agentCache.delete(sessionId);
       cached = undefined;
-      // Pass preserved messages as initialMessages if none were provided
-      if (!initialMessages || initialMessages.length === 0) {
-        initialMessages = preservedMessages;
-      }
     }
   }
 
   if (!cached) {
     const model = resolveModel(config.model);
+    const restoredMessages = await restoreAgentMessagesFromTranscript(projectRoot, sessionId);
+    const initialAgentMessages = restoredMessages.length > 0
+      ? restoredMessages
+      : initialMessages && initialMessages.length > 0
+        ? plainToAgentMessages(initialMessages)
+        : [];
     const agent = new Agent({
       initialState: {
         model,
@@ -264,6 +353,7 @@ export async function runAgentSession(
           createGrepTool(projectRoot),
           createLsTool(projectRoot),
         ],
+        messages: initialAgentMessages,
       },
       transformContext: createBookContextTransform(bookId, projectRoot),
       streamFn: streamSimple,
@@ -273,12 +363,13 @@ export async function runAgentSession(
       },
     });
 
-    // Restore prior conversation if provided.
-    if (initialMessages && initialMessages.length > 0) {
-      agent.state.messages = plainToAgentMessages(initialMessages);
-    }
-
-    cached = { agent, bookId, lastActive: Date.now() };
+    cached = {
+      agent,
+      bookId,
+      modelId: model.id ?? requestedModelId,
+      lastCommittedSeq: await latestCommittedSeq(projectRoot, sessionId),
+      lastActive: Date.now(),
+    };
     agentCache.set(sessionId, cached);
     ensureCleanupTimer();
   }
@@ -286,27 +377,96 @@ export async function runAgentSession(
   cached.lastActive = Date.now();
   const { agent } = cached;
 
-  // ----- Subscribe to events (for SSE streaming to frontend) -----
-  let unsubscribe: (() => void) | undefined;
-  if (onEvent) {
-    unsubscribe = agent.subscribe((event: AgentEvent) => {
-      onEvent(event);
+  // ----- Prepare transcript persistence -----
+  const requestId = randomUUID();
+  let seq = await ensureSessionCreatedEvent(projectRoot, sessionId, bookId);
+  await appendTranscriptEvent(projectRoot, {
+    type: "request_started",
+    version: 1,
+    sessionId,
+    requestId,
+    seq: seq++,
+    timestamp: Date.now(),
+    input: userMessage,
+  });
+
+  let parentUuid: string | null = null;
+  let piTurnIndex = 0;
+  let lastAssistantUuid: string | null = null;
+
+  const persistAgentEvent = async (event: AgentEvent): Promise<void> => {
+    if (event.type === "turn_start") {
+      piTurnIndex += 1;
+      return;
+    }
+    if (event.type !== "message_end") return;
+
+    const role = transcriptRoleForMessage(event.message);
+    if (!role) return;
+
+    const uuid = randomUUID();
+    const isToolResult = role === "toolResult";
+    const toolCallId = toolCallIdForMessage(event.message);
+    await appendTranscriptEvent(projectRoot, {
+      type: "message",
+      version: 1,
+      sessionId,
+      requestId,
+      uuid,
+      parentUuid: isToolResult && lastAssistantUuid ? lastAssistantUuid : parentUuid,
+      seq: seq++,
+      role,
+      timestamp: messageTimestamp(event.message),
+      piTurnIndex,
+      ...(toolCallId ? { toolCallId } : {}),
+      ...(isToolResult && lastAssistantUuid
+        ? { sourceToolAssistantUuid: lastAssistantUuid }
+        : {}),
+      message: event.message,
     });
-  }
+
+    if (role === "assistant") lastAssistantUuid = uuid;
+    parentUuid = uuid;
+  };
+
+  // ----- Subscribe to events (transcript persistence + SSE forwarding) -----
+  const unsubscribe = agent.subscribe(async (event: AgentEvent) => {
+    await persistAgentEvent(event);
+    onEvent?.(event);
+  });
 
   // ----- Execute the turn -----
   try {
     await agent.prompt(userMessage);
+    await appendTranscriptEvent(projectRoot, {
+      type: "request_committed",
+      version: 1,
+      sessionId,
+      requestId,
+      seq: seq++,
+      timestamp: Date.now(),
+    });
+    cached.lastCommittedSeq = seq - 1;
+  } catch (error) {
+    await appendTranscriptEvent(projectRoot, {
+      type: "request_failed",
+      version: 1,
+      sessionId,
+      requestId,
+      seq: seq++,
+      timestamp: Date.now(),
+      error: error instanceof Error ? error.message : String(error),
+    });
+    throw error;
   } finally {
-    unsubscribe?.();
+    unsubscribe();
   }
 
   // ----- Extract result -----
   const allMessages = agent.state.messages;
   const responseText = extractResponseText(allMessages);
-  const plainMessages = agentMessagesToPlain(allMessages);
 
-  return { responseText, messages: plainMessages };
+  return { responseText, messages: allMessages.slice() };
 }
 
 /**