Sfoglia il codice sorgente

fix: isolate agent transcript persistence

fanghanjun 6 mesi fa
parent
commit
255d54ddd2

+ 101 - 2
packages/core/src/__tests__/agent-session.test.ts

@@ -12,9 +12,11 @@ const EMPTY_USAGE = {
   cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
 };
 
-const { agentInstances, streamCalls } = vi.hoisted(() => ({
+const { agentInstances, streamCalls, heldStreamCompletions, heldStreamWaiters } = vi.hoisted(() => ({
   agentInstances: [] as any[],
   streamCalls: [] as Array<{ model: any; context: any }>,
+  heldStreamCompletions: [] as Array<() => void>,
+  heldStreamWaiters: [] as Array<() => void>,
 }));
 
 vi.mock("@mariozechner/pi-agent-core", async () => {
@@ -106,7 +108,10 @@ vi.mock("@mariozechner/pi-ai", async () => {
       reason: message.stopReason === "toolUse" ? "toolUse" : "stop",
       message,
     });
-    if (prompt.startsWith("slow ")) {
+    if (prompt === "hold for interleave") {
+      heldStreamCompletions.push(done);
+      heldStreamWaiters.splice(0).forEach((resolve) => resolve());
+    } else if (prompt.startsWith("slow ")) {
       setTimeout(done, prompt.includes("first") ? 20 : 0);
     } else {
       done();
@@ -137,15 +142,18 @@ import { runAgentSession, evictAgentCache } from "../agent/agent-session.js";
 import {
   appendManualSessionMessages,
   appendTranscriptEvent,
+  appendTranscriptEvents,
   readTranscriptEvents,
 } from "../interaction/session-transcript.js";
 import { restoreAgentMessagesFromTranscript } from "../interaction/session-transcript-restore.js";
 
 describe("runAgentSession cache — bookId switch", () => {
   let projectRoot: string;
+  let otherProjectRoot: string | null;
 
   beforeEach(async () => {
     projectRoot = await mkdtemp(join(tmpdir(), "inkos-agent-cache-"));
+    otherProjectRoot = null;
     await mkdir(join(projectRoot, "books", "book-a", "story"), { recursive: true });
     await writeFile(
       join(projectRoot, "books", "book-a", "story", "story_bible.md"),
@@ -158,13 +166,18 @@ describe("runAgentSession cache — bookId switch", () => {
     );
     agentInstances.length = 0;
     streamCalls.length = 0;
+    heldStreamCompletions.length = 0;
+    heldStreamWaiters.length = 0;
   });
 
   afterEach(async () => {
     evictAgentCache("s1");
     evictAgentCache("s-cache-seq");
     evictAgentCache("s-error");
+    evictAgentCache("s-project-root-cache");
+    evictAgentCache("s-interleave-seq");
     await rm(projectRoot, { recursive: true, force: true });
+    if (otherProjectRoot) await rm(otherProjectRoot, { recursive: true, force: true });
   });
 
   it("rebuilds Agent when bookId changes for same sessionId", async () => {
@@ -243,6 +256,42 @@ describe("runAgentSession cache — bookId switch", () => {
     expect(agentInstances).toHaveLength(1);
   });
 
+  it("keeps cached Agents isolated by projectRoot for the same sessionId", async () => {
+    const model = { provider: "x", id: "y", api: "anthropic-messages" } as any;
+    const pipeline = {} as any;
+    otherProjectRoot = await mkdtemp(join(tmpdir(), "inkos-agent-cache-other-"));
+    await mkdir(join(otherProjectRoot, "books", "book-a", "story"), { recursive: true });
+    await writeFile(
+      join(otherProjectRoot, "books", "book-a", "story", "story_bible.md"),
+      "另一个 projectRoot 的真相",
+    );
+
+    await runAgentSession(
+      { sessionId: "s-project-root-cache", bookId: "book-a", language: "zh", pipeline, projectRoot, model },
+      "root A",
+    );
+    await runAgentSession(
+      {
+        sessionId: "s-project-root-cache",
+        bookId: "book-a",
+        language: "zh",
+        pipeline,
+        projectRoot: otherProjectRoot,
+        model,
+      },
+      "root B",
+    );
+    await runAgentSession(
+      { sessionId: "s-project-root-cache", bookId: "book-a", language: "zh", pipeline, projectRoot, model },
+      "root A again",
+    );
+
+    expect(agentInstances).toHaveLength(2);
+    const body = JSON.stringify(streamCalls.at(-1)?.context.messages);
+    expect(body).toContain("书A 的真相");
+    expect(body).not.toContain("另一个 projectRoot 的真相");
+  });
+
   it("rebuilds Agent when model id is unchanged but API protocol changes", async () => {
     const pipeline = {} as any;
     const legacyGoogle = {
@@ -689,4 +738,54 @@ describe("runAgentSession cache — bookId switch", () => {
     expect(events.filter((event) => event.type === "request_committed")).toHaveLength(2);
     expect(events.map((event) => event.seq)).toEqual(events.map((_, index) => index + 1));
   });
+
+  it("assigns transcript seq after interleaved non-agent writes", async () => {
+    const model = { provider: "x", id: "y", api: "anthropic-messages" } as any;
+    const pipeline = {} as any;
+    let resolveTurnStarted!: () => void;
+    const turnStarted = new Promise<void>((resolve) => {
+      resolveTurnStarted = resolve;
+    });
+    let interleavedWrite: Promise<unknown> | null = null;
+
+    const running = runAgentSession(
+      {
+        sessionId: "s-interleave-seq",
+        bookId: "book-a",
+        language: "zh",
+        pipeline,
+        projectRoot,
+        model,
+        onEvent: (event) => {
+          if (event.type !== "turn_start" || interleavedWrite) return;
+          interleavedWrite = appendTranscriptEvents(projectRoot, "s-interleave-seq", ({ nextSeq }) => [{
+            type: "session_metadata_updated",
+            version: 1,
+            sessionId: "s-interleave-seq",
+            seq: nextSeq,
+            timestamp: Date.now(),
+            updatedAt: Date.now(),
+            title: "interleaved update",
+          }]);
+          resolveTurnStarted();
+        },
+      },
+      "hold for interleave",
+    );
+
+    await turnStarted;
+    await interleavedWrite;
+    if (heldStreamCompletions.length === 0) {
+      await new Promise<void>((resolve) => {
+        heldStreamWaiters.push(resolve);
+      });
+    }
+    const finishStream = heldStreamCompletions.shift();
+    expect(finishStream).toBeTypeOf("function");
+    finishStream?.();
+    await running;
+
+    const events = await readTranscriptEvents(projectRoot, "s-interleave-seq");
+    expect(events.map((event) => event.seq)).toEqual(events.map((_, index) => index + 1));
+  });
 });

+ 146 - 0
packages/core/src/__tests__/session-transcript-restore.test.ts

@@ -1021,4 +1021,150 @@ describe("session transcript restore", () => {
 
     expect(session?.messages.map((message) => message.content)).toEqual(["先问", "后答"]);
   });
+
+  it("does not carry pending tool executions or thinking across request boundaries", async () => {
+    await appendTranscriptEvent(projectRoot, {
+      type: "session_created",
+      version: 1,
+      sessionId: "s1",
+      seq: 1,
+      timestamp: 1,
+      bookId: "book-a",
+      title: null,
+      createdAt: 1,
+      updatedAt: 1,
+    });
+    await appendTranscriptEvent(projectRoot, {
+      type: "request_started",
+      version: 1,
+      sessionId: "s1",
+      requestId: "r1",
+      seq: 2,
+      timestamp: 2,
+      input: "列目录",
+    });
+    await appendTranscriptEvent(projectRoot, {
+      type: "message",
+      version: 1,
+      sessionId: "s1",
+      requestId: "r1",
+      uuid: "u1",
+      parentUuid: null,
+      seq: 3,
+      role: "user",
+      timestamp: 3,
+      message: { role: "user", content: "列目录", timestamp: 3 },
+    } as MessageEvent);
+    await appendTranscriptEvent(projectRoot, {
+      type: "message",
+      version: 1,
+      sessionId: "s1",
+      requestId: "r1",
+      uuid: "a1",
+      parentUuid: "u1",
+      seq: 4,
+      role: "assistant",
+      timestamp: 4,
+      toolCallId: "ls-1",
+      message: {
+        role: "assistant",
+        content: [
+          { type: "thinking", thinking: "第一轮工具前思考" },
+          { type: "toolCall", id: "ls-1", name: "ls", arguments: { subdir: "story" } },
+        ],
+        api: "anthropic-messages",
+        provider: "anthropic",
+        model: "claude",
+        usage,
+        stopReason: "toolUse",
+        timestamp: 4,
+      },
+    } as MessageEvent);
+    await appendTranscriptEvent(projectRoot, {
+      type: "message",
+      version: 1,
+      sessionId: "s1",
+      requestId: "r1",
+      uuid: "t1",
+      parentUuid: "a1",
+      seq: 5,
+      role: "toolResult",
+      timestamp: 5,
+      toolCallId: "ls-1",
+      sourceToolAssistantUuid: "a1",
+      message: {
+        role: "toolResult",
+        toolCallId: "ls-1",
+        toolName: "ls",
+        content: [{ type: "text", text: "roles/" }],
+        isError: false,
+        timestamp: 5,
+      },
+    } as MessageEvent);
+    await appendTranscriptEvent(projectRoot, {
+      type: "request_committed",
+      version: 1,
+      sessionId: "s1",
+      requestId: "r1",
+      seq: 6,
+      timestamp: 6,
+    });
+    await appendTranscriptEvent(projectRoot, {
+      type: "request_started",
+      version: 1,
+      sessionId: "s1",
+      requestId: "r2",
+      seq: 7,
+      timestamp: 7,
+      input: "继续",
+    });
+    await appendTranscriptEvent(projectRoot, {
+      type: "message",
+      version: 1,
+      sessionId: "s1",
+      requestId: "r2",
+      uuid: "u2",
+      parentUuid: null,
+      seq: 8,
+      role: "user",
+      timestamp: 8,
+      message: { role: "user", content: "继续", timestamp: 8 },
+    } as MessageEvent);
+    await appendTranscriptEvent(projectRoot, {
+      type: "message",
+      version: 1,
+      sessionId: "s1",
+      requestId: "r2",
+      uuid: "a2",
+      parentUuid: "u2",
+      seq: 9,
+      role: "assistant",
+      timestamp: 9,
+      message: {
+        role: "assistant",
+        content: [{ type: "text", text: "第二轮回答" }],
+        api: "anthropic-messages",
+        provider: "anthropic",
+        model: "claude",
+        usage,
+        stopReason: "stop",
+        timestamp: 9,
+      },
+    } as MessageEvent);
+    await appendTranscriptEvent(projectRoot, {
+      type: "request_committed",
+      version: 1,
+      sessionId: "s1",
+      requestId: "r2",
+      seq: 10,
+      timestamp: 10,
+    });
+
+    const session = await deriveBookSessionFromTranscript(projectRoot, "s1");
+    const secondAssistant = session?.messages.find((message) => message.content === "第二轮回答");
+
+    expect(secondAssistant).toMatchObject({ role: "assistant", content: "第二轮回答" });
+    expect(secondAssistant).not.toHaveProperty("thinking");
+    expect(secondAssistant).not.toHaveProperty("toolExecutions");
+  });
 });

+ 54 - 35
packages/core/src/agent/agent-session.ts

@@ -19,7 +19,6 @@ import {
 import { createBookContextTransform } from "./context-transform.js";
 import {
   appendTranscriptEvents,
-  appendTranscriptEvent,
   readTranscriptEvents,
 } from "../interaction/session-transcript.js";
 import {
@@ -27,7 +26,7 @@ import {
   adaptRestoredAgentMessagesForModel,
   restoreAgentMessagesFromTranscript,
 } from "../interaction/session-transcript-restore.js";
-import type { TranscriptRole } from "../interaction/session-transcript-schema.js";
+import type { TranscriptEvent, TranscriptRole } from "../interaction/session-transcript-schema.js";
 
 // ---------------------------------------------------------------------------
 // Types
@@ -69,6 +68,7 @@ export interface AgentSessionResult {
 
 interface CachedAgent {
   agent: Agent;
+  sessionId: string;
   projectRoot: string;
   bookId: string | null;
   language: string;
@@ -148,6 +148,10 @@ function sessionQueueKey(projectRoot: string, sessionId: string): string {
   return `${projectRoot}\0${sessionId}`;
 }
 
+function agentCacheKey(projectRoot: string, sessionId: string): string {
+  return sessionQueueKey(projectRoot, sessionId);
+}
+
 async function runInAgentSessionQueue<T>(
   projectRoot: string,
   sessionId: string,
@@ -225,9 +229,9 @@ async function ensureSessionCreatedEvent(
   projectRoot: string,
   sessionId: string,
   bookId: string | null,
-): Promise<number> {
-  const created = await appendTranscriptEvents(projectRoot, sessionId, ({ events, nextSeq }) => {
-    if (events.length > 0) return [];
+): Promise<void> {
+  await appendTranscriptEvents(projectRoot, sessionId, ({ events, nextSeq }) => {
+    if (events.some((event) => event.type === "session_created")) return [];
 
     const now = Date.now();
     return [{
@@ -242,13 +246,19 @@ async function ensureSessionCreatedEvent(
       updatedAt: now,
     }];
   });
+}
 
-  if (created.length > 0) {
-    return created.reduce((max, event) => Math.max(max, event.seq), 0) + 1;
-  }
-
-  const events = await readTranscriptEvents(projectRoot, sessionId);
-  return events.reduce((max, event) => Math.max(max, event.seq), 0) + 1;
+async function appendAgentTranscriptEvent(
+  projectRoot: string,
+  sessionId: string,
+  buildEvent: (seq: number) => TranscriptEvent,
+): Promise<TranscriptEvent> {
+  const events = await appendTranscriptEvents(projectRoot, sessionId, ({ nextSeq }) => [
+    buildEvent(nextSeq),
+  ]);
+  const event = events[0];
+  if (!event) throw new Error(`Failed to append transcript event for session "${sessionId}"`);
+  return event;
 }
 
 /**
@@ -486,9 +496,10 @@ async function runAgentSessionUnlocked(
   const model = resolveModel(config.model);
   const requestedModelIdentity = agentModelIdentity(model);
   const allowSystemFileRead = config.allowSystemFileRead ?? envFlagEnabled(process.env.INKOS_AGENT_ALLOW_SYSTEM_READ, true);
+  const cacheKey = agentCacheKey(projectRoot, sessionId);
 
   // ----- Resolve or create Agent -----
-  let cached = agentCache.get(sessionId);
+  let cached = agentCache.get(cacheKey);
   let currentCommittedSeq: number | undefined;
 
   if (cached) {
@@ -514,7 +525,7 @@ async function runAgentSessionUnlocked(
       readPermissionChanged ||
       transcriptChanged
     ) {
-      agentCache.delete(sessionId);
+      agentCache.delete(cacheKey);
       cached = undefined;
     }
   }
@@ -557,6 +568,7 @@ async function runAgentSessionUnlocked(
 
     cached = {
       agent,
+      sessionId,
       projectRoot,
       bookId,
       language,
@@ -566,7 +578,7 @@ async function runAgentSessionUnlocked(
       lastCommittedSeq: currentCommittedSeq ?? await latestCommittedSeq(projectRoot, sessionId),
       lastActive: Date.now(),
     };
-    agentCache.set(sessionId, cached);
+    agentCache.set(cacheKey, cached);
     ensureCleanupTimer();
   }
 
@@ -575,16 +587,16 @@ async function runAgentSessionUnlocked(
 
   // ----- Prepare transcript persistence -----
   const requestId = randomUUID();
-  let seq = await ensureSessionCreatedEvent(projectRoot, sessionId, bookId);
-  await appendTranscriptEvent(projectRoot, {
+  await ensureSessionCreatedEvent(projectRoot, sessionId, bookId);
+  await appendAgentTranscriptEvent(projectRoot, sessionId, (seq) => ({
     type: "request_started",
     version: 1,
     sessionId,
     requestId,
-    seq: seq++,
+    seq,
     timestamp: Date.now(),
     input: userMessage,
-  });
+  }));
 
   let parentUuid: string | null = null;
   let piTurnIndex = 0;
@@ -603,14 +615,14 @@ async function runAgentSessionUnlocked(
     const uuid = randomUUID();
     const isToolResult = role === "toolResult";
     const toolCallId = toolCallIdForMessage(event.message);
-    await appendTranscriptEvent(projectRoot, {
+    await appendAgentTranscriptEvent(projectRoot, sessionId, (seq) => ({
       type: "message",
       version: 1,
       sessionId,
       requestId,
       uuid,
       parentUuid: isToolResult && lastAssistantUuid ? lastAssistantUuid : parentUuid,
-      seq: seq++,
+      seq,
       role,
       timestamp: messageTimestamp(event.message),
       piTurnIndex,
@@ -619,7 +631,7 @@ async function runAgentSessionUnlocked(
         ? { sourceToolAssistantUuid: lastAssistantUuid }
         : {}),
       message: event.message,
-    });
+    }));
 
     if (role === "assistant") lastAssistantUuid = uuid;
     parentUuid = uuid;
@@ -641,38 +653,39 @@ async function runAgentSessionUnlocked(
     finalAssistant = lastAssistantMessage(agent.state.messages);
     errorMessage = assistantErrorMessage(finalAssistant);
     if (errorMessage) {
-      await appendTranscriptEvent(projectRoot, {
+      const failedError = errorMessage;
+      await appendAgentTranscriptEvent(projectRoot, sessionId, (seq) => ({
         type: "request_failed",
         version: 1,
         sessionId,
         requestId,
-        seq: seq++,
+        seq,
         timestamp: Date.now(),
-        error: errorMessage,
-      });
-      agentCache.delete(sessionId);
+        error: failedError,
+      }));
+      agentCache.delete(cacheKey);
     } else {
-      await appendTranscriptEvent(projectRoot, {
+      const committed = await appendAgentTranscriptEvent(projectRoot, sessionId, (seq) => ({
         type: "request_committed",
         version: 1,
         sessionId,
         requestId,
-        seq: seq++,
+        seq,
         timestamp: Date.now(),
-      });
-      cached.lastCommittedSeq = seq - 1;
+      }));
+      cached.lastCommittedSeq = committed.seq;
     }
   } catch (error) {
-    await appendTranscriptEvent(projectRoot, {
+    await appendAgentTranscriptEvent(projectRoot, sessionId, (seq) => ({
       type: "request_failed",
       version: 1,
       sessionId,
       requestId,
-      seq: seq++,
+      seq,
       timestamp: Date.now(),
       error: error instanceof Error ? error.message : String(error),
-    });
-    agentCache.delete(sessionId);
+    }));
+    agentCache.delete(cacheKey);
     throw error;
   } finally {
     unsubscribe();
@@ -697,5 +710,11 @@ async function runAgentSessionUnlocked(
 
 /** Manually evict a cached Agent session. */
 export function evictAgentCache(sessionId: string): boolean {
-  return agentCache.delete(sessionId);
+  let deleted = agentCache.delete(sessionId);
+  for (const [key, entry] of agentCache) {
+    if (entry.sessionId !== sessionId) continue;
+    agentCache.delete(key);
+    deleted = true;
+  }
+  return deleted;
 }

+ 12 - 2
packages/core/src/interaction/session-transcript-restore.ts

@@ -437,6 +437,12 @@ function messageEventsToInteractionMessages(events: MessageEvent[]): Interaction
   const toolCalls = new Map<string, RestoredToolCall>();
   let pendingToolExecutions: ToolExecution[] = [];
   let pendingThinking: string[] = [];
+  let activeRequestId: string | null = null;
+
+  const clearPending = () => {
+    pendingToolExecutions = [];
+    pendingThinking = [];
+  };
 
   const objectArgs = (value: unknown): Record<string, unknown> | undefined => {
     if (isObject(value)) return value;
@@ -509,6 +515,10 @@ function messageEventsToInteractionMessages(events: MessageEvent[]): Interaction
 
   for (const event of events) {
     const raw = event.message as Record<string, unknown>;
+    if (activeRequestId !== event.requestId) {
+      activeRequestId = event.requestId;
+      clearPending();
+    }
 
     if (event.role === "assistant" && isObject(raw)) {
       rememberToolCalls(event, raw);
@@ -520,8 +530,7 @@ function messageEventsToInteractionMessages(events: MessageEvent[]): Interaction
       );
       if (message) {
         messages.push(message);
-        pendingToolExecutions = [];
-        pendingThinking = [];
+        clearPending();
       } else if (hasToolCallContent(raw) && currentThinking) {
         pendingThinking.push(currentThinking);
       }
@@ -534,6 +543,7 @@ function messageEventsToInteractionMessages(events: MessageEvent[]): Interaction
       continue;
     }
 
+    clearPending();
     const message = messageEventToInteractionMessage(event);
     if (message) messages.push(message);
   }