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.
This commit is contained in:
@@ -22,14 +22,7 @@ import {
|
|||||||
type WriteToolInput,
|
type WriteToolInput,
|
||||||
} from "@mariozechner/pi-coding-agent";
|
} from "@mariozechner/pi-coding-agent";
|
||||||
import type { ThinkingLevel } from "@mariozechner/pi-agent-core";
|
import type { ThinkingLevel } from "@mariozechner/pi-agent-core";
|
||||||
import type {
|
import type { Api, ImageContent, Model, TextContent } from "@mariozechner/pi-ai";
|
||||||
Api,
|
|
||||||
ImageContent,
|
|
||||||
Model,
|
|
||||||
TextContent,
|
|
||||||
ThinkingContent,
|
|
||||||
ToolCall,
|
|
||||||
} from "@mariozechner/pi-ai";
|
|
||||||
import { z } from "zod";
|
import { z } from "zod";
|
||||||
|
|
||||||
import type {
|
import type {
|
||||||
@@ -733,16 +726,6 @@ function getUserMessageText(content: string | (TextContent | ImageContent)[]): s
|
|||||||
return textParts.join("\n\n");
|
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 {
|
function parsePersistenceMetadata(metadata: AgentMetadata | undefined): PiPersistenceMetadata {
|
||||||
const parsed = PiPersistenceMetadataSchema.safeParse(metadata);
|
const parsed = PiPersistenceMetadataSchema.safeParse(metadata);
|
||||||
if (parsed.success) {
|
if (parsed.success) {
|
||||||
@@ -1112,6 +1095,9 @@ export class PiDirectAgentSession implements AgentSession {
|
|||||||
}
|
}
|
||||||
|
|
||||||
async *streamHistory(): AsyncGenerator<AgentStreamEvent> {
|
async *streamHistory(): AsyncGenerator<AgentStreamEvent> {
|
||||||
|
const pendingToolCalls = new Map<string, PiTrackedToolCall>();
|
||||||
|
let userIndex = 0;
|
||||||
|
|
||||||
for (const message of this.session.messages) {
|
for (const message of this.session.messages) {
|
||||||
if (message.role === "user") {
|
if (message.role === "user") {
|
||||||
const text = getUserMessageText(message.content);
|
const text = getUserMessageText(message.content);
|
||||||
@@ -1119,38 +1105,120 @@ export class PiDirectAgentSession implements AgentSession {
|
|||||||
yield {
|
yield {
|
||||||
type: "timeline",
|
type: "timeline",
|
||||||
provider: PI_PROVIDER,
|
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;
|
continue;
|
||||||
}
|
}
|
||||||
|
|
||||||
if (message.role !== "assistant") {
|
if (message.role === "bashExecution") {
|
||||||
continue;
|
const callId = `pi-bash-${message.timestamp}`;
|
||||||
}
|
const exitCode = message.exitCode ?? null;
|
||||||
|
const detail: ToolCallDetail = {
|
||||||
for (const content of message.content) {
|
type: "shell",
|
||||||
const text = getAssistantContentText(content);
|
command: message.command,
|
||||||
if (!text) {
|
output: message.output,
|
||||||
continue;
|
exitCode,
|
||||||
}
|
};
|
||||||
|
yield {
|
||||||
if (content.type === "text") {
|
type: "timeline",
|
||||||
yield {
|
provider: PI_PROVIDER,
|
||||||
type: "timeline",
|
item: {
|
||||||
provider: PI_PROVIDER,
|
type: "tool_call",
|
||||||
item: { type: "assistant_message", text },
|
callId,
|
||||||
};
|
name: "bash",
|
||||||
continue;
|
status: message.cancelled ? "canceled" : "completed",
|
||||||
}
|
detail,
|
||||||
|
error: null,
|
||||||
if (content.type === "thinking") {
|
},
|
||||||
yield {
|
};
|
||||||
type: "timeline",
|
|
||||||
provider: PI_PROVIDER,
|
|
||||||
item: { type: "reasoning", text },
|
|
||||||
};
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -5,7 +5,11 @@ import { randomUUID } from "node:crypto";
|
|||||||
import { beforeAll, beforeEach, describe, expect, test } from "vitest";
|
import { beforeAll, beforeEach, describe, expect, test } from "vitest";
|
||||||
import pino from "pino";
|
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 { PiDirectAgentClient } from "../agent/providers/pi-direct-agent.js";
|
||||||
import { DaemonClient } from "../test-utils/daemon-client.js";
|
import { DaemonClient } from "../test-utils/daemon-client.js";
|
||||||
import { createTestPaseoDaemon } from "../test-utils/paseo-daemon.js";
|
import { createTestPaseoDaemon } from "../test-utils/paseo-daemon.js";
|
||||||
@@ -362,6 +366,87 @@ describe(
|
|||||||
PI_TEST_TIMEOUT_MS,
|
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(
|
test(
|
||||||
"PiDirectAgentClient.listModels returns non-empty Pi model definitions",
|
"PiDirectAgentClient.listModels returns non-empty Pi model definitions",
|
||||||
async () => {
|
async () => {
|
||||||
|
|||||||
Reference in New Issue
Block a user