🐛 fix(sdk): engine 执行循环九项修复——运行锁 TOCTOU、缓存计价、压缩边界 uuid 回填

packages/sdk/src/agent.test.ts

@@ -2,7 +2,7 @@ import { afterEach, describe, expect, test } from "bun:test"
 import { mkdirSync, mkdtempSync, rmSync, writeFileSync } from "node:fs"
 import { tmpdir } from "node:os"
 import { join } from "node:path"
-import { createAgent } from "./agent.js"
+import { createAgent, sessionMessagesFromHistory } from "./agent.js"
 import { SkillTool } from "./tools/skill-tool.js"
 import type { SDKMessage, ToolDefinition } from "./types.js"
 import type { CreateMessageParams, CreateMessageResponse, LLMProvider } from "./providers/types.js"
@@ -1383,3 +1383,127 @@ describe("auth_status emission", () => {
     await agent.close()
   })
 })
+
+describe("Agent concurrent run lock (#357)", () => {
+  test("rejects a second query while the first run is still initializing", async () => {
+    const provider = new CapturingProvider()
+    const agent = createAgent({
+      persistSession: false,
+      tools: [],
+      provider,
+      model: "host/model-a",
+    })
+    await agent.getInitializationResult()
+
+    const first = (agent.query("first") as any)[Symbol.asyncIterator]() as AsyncIterator<SDKMessage>
+    const second = (agent.query("second") as any)[Symbol.asyncIterator]() as AsyncIterator<SDKMessage>
+    // Start both generators in the same tick: the first one grabs the run
+    // lock synchronously, so the second must be rejected before either
+    // engine is even constructed.
+    const firstStart = first.next()
+    const secondOutcome = await second.next().then(
+      () => "allowed",
+      (error: Error) => error.message,
+    )
+    expect(secondOutcome).toBe("agent is running")
+
+    // The first run completes normally.
+    let done = await firstStart
+    while (!done.done) {
+      done = await first.next()
+    }
+    expect(provider.requests).toHaveLength(1)
+
+    // The lock is released: a follow-up query runs fine.
+    for await (const _event of agent.query("third")) {
+      // drain
+    }
+    expect(provider.requests).toHaveLength(2)
+    await agent.close()
+  })
+})
+
+describe("Agent session message uuid realignment (#363)", () => {
+  test("rebuilds session messages after a compaction boundary without rotating user uuids", async () => {
+    const agent = createAgent({
+      persistSession: false,
+      tools: [],
+      provider: new CapturingProvider(),
+      model: "host/model-a",
+    })
+    await agent.getInitializationResult()
+
+    for await (const _event of agent.query("checkpoint anchor request", {
+      contextController: {
+        shouldAutoCompact: () => true,
+        async compactConversation() {
+          return {
+            compactedMessages: [
+              { role: "user", content: "[Previous conversation summary]\n\nanchor summary" },
+            ],
+            summary: "anchor summary",
+          }
+        },
+      },
+    })) {
+      // drain
+    }
+
+    const loggedUserUuid = ((agent.getMessages().find((message) => message.type === "user") as any) as { uuid: string }).uuid
+    const rebuilt = (agent as any).sessionMessages as Array<{ uuid: string; role: string; content: unknown }>
+    expect(rebuilt.map((message) => message.role)).toEqual(["user", "user", "assistant"])
+    // The rebuilt latest user message keeps its original uuid, so
+    // fileCheckpointState lookups keyed by that uuid still hit.
+    expect(rebuilt[1]!.uuid).toBe(loggedUserUuid)
+    expect(JSON.stringify(rebuilt[1]!.content)).toContain("checkpoint anchor request")
+
+    await agent.close()
+  })
+
+  test("pairs roles from the end and falls back to fresh uuids past the old list", () => {
+    const history = [
+      { role: "user", content: "one" },
+      { role: "assistant", content: "two" },
+      { role: "user", content: "three" },
+    ] as any[]
+    const previous = [
+      { uuid: "u-1", role: "user", timestamp: "t", content: "one" },
+      { uuid: "a-1", role: "assistant", timestamp: "t", content: "two" },
+    ] as any[]
+
+    const rebuilt = sessionMessagesFromHistory(history, previous)
+
+    // Trailing messages map onto the previous list; the leading extra user
+    // (e.g. a synthetic compaction summary) gets a fresh uuid.
+    expect(rebuilt.map((message) => message.uuid)).toEqual([expect.any(String), "a-1", "u-1"])
+  })
+
+  test("never hands an old uuid to a synthetic compaction summary when history shrank (#363)", () => {
+    const oldUuids = ["u-1", "u-2", "u-3", "u-4", "u-5", "u-6"]
+    const history = [
+      {
+        role: "user",
+        content: [{ type: "text", text: "checkpoint summary", _meta: { contextBlock: "compaction" } }],
+      },
+      { role: "assistant", content: "mid answer" },
+      { role: "user", content: "latest question" },
+    ] as any[]
+    const previous = [
+      ...oldUuids.map((uuid) => ({ uuid, role: "user", timestamp: "t", content: `request ${uuid}` })),
+      { uuid: "a-mid", role: "assistant", timestamp: "t", content: "older answer" },
+    ] as any[]
+
+    const rebuilt = sessionMessagesFromHistory(history, previous)
+
+    // The summary is synthetic — it must take a fresh uuid instead of stealing
+    // one from a swallowed user message (rewindFiles would otherwise restore
+    // unrelated snapshots through it).
+    const summaryUuid = rebuilt[0]!.uuid
+    for (const uuid of oldUuids) {
+      expect(summaryUuid).not.toBe(uuid)
+    }
+    // The surviving real messages keep their own uuids.
+    expect(rebuilt[1]!.uuid).toBe("a-mid")
+    expect(rebuilt[2]!.uuid).toBe("u-6")
+  })
+})

