fix(hub): resume Claude SDK sessions

2026-07-08 11:48:20 +08:00
parent 012c8a7d89
commit c73b41e1da
5 changed files with 108 additions and 50 deletions
@@ -1,47 +1,89 @@
-import { describe, expect, it } from "vitest";
-import { mkdtemp, rm, writeFile } from "node:fs/promises";
-import { tmpdir } from "node:os";
-import { join } from "node:path";
+import { describe, expect, it, vi, beforeEach } from "vitest";
 import { runAgent } from "../../src/agent/runner.js";
-import { ToolRegistry } from "../../src/agent/tools.js";
-import { readFileTool } from "../../src/agent/workspace.js";
-import { createMockModelFactory } from "../integration/helpers.js";

-describe("runAgent with AI SDK tool loop", () => {
-  it("executes SDK tools and returns the final assistant response", async () => {
-    const ws = await mkdtemp(join(tmpdir(), "hub-runner-"));
-    try {
-      await writeFile(join(ws, "lesson.txt"), "hello lesson", "utf8");
-      const tools = new ToolRegistry();
-      tools.register("read_file", readFileTool);
-      const { modelFactory, models } = createMockModelFactory([
-        {
-          toolCalls: [{ toolCallId: "call-1", toolName: "read_file", input: { path: "lesson.txt" } }],
-        },
-        { text: "done", finishReason: "stop" },
-      ]);
+const queryMock = vi.hoisted(() => vi.fn());

-      const stubPrisma = {
-        projectAgentLock: { update: async () => ({}) },
-        agentMessage: { create: async () => ({}) },
-        agentFileChange: { createMany: async () => ({}) },
-      } as unknown as import("@prisma/client").PrismaClient;
-      const result = await runAgent(modelFactory, tools, {
-        prompt: "read lesson.txt",
-        model: "mock-model",
-        project: { projectId: "p", boundChatId: "c", workspaceDir: ws },
-        systemPrompt: "system",
-        maxIterations: 5,
-        runId: "test-run",
-        sessionId: "test-session",
-        prisma: stubPrisma,
-      });
+vi.mock("@anthropic-ai/claude-agent-sdk", () => ({
+  query: queryMock,
+}));

-      console.log("DEBUG result:", result.status, result.error, result.messages.length);
-      expect(result.messages.at(-1)).toMatchObject({ role: "assistant" });
-      expect(models.get("mock-model")?.calls).toHaveLength(2);
-    } finally {
-      await rm(ws, { recursive: true, force: true });
-    }
+const stubPrisma = {
+  projectAgentLock: { update: async () => ({}) },
+} as unknown as import("@prisma/client").PrismaClient;
+
+function assistantMessage(text: string) {
+  return {
+    type: "assistant",
+    message: {
+      content: [{ type: "text", text }],
+      usage: { input_tokens: 3, output_tokens: 5 },
+    },
+  };
+}
+
+function resultMessage(sessionId: string) {
+  return {
+    type: "result",
+    subtype: "success",
+    session_id: sessionId,
+  };
+}
+
+function messages(...items: unknown[]) {
+  return (async function* () {
+    for (const item of items) yield item;
+  })();
+}
+
+describe("runAgent", () => {
+  beforeEach(() => {
+    queryMock.mockReset();
+  });
+
+  it("passes the previous Claude SDK session id through resume", async () => {
+    queryMock.mockReturnValue(messages(assistantMessage("ok"), resultMessage("sdk-session-new")));
+
+    const result = await runAgent({
+      prompt: "再发一次",
+      model: "anthropic/claude-sonnet-5",
+      project: { projectId: "p", boundChatId: "c", workspaceDir: "/tmp/ws" },
+      systemPrompt: undefined,
+      resumeSessionId: "sdk-session-old",
+      runId: "run-1",
+      sessionId: "hub-session-1",
+      prisma: stubPrisma,
+    });
+
+    expect(queryMock).toHaveBeenCalledWith({
+      prompt: "再发一次",
+      options: expect.objectContaining({
+        cwd: "/tmp/ws",
+        model: "anthropic/claude-sonnet-5",
+        resume: "sdk-session-old",
+      }),
+    });
+    expect(result).toMatchObject({
+      status: "completed",
+      text: "ok",
+      sdkSessionId: "sdk-session-new",
+      usage: { inputTokens: 3, outputTokens: 5 },
+    });
+  });
+
+  it("does not send resume for a fresh Hub session", async () => {
+    queryMock.mockReturnValue(messages(assistantMessage("fresh"), resultMessage("sdk-session-1")));
+
+    await runAgent({
+      prompt: "你好",
+      model: undefined,
+      project: { projectId: "p", boundChatId: "c", workspaceDir: "/tmp/ws" },
+      systemPrompt: undefined,
+      runId: "run-1",
+      sessionId: "hub-session-1",
+      prisma: stubPrisma,
+    });
+
+    const call = queryMock.mock.calls[0]?.[0] as { options?: Record<string, unknown> } | undefined;
+    expect(call?.options).not.toHaveProperty("resume");
  });
 });