瀏覽代碼

feat: restore committed transcript messages

fanghanjun 6 月之前
父節點
當前提交
2c59502260

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

@@ -0,0 +1,213 @@
+import { describe, expect, it, beforeEach, afterEach } from "vitest";
+import { mkdtemp, rm } from "node:fs/promises";
+import { join } from "node:path";
+import { tmpdir } from "node:os";
+import { appendTranscriptEvent } from "../interaction/session-transcript.js";
+import { restoreAgentMessagesFromTranscript } from "../interaction/session-transcript-restore.js";
+import type { MessageEvent } from "../interaction/session-transcript-schema.js";
+
+const usage = {
+  input: 0,
+  output: 0,
+  cacheRead: 0,
+  cacheWrite: 0,
+  totalTokens: 0,
+  cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+};
+
+describe("session transcript restore", () => {
+  let projectRoot: string;
+
+  beforeEach(async () => {
+    projectRoot = await mkdtemp(join(tmpdir(), "inkos-restore-"));
+  });
+
+  afterEach(async () => {
+    await rm(projectRoot, { recursive: true, force: true });
+  });
+
+  it("只恢复已 committed request 内的 message", async () => {
+    await appendTranscriptEvent(projectRoot, {
+      type: "request_started",
+      version: 1,
+      sessionId: "s1",
+      requestId: "r1",
+      seq: 1,
+      timestamp: 1,
+      input: "hi",
+    });
+    await appendTranscriptEvent(projectRoot, {
+      type: "message",
+      version: 1,
+      sessionId: "s1",
+      requestId: "r1",
+      uuid: "u1",
+      parentUuid: null,
+      seq: 2,
+      role: "user",
+      timestamp: 2,
+      message: { role: "user", content: "hi", timestamp: 2 },
+    } as MessageEvent);
+    await appendTranscriptEvent(projectRoot, {
+      type: "request_committed",
+      version: 1,
+      sessionId: "s1",
+      requestId: "r1",
+      seq: 3,
+      timestamp: 3,
+    });
+    await appendTranscriptEvent(projectRoot, {
+      type: "request_started",
+      version: 1,
+      sessionId: "s1",
+      requestId: "r2",
+      seq: 4,
+      timestamp: 4,
+      input: "lost",
+    });
+    await appendTranscriptEvent(projectRoot, {
+      type: "message",
+      version: 1,
+      sessionId: "s1",
+      requestId: "r2",
+      uuid: "u2",
+      parentUuid: "u1",
+      seq: 5,
+      role: "user",
+      timestamp: 5,
+      message: { role: "user", content: "lost", timestamp: 5 },
+    } as MessageEvent);
+
+    const restored = await restoreAgentMessagesFromTranscript(projectRoot, "s1");
+
+    expect(restored).toHaveLength(1);
+    expect(restored[0]).toMatchObject({ role: "user", content: "hi" });
+  });
+
+  it("保留 committed toolResult 和 assistant thinking signature", async () => {
+    await appendTranscriptEvent(projectRoot, {
+      type: "request_started",
+      version: 1,
+      sessionId: "s1",
+      requestId: "r1",
+      seq: 1,
+      timestamp: 1,
+      input: "tool",
+    });
+    await appendTranscriptEvent(projectRoot, {
+      type: "message",
+      version: 1,
+      sessionId: "s1",
+      requestId: "r1",
+      uuid: "a1",
+      parentUuid: null,
+      seq: 2,
+      role: "assistant",
+      timestamp: 2,
+      toolCallId: "tool-1",
+      message: {
+        role: "assistant",
+        content: [
+          { type: "thinking", thinking: "需要查资料", signature: "sig" },
+          { type: "toolCall", id: "tool-1", name: "read", arguments: { path: "a.md" } },
+        ],
+        api: "anthropic-messages",
+        provider: "anthropic",
+        model: "claude",
+        usage,
+        stopReason: "tool_use",
+        timestamp: 2,
+      },
+    } as MessageEvent);
+    await appendTranscriptEvent(projectRoot, {
+      type: "message",
+      version: 1,
+      sessionId: "s1",
+      requestId: "r1",
+      uuid: "t1",
+      parentUuid: "a1",
+      seq: 3,
+      role: "toolResult",
+      timestamp: 3,
+      toolCallId: "tool-1",
+      sourceToolAssistantUuid: "a1",
+      message: {
+        role: "toolResult",
+        toolCallId: "tool-1",
+        toolName: "read",
+        content: [{ type: "text", text: "资料" }],
+        details: { path: "a.md" },
+        isError: false,
+        timestamp: 3,
+      },
+    } as MessageEvent);
+    await appendTranscriptEvent(projectRoot, {
+      type: "request_committed",
+      version: 1,
+      sessionId: "s1",
+      requestId: "r1",
+      seq: 4,
+      timestamp: 4,
+    });
+
+    const restored = await restoreAgentMessagesFromTranscript(projectRoot, "s1");
+
+    expect(restored).toHaveLength(2);
+    expect(restored[0]).toMatchObject({
+      role: "assistant",
+      content: [
+        { type: "thinking", thinking: "需要查资料", signature: "sig" },
+        { type: "toolCall", id: "tool-1" },
+      ],
+    });
+    expect(restored[1]).toMatchObject({ role: "toolResult", toolCallId: "tool-1", toolName: "read" });
+  });
+
+  it("移除最后 assistant message 的 trailing thinking block", async () => {
+    await appendTranscriptEvent(projectRoot, {
+      type: "request_started",
+      version: 1,
+      sessionId: "s1",
+      requestId: "r1",
+      seq: 1,
+      timestamp: 1,
+      input: "hi",
+    });
+    await appendTranscriptEvent(projectRoot, {
+      type: "message",
+      version: 1,
+      sessionId: "s1",
+      requestId: "r1",
+      uuid: "a1",
+      parentUuid: null,
+      seq: 2,
+      role: "assistant",
+      timestamp: 2,
+      message: {
+        role: "assistant",
+        content: [
+          { type: "text", text: "回答" },
+          { type: "thinking", thinking: "尾部", signature: "sig" },
+        ],
+        api: "anthropic-messages",
+        provider: "anthropic",
+        model: "claude",
+        usage,
+        stopReason: "stop",
+        timestamp: 2,
+      },
+    } as MessageEvent);
+    await appendTranscriptEvent(projectRoot, {
+      type: "request_committed",
+      version: 1,
+      sessionId: "s1",
+      requestId: "r1",
+      seq: 3,
+      timestamp: 3,
+    });
+
+    const restored = await restoreAgentMessagesFromTranscript(projectRoot, "s1");
+
+    expect((restored[0] as any).content).toEqual([{ type: "text", text: "回答" }]);
+  });
+});