packages/sdk/src/agent.ts

@@ -80,9 +80,10 @@ type QueryInput = string | ContentBlockParam[] | SDKUserMessage
 function toSessionMessage(
   role: SessionMessage['role'],
   content: unknown,
+  uuid?: string,
 ): SessionMessage {
   return {
-    uuid: crypto.randomUUID(),
+    uuid: uuid ?? crypto.randomUUID(),
     role,
     timestamp: new Date().toISOString(),
     content,
@@ -240,10 +241,55 @@ function normalizeSessionMessageContent(
   return message.content as NormalizedMessageParam['content']
 }
 
-function sessionMessagesFromHistory(
+function isCompactionSummaryMessage(message: NormalizedMessageParam): boolean {
+  return Array.isArray(message.content)
+    && message.content.some((block: any) =>
+      block?.type === 'text' && block?._meta?.contextBlock === 'compaction')
+}
+
+export function sessionMessagesFromHistory(
   messages: NormalizedMessageParam[],
+  previous?: SessionMessage[],
 ): SessionMessage[] {
-  return messages.map((message) => toSessionMessage(message.role, message.content))
+  // Realign with the previous list so rebuilt messages keep their original
+  // uuids — fileCheckpointState is keyed by user-message uuid, and fresh
+  // uuids here would orphan every checkpoint (#363). Alignment pairs each
+  // role from the END: compaction prepends a synthetic summary user message,
+  // so only the trailing messages correspond 1:1 with what came before.
+  // Synthetic summaries never participate in pairing: when the previous list
+  // holds more same-role entries than the rebuilt history (the normal
+  // compaction shape), tail pairing would hand them a swallowed message's
+  // uuid and let rewindFiles restore unrelated snapshots.
+  const previousUuidsByRole = new Map<string, string[]>()
+  for (const message of previous ?? []) {
+    const uuids = previousUuidsByRole.get(message.role) ?? []
+    uuids.push(message.uuid)
+    previousUuidsByRole.set(message.role, uuids)
+  }
+  const indicesByRole = new Map<string, number[]>()
+  messages.forEach((message, index) => {
+    if (isCompactionSummaryMessage(message)) return
+    const indices = indicesByRole.get(message.role) ?? []
+    indices.push(index)
+    indicesByRole.set(message.role, indices)
+  })
+  const uuidByIndex = new Map<number, string>()
+  for (const [role, indices] of indicesByRole) {
+    const uuids = previousUuidsByRole.get(role) ?? []
+    const paired = Math.min(indices.length, uuids.length)
+    for (let offset = 0; offset < paired; offset++) {
+      uuidByIndex.set(
+        indices[indices.length - paired + offset]!,
+        uuids[uuids.length - paired + offset]!,
+      )
+    }
+  }
+  return messages.map((message, index) =>
+    toSessionMessage(
+      message.role as SessionMessage['role'],
+      message.content,
+      uuidByIndex.get(index),
+    ))
 }
 
 export class Agent {
@@ -265,6 +311,8 @@ export class Agent {
   private sid: string
   private abortCtrl: AbortController | null = null
   private currentEngine: QueryEngine | null = null
+  /** Synchronous in-flight marker: set at run entry, cleared when the run ends. */
+  private runLocked = false
   private hookRegistry: HookRegistry
   private loadedSettings: LoadedSettingsSource[] = []
   private loadedPlugins: LoadedPlugin[] = []
@@ -870,10 +918,25 @@ export class Agent {
   private async *runSinglePrompt(
     prompt: QueryInput,
     overrides?: Partial<AgentOptions>,
+  ): AsyncGenerator<SDKMessage, void> {
+    // A synchronous flag closes the TOCTOU window: the engine assignment sits
+    // several awaits deep, so two lazily-started queries could both pass a
+    // pure currentEngine check and fork the same session into two engines (#357).
+    if (this.runLocked || this.currentEngine) throw new Error('agent is running')
+    this.runLocked = true
+    try {
+      yield* this.runSinglePromptLocked(prompt, overrides)
+    } finally {
+      this.runLocked = false
+    }
+  }
+
+  private async *runSinglePromptLocked(
+    prompt: QueryInput,
+    overrides?: Partial<AgentOptions>,
   ): AsyncGenerator<SDKMessage, void> {
     // currentEngine (not abortCtrl, which is never cleared after a run) is
     // the accurate in-flight marker: set before the loop, cleared in finally.
-    if (this.currentEngine) throw new Error('agent is running')
     await this.setupDone
 
     // Fail fast before any listener is attached, the user message is
@@ -1117,10 +1180,10 @@ export class Agent {
       this.lastContextUsage = engine.getContextUsage()
       this.currentEngine = null
       if (compactionBoundarySeen) {
-        this.sessionMessages = sessionMessagesFromHistory(this.history)
+        this.sessionMessages = sessionMessagesFromHistory(this.history, this.sessionMessages)
       }
       if (opts.toolContinuations?.length) {
-        this.sessionMessages = sessionMessagesFromHistory(this.history)
+        this.sessionMessages = sessionMessagesFromHistory(this.history, this.sessionMessages)
       }
       persistedSessionEvent = await this.persistCurrentSession(cwd, opts)
     }