feat(server): replay pi tool calls in streamHistory for resume

Pi's streamHistory only emitted text, so resuming a session lost tool calls, tool results, and bash executions from the replayed timeline. Iterate session.messages to also emit running/completed/failed tool_call items and bashExecution shell calls, and tag user messages with a stable messageId so hydrateTimelineFromLegacyProviderHistory can dedupe. Cover it with a real OpenRouter-backed e2e test that runs a bash tool call, resumes via PiDirectAgentClient.resumeSession, and asserts the replayed timeline contains the user message, assistant message, and completed shell tool_call with the expected output.
2026-04-19 23:17:40 +07:00
parent 75f83b0b9b
commit 8165de7b06
2 changed files with 199 additions and 46 deletions
@@ -22,14 +22,7 @@ import {
  type WriteToolInput,
 } from "@mariozechner/pi-coding-agent";
 import type { ThinkingLevel } from "@mariozechner/pi-agent-core";
-import type {
-  Api,
-  ImageContent,
-  Model,
-  TextContent,
-  ThinkingContent,
-  ToolCall,
-} from "@mariozechner/pi-ai";
+import type { Api, ImageContent, Model, TextContent } from "@mariozechner/pi-ai";
 import { z } from "zod";

 import type {
@@ -733,16 +726,6 @@ function getUserMessageText(content: string | (TextContent | ImageContent)[]): s
  return textParts.join("\n\n");
 }

-function getAssistantContentText(content: TextContent | ThinkingContent | ToolCall): string | null {
-  if (content.type === "text") {
-    return content.text || null;
-  }
-  if (content.type === "thinking") {
-    return content.thinking || null;
-  }
-  return null;
-}
-
 function parsePersistenceMetadata(metadata: AgentMetadata | undefined): PiPersistenceMetadata {
  const parsed = PiPersistenceMetadataSchema.safeParse(metadata);
  if (parsed.success) {
@@ -1112,6 +1095,9 @@ export class PiDirectAgentSession implements AgentSession {
  }

  async *streamHistory(): AsyncGenerator<AgentStreamEvent> {
+    const pendingToolCalls = new Map<string, PiTrackedToolCall>();
+    let userIndex = 0;
+
    for (const message of this.session.messages) {
      if (message.role === "user") {
        const text = getUserMessageText(message.content);
@@ -1119,38 +1105,120 @@ export class PiDirectAgentSession implements AgentSession {
          yield {
            type: "timeline",
            provider: PI_PROVIDER,
-            item: { type: "user_message", text },
+            item: {
+              type: "user_message",
+              text,
+              messageId: `pi-user-${userIndex}`,
+            },
+          };
+        }
+        userIndex += 1;
+        continue;
+      }
+
+      if (message.role === "assistant") {
+        for (const content of message.content) {
+          if (content.type === "text") {
+            if (content.text) {
+              yield {
+                type: "timeline",
+                provider: PI_PROVIDER,
+                item: { type: "assistant_message", text: content.text },
+              };
+            }
+            continue;
+          }
+
+          if (content.type === "thinking") {
+            if (content.thinking) {
+              yield {
+                type: "timeline",
+                provider: PI_PROVIDER,
+                item: { type: "reasoning", text: content.thinking },
+              };
+            }
+            continue;
+          }
+
+          if (content.type === "toolCall") {
+            const tracked = parseToolArgs(content.name, content.arguments);
+            pendingToolCalls.set(content.id, tracked);
+            yield {
+              type: "timeline",
+              provider: PI_PROVIDER,
+              item: {
+                type: "tool_call",
+                callId: content.id,
+                name: tracked.toolName,
+                status: "running",
+                detail: mapToolDetail(tracked, null),
+                error: null,
+              },
+            };
+          }
+        }
+        continue;
+      }
+
+      if (message.role === "toolResult") {
+        const tracked =
+          pendingToolCalls.get(message.toolCallId) ?? parseToolArgs(message.toolName, null);
+        pendingToolCalls.delete(message.toolCallId);
+        const result = parseToolResult({ content: message.content });
+        const detail = mapToolDetail(tracked, result);
+
+        if (message.isError) {
+          const errorText = extractTextFromToolResult(result) ?? "Tool call failed";
+          yield {
+            type: "timeline",
+            provider: PI_PROVIDER,
+            item: {
+              type: "tool_call",
+              callId: message.toolCallId,
+              name: tracked.toolName,
+              status: "failed",
+              detail,
+              error: errorText,
+            },
+          };
+        } else {
+          yield {
+            type: "timeline",
+            provider: PI_PROVIDER,
+            item: {
+              type: "tool_call",
+              callId: message.toolCallId,
+              name: tracked.toolName,
+              status: "completed",
+              detail,
+              error: null,
+            },
          };
        }
        continue;
      }

-      if (message.role !== "assistant") {
-        continue;
-      }
-
-      for (const content of message.content) {
-        const text = getAssistantContentText(content);
-        if (!text) {
-          continue;
-        }
-
-        if (content.type === "text") {
-          yield {
-            type: "timeline",
-            provider: PI_PROVIDER,
-            item: { type: "assistant_message", text },
-          };
-          continue;
-        }
-
-        if (content.type === "thinking") {
-          yield {
-            type: "timeline",
-            provider: PI_PROVIDER,
-            item: { type: "reasoning", text },
-          };
-        }
+      if (message.role === "bashExecution") {
+        const callId = `pi-bash-${message.timestamp}`;
+        const exitCode = message.exitCode ?? null;
+        const detail: ToolCallDetail = {
+          type: "shell",
+          command: message.command,
+          output: message.output,
+          exitCode,
+        };
+        yield {
+          type: "timeline",
+          provider: PI_PROVIDER,
+          item: {
+            type: "tool_call",
+            callId,
+            name: "bash",
+            status: message.cancelled ? "canceled" : "completed",
+            detail,
+            error: null,
+          },
+        };
      }
    }
  }
@@ -5,7 +5,11 @@ import { randomUUID } from "node:crypto";
 import { beforeAll, beforeEach, describe, expect, test } from "vitest";
 import pino from "pino";

-import type { AgentPersistenceHandle, AgentTimelineItem } from "../agent/agent-sdk-types.js";
+import type {
+  AgentPersistenceHandle,
+  AgentStreamEvent,
+  AgentTimelineItem,
+} from "../agent/agent-sdk-types.js";
 import { PiDirectAgentClient } from "../agent/providers/pi-direct-agent.js";
 import { DaemonClient } from "../test-utils/daemon-client.js";
 import { createTestPaseoDaemon } from "../test-utils/paseo-daemon.js";
@@ -362,6 +366,87 @@ describe(
      PI_TEST_TIMEOUT_MS,
    );

+    test(
+      "streamHistory replays user/assistant/tool_call timeline after resume",
+      async () => {
+        const cwd = tmpCwd("pi-history-prime-");
+        const marker = "HISTORY_PRIME_MARKER_4242";
+
+        const piClient = createPiClient();
+        const session = await piClient.createSession({
+          provider: "pi",
+          cwd,
+          model: PI_REAL_TEST_MODEL,
+        });
+
+        let handle: AgentPersistenceHandle | null = null;
+
+        try {
+          const result = await session.run(
+            `Use the bash tool to run this exact bash command: echo ${marker}`,
+          );
+
+          const hadShellToolCall = result.timeline.some(
+            (item) =>
+              item.type === "tool_call" &&
+              item.status === "completed" &&
+              item.detail.type === "shell" &&
+              item.detail.command.includes(`echo ${marker}`),
+          );
+          expect(hadShellToolCall).toBe(true);
+
+          handle = session.describePersistence();
+          expect(handle).toBeTruthy();
+        } finally {
+          await session.close();
+        }
+
+        const resumed = await piClient.resumeSession(handle as AgentPersistenceHandle);
+
+        try {
+          const events: AgentStreamEvent[] = [];
+          for await (const event of resumed.streamHistory()) {
+            events.push(event);
+          }
+
+          const items = events
+            .filter(
+              (event): event is Extract<AgentStreamEvent, { type: "timeline" }> =>
+                event.type === "timeline",
+            )
+            .map((event) => event.item);
+
+          const userItems = items.filter(
+            (item): item is Extract<AgentTimelineItem, { type: "user_message" }> =>
+              item.type === "user_message",
+          );
+          expect(userItems.length).toBeGreaterThan(0);
+          expect(userItems.some((item) => item.text.includes(marker))).toBe(true);
+          expect(userItems.every((item) => typeof item.messageId === "string")).toBe(true);
+
+          const assistantItems = items.filter(
+            (item): item is Extract<AgentTimelineItem, { type: "assistant_message" }> =>
+              item.type === "assistant_message",
+          );
+          expect(assistantItems.length).toBeGreaterThan(0);
+
+          const toolCallItem = findCompletedToolCall(
+            items,
+            (item) =>
+              item.detail.type === "shell" && item.detail.command.includes(`echo ${marker}`),
+          );
+          expect(toolCallItem).toBeDefined();
+          if (toolCallItem?.detail.type === "shell") {
+            expect(toolCallItem.detail.output).toContain(marker);
+          }
+        } finally {
+          await resumed.close();
+          rmSync(cwd, { recursive: true, force: true });
+        }
+      },
+      PI_TEST_TIMEOUT_MS,
+    );
+
    test(
      "PiDirectAgentClient.listModels returns non-empty Pi model definitions",
      async () => {