Parcourir la source

fix(core): evict agent cache when bookId changes on same sessionId (#204)

* fix(core): evict agent cache when bookId changes on same sessionId

同一 sessionId 切书时,agentCache 命中路径只检查 model、忽略 bookId。
Agent 构造时 bookId 被闭包进 systemPrompt / tools / transformContext,
命中后继续读旧书的 story 目录,造成跨书真相文件错乱。建书流程里
session 的 bookId 从 null 迁到真实 id 时也被同一问题击中。

CachedAgent 新增 bookId 字段;命中分支把 bookChanged 并入已有的
modelChanged 判断,走同一条 evict-and-preserve-messages 路径。
新增 agent-session.test.ts 覆盖三种情况:跨书重建、null→realId
重建、同 bookId 复用。

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

* fix(core): review follow-ups for agent-session bookId cache

- 在 runAgentSession 入口把 bookId 规范化为 `string | null`,调用方
  传入 undefined(绕过类型)时既不会在 path.join 抛错,也不会因为
  `null !== undefined` 触发无效的缓存重建。
- 切书测试新增一条断言:切书后 rebuilt Agent 必须含有之前用户的
  提问——验证 initialMessages 重播路径真的生效,而不只是新 Agent
  被构造。
- 新增一条测试专门覆盖"undefined 等同 null"不重建的不变式。
- CachedAgent 接口上加注释说明 projectRoot / language / pipeline
  虽然也被 Agent 闭包捕获,但单进程内视为稳定,故未纳入 change
  detection;将来若这些字段也会跨轮变化,需要一起加回检测。

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

---------

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
hanjun fang il y a 6 mois
Parent
commit
5abd37cdbf

+ 177 - 0
packages/core/src/__tests__/agent-session.test.ts

@@ -0,0 +1,177 @@
+import { describe, it, expect, beforeEach, afterEach, vi } from "vitest";
+import { mkdtemp, mkdir, writeFile, rm } from "node:fs/promises";
+import { join } from "node:path";
+import { tmpdir } from "node:os";
+
+const { agentInstances } = vi.hoisted(() => ({
+  agentInstances: [] as 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;
+    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;
+      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 };
+});
+
+vi.mock("@mariozechner/pi-ai", async () => {
+  const actual = await vi.importActual<any>("@mariozechner/pi-ai");
+  return {
+    ...actual,
+    streamSimple: vi.fn(),
+    getEnvApiKey: vi.fn(() => "fake-key"),
+    getModel: vi.fn((provider: string, id: string) => ({
+      provider,
+      id,
+      api: "anthropic-messages",
+    })),
+  };
+});
+
+import { runAgentSession, evictAgentCache } from "../agent/agent-session.js";
+
+describe("runAgentSession cache — bookId switch", () => {
+  let projectRoot: string;
+
+  beforeEach(async () => {
+    projectRoot = await mkdtemp(join(tmpdir(), "inkos-agent-cache-"));
+    await mkdir(join(projectRoot, "books", "book-a", "story"), { recursive: true });
+    await writeFile(
+      join(projectRoot, "books", "book-a", "story", "story_bible.md"),
+      "书A 的真相",
+    );
+    await mkdir(join(projectRoot, "books", "book-b", "story"), { recursive: true });
+    await writeFile(
+      join(projectRoot, "books", "book-b", "story", "story_bible.md"),
+      "书B 的真相",
+    );
+    agentInstances.length = 0;
+  });
+
+  afterEach(async () => {
+    evictAgentCache("s1");
+    await rm(projectRoot, { recursive: true, force: true });
+  });
+
+  it("rebuilds Agent when bookId changes for same sessionId", 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 },
+      "earlier question about book A",
+    );
+    expect(agentInstances).toHaveLength(1);
+
+    await runAgentSession(
+      { sessionId: "s1", bookId: "book-b", language: "zh", pipeline, projectRoot, model },
+      "new question",
+    );
+
+    expect(agentInstances).toHaveLength(2);
+
+    const injected = await agentInstances[1].transformContext([]);
+    const body = JSON.stringify(injected);
+    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();
+  });
+
+  it("rebuilds Agent when bookId goes from null to a real book", async () => {
+    const model = { provider: "x", id: "y", api: "anthropic-messages" } as any;
+    const pipeline = {} as any;
+
+    await runAgentSession(
+      { sessionId: "s1", bookId: null, language: "zh", pipeline, projectRoot, model },
+      "hi",
+    );
+    expect(agentInstances).toHaveLength(1);
+
+    await runAgentSession(
+      { sessionId: "s1", bookId: "book-a", language: "zh", pipeline, projectRoot, model },
+      "hi",
+    );
+
+    expect(agentInstances).toHaveLength(2);
+    const injected = await agentInstances[1].transformContext([]);
+    expect(JSON.stringify(injected)).toContain("书A 的真相");
+  });
+
+  it("treats undefined bookId as null (no spurious rebuild)", async () => {
+    const model = { provider: "x", id: "y", api: "anthropic-messages" } as any;
+    const pipeline = {} as any;
+
+    await runAgentSession(
+      { sessionId: "s1", bookId: null, language: "zh", pipeline, projectRoot, model },
+      "hi",
+    );
+    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",
+    );
+
+    expect(agentInstances).toHaveLength(1);
+  });
+
+  it("reuses Agent when bookId unchanged on same sessionId", 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 },
+      "hi",
+    );
+    await runAgentSession(
+      { sessionId: "s1", bookId: "book-a", language: "zh", pipeline, projectRoot, model },
+      "hi2",
+    );
+
+    expect(agentInstances).toHaveLength(1);
+  });
+});