+ 104 - 0
packages/core/src/interaction/session-transcript-restore.ts

@@ -0,0 +1,104 @@
+import type { AgentMessage } from "@mariozechner/pi-agent-core";
+import { readTranscriptEvents } from "./session-transcript.js";
+import type { MessageEvent, TranscriptEvent } from "./session-transcript-schema.js";
+
+function isObject(value: unknown): value is Record<string, unknown> {
+  return !!value && typeof value === "object";
+}
+
+function contentBlocks(message: Record<string, unknown>): unknown[] {
+  return Array.isArray(message.content) ? message.content : [];
+}
+
+function hasTextContent(message: Record<string, unknown>): boolean {
+  return contentBlocks(message).some(
+    (block) =>
+      isObject(block) &&
+      block.type === "text" &&
+      typeof block.text === "string" &&
+      block.text.length > 0,
+  );
+}
+
+function hasToolCallContent(message: Record<string, unknown>): boolean {
+  return contentBlocks(message).some(
+    (block) => isObject(block) && block.type === "toolCall" && typeof block.id === "string",
+  );
+}
+
+function toolCallIds(message: Record<string, unknown>): string[] {
+  return contentBlocks(message)
+    .filter(
+      (block): block is Record<string, unknown> =>
+        isObject(block) && block.type === "toolCall" && typeof block.id === "string",
+    )
+    .map((block) => block.id as string);
+}
+
+function isThinkingBlock(block: unknown): boolean {
+  return isObject(block) && (block.type === "thinking" || block.type === "redacted_thinking");
+}
+
+function removeTrailingThinking(message: AgentMessage): AgentMessage {
+  if (!isObject(message) || message.role !== "assistant" || !Array.isArray(message.content)) {
+    return message;
+  }
+
+  const content = [...message.content];
+  while (content.length > 0 && isThinkingBlock(content[content.length - 1])) {
+    content.pop();
+  }
+
+  if (content.length === message.content.length) return message;
+  return { ...message, content } as AgentMessage;
+}
+
+export function cleanRestoredAgentMessages(messages: AgentMessage[]): AgentMessage[] {
+  const availableToolCalls = new Set<string>();
+  for (const message of messages) {
+    if (isObject(message) && message.role === "assistant") {
+      for (const id of toolCallIds(message)) availableToolCalls.add(id);
+    }
+  }
+
+  const cleaned = messages.filter((message) => {
+    if (!isObject(message)) return false;
+    if (message.role === "toolResult") {
+      return typeof message.toolCallId === "string" && availableToolCalls.has(message.toolCallId);
+    }
+    if (message.role === "assistant") {
+      return hasTextContent(message) || hasToolCallContent(message);
+    }
+    return message.role === "user" || message.role === "system";
+  });
+
+  if (cleaned.length === 0) return cleaned;
+  const last = cleaned[cleaned.length - 1];
+  if (isObject(last) && last.role === "assistant") {
+    cleaned[cleaned.length - 1] = removeTrailingThinking(last);
+  }
+
+  return cleaned;
+}
+
+export function committedMessageEvents(events: TranscriptEvent[]): MessageEvent[] {
+  const committed = new Set(
+    events
+      .filter((event) => event.type === "request_committed")
+      .map((event) => event.requestId),
+  );
+
+  return events
+    .filter((event): event is MessageEvent => event.type === "message" && committed.has(event.requestId))
+    .sort((a, b) => a.seq - b.seq);
+}
+
+export async function restoreAgentMessagesFromTranscript(
+  projectRoot: string,
+  sessionId: string,
+): Promise<AgentMessage[]> {
+  const events = await readTranscriptEvents(projectRoot, sessionId);
+  return cleanRestoredAgentMessages(
+    committedMessageEvents(events).map((event) => event.message as AgentMessage),
+  );
+}