diff --git a/packages/server/src/server/agent/providers/pi-direct-agent.ts b/packages/server/src/server/agent/providers/pi-direct-agent.ts index 66f64c925..ed8ee13f6 100644 --- a/packages/server/src/server/agent/providers/pi-direct-agent.ts +++ b/packages/server/src/server/agent/providers/pi-direct-agent.ts @@ -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 { + const pendingToolCalls = new Map(); + 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, + }, + }; } } } diff --git a/packages/server/src/server/daemon-e2e/pi.real.e2e.test.ts b/packages/server/src/server/daemon-e2e/pi.real.e2e.test.ts index 63a5cac7c..8f2cd5172 100644 --- a/packages/server/src/server/daemon-e2e/pi.real.e2e.test.ts +++ b/packages/server/src/server/daemon-e2e/pi.real.e2e.test.ts @@ -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 => + event.type === "timeline", + ) + .map((event) => event.item); + + const userItems = items.filter( + (item): item is Extract => + 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 => + 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 () => {