+ 23 - 4
packages/core/src/agent/agent-session.ts

@@ -51,8 +51,14 @@ export interface AgentSessionResult {
 // Cache
 // ---------------------------------------------------------------------------
 
+// We only record fields that can realistically change between turns on the
+// same sessionId and are captured into the Agent at construction time.
+// `projectRoot`, `language`, and `pipeline` are also closure-captured by the
+// Agent (into systemPrompt / tools / transformContext), but within a single
+// server process they're treated as stable — we don't re-check them.
 interface CachedAgent {
   agent: Agent;
+  bookId: string | null;
   lastActive: number;
 }
 
@@ -205,18 +211,31 @@ export async function runAgentSession(
   userMessage: string,
   initialMessages?: Array<{ role: string; content: string }>,
 ): Promise<AgentSessionResult> {
-  const { sessionId, bookId, language, pipeline, projectRoot, onEvent } = config;
+  const { sessionId, language, pipeline, projectRoot, onEvent } = config;
+  // Normalize at the entry point so downstream comparisons, closures, and
+  // fs paths never see `undefined`. The type is already `string | null`, but
+  // some callers may bypass the type system (e.g. `activeBookId ?? null` gets
+  // 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;
 
   // ----- Resolve or create Agent -----
   let cached = agentCache.get(sessionId);
 
   if (cached) {
-    // Check if model changed — evict and rebuild if so
+    // Evict and rebuild if model OR bookId changed. Both are captured into the
+    // Agent at construction time (model via initialState, bookId via closures
+    // 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;
-    if (currentModelId && newModelId && currentModelId !== newModelId) {
+    const modelChanged = !!(currentModelId && newModelId && currentModelId !== newModelId);
+    const bookChanged = cached.bookId !== bookId;
+
+    if (modelChanged || bookChanged) {
       // Preserve conversation messages for re-injection
       const preservedMessages = agentMessagesToPlain(cached.agent.state.messages);
       agentCache.delete(sessionId);
@@ -259,7 +278,7 @@ export async function runAgentSession(
       agent.state.messages = plainToAgentMessages(initialMessages);
     }
 
-    cached = { agent, lastActive: Date.now() };
+    cached = { agent, bookId, lastActive: Date.now() };
     agentCache.set(sessionId, cached);
     ensureCleanupTimer();
   }