diff --git a/packages/server/src/server/agent/providers/pi/agent.test.ts b/packages/server/src/server/agent/providers/pi/agent.test.ts index 88a3f6cb5..daf25470c 100644 --- a/packages/server/src/server/agent/providers/pi/agent.test.ts +++ b/packages/server/src/server/agent/providers/pi/agent.test.ts @@ -57,14 +57,15 @@ function readUtf8File(pathname: string): string { } } -async function applyPaseoExtensionSystemPrompt( +type PaseoExtensionListener = (event: unknown, context?: unknown) => unknown; + +async function loadPaseoExtensionListeners( extensionPath: string, - systemPrompt: string, -): Promise { - const listeners = new Map unknown>(); +): Promise> { + const listeners = new Map(); const extension = (await import(pathToFileURL(extensionPath).href)) as { default: (piApi: { - on: (event: string, listener: (event: { systemPrompt: string }) => unknown) => void; + on: (event: string, listener: PaseoExtensionListener) => void; registerCommand: () => void; }) => void; }; @@ -72,6 +73,14 @@ async function applyPaseoExtensionSystemPrompt( on: (event, listener) => listeners.set(event, listener), registerCommand: () => undefined, }); + return listeners; +} + +async function applyPaseoExtensionSystemPrompt( + extensionPath: string, + systemPrompt: string, +): Promise { + const listeners = await loadPaseoExtensionListeners(extensionPath); const result = await listeners.get("before_agent_start")?.({ systemPrompt }); return (result as { systemPrompt?: string } | undefined)?.systemPrompt; } @@ -594,15 +603,15 @@ describe("PiRpcAgentSession", () => { ]); }); - test("emits live user messages with captured Pi tree entry ids", async () => { + test("emits live user messages with submitted Pi tree entry ids", async () => { const { pi, session, events } = await createSession(); const fakeSession = pi.latestSession(); - fakeSession.capturedUserEntries = [{ id: "entry-user-1", parentId: null, text: "hello" }]; await session.startTurn("hello"); - fakeSession.emit({ - type: "message_end", - message: { role: "user", content: "hello" }, + fakeSession.finishSubmittedUserMessage({ + id: "entry-user-1", + parentId: null, + text: "hello", }); await events.nextTimelineEvent(); @@ -612,6 +621,38 @@ describe("PiRpcAgentSession", () => { ]); }); + test("uses the Pi entry attached to a submitted prompt after resuming old history", async () => { + const pi = new FakePi(); + const client = createClient(pi); + const session = (await client.resumeSession({ + provider: "pi", + sessionId: "pi-session-1", + nativeHandle: "/tmp/native-pi-session", + metadata: { cwd: "/workspace/project" }, + })) as PiRpcAgentSession; + const events = new SessionEvents(session); + const fakeSession = pi.latestSession(); + fakeSession.capturedUserEntries = [{ id: "entry-old", parentId: null, text: "old prompt" }]; + + await session.startTurn("new prompt", { clientMessageId: "client-new" }); + fakeSession.finishSubmittedUserMessage({ + id: "entry-new", + parentId: "entry-old-assistant", + text: "new prompt", + }); + + await events.nextTimelineEvent(); + + expect(events.timelineItems()).toEqual([ + { + type: "user_message", + text: "new prompt", + messageId: "entry-new", + clientMessageId: "client-new", + }, + ]); + }); + test("surfaces Pi extension command messages and completes when no agent turn starts", async () => { const { pi, session, events } = await createSession(); const fakeSession = pi.latestSession(); @@ -734,6 +775,52 @@ describe("PiRpcAgentSession", () => { ]); }); + test("reports the persisted Pi entry attached to the submitted message", async () => { + const pi = new FakePi(); + const client = createClient(pi); + const session = await client.createSession(createConfig()); + const extensionPath = pi.recordedLaunches[0]?.extensionPaths[0]; + expect(extensionPath).toBeDefined(); + const listeners = await loadPaseoExtensionListeners(extensionPath!); + const submittedMessage = { role: "user", content: "new prompt" }; + const entries: Array<{ + type: string; + id: string; + parentId: string | null; + message: { role: string; content: string }; + }> = [ + { + type: "message", + id: "entry-old", + parentId: null, + message: { role: "user", content: "old prompt" }, + }, + ]; + const notifications: string[] = []; + const context = { + sessionManager: { getEntries: () => entries }, + ui: { notify: (message: string) => notifications.push(message) }, + }; + + await listeners.get("message_end")?.({ message: submittedMessage }, context); + entries.push({ + type: "message", + id: "entry-new", + parentId: "entry-old-assistant", + message: submittedMessage, + }); + await listeners.get("message_start")?.( + { message: { role: "assistant", content: [] } }, + context, + ); + + expect(notifications).toEqual([ + 'PASEO_SUBMITTED_USER_ENTRY {"entry":{"id":"entry-new","parentId":"entry-old-assistant","text":"new prompt"}}', + ]); + + await session.close(); + }); + test("appends agent and daemon prompts after Pi's discovered system prompt", async () => { const pi = new FakePi(); const client = createClient(pi); diff --git a/packages/server/src/server/agent/providers/pi/agent.ts b/packages/server/src/server/agent/providers/pi/agent.ts index f1f55d600..66237340d 100644 --- a/packages/server/src/server/agent/providers/pi/agent.ts +++ b/packages/server/src/server/agent/providers/pi/agent.ts @@ -92,6 +92,7 @@ const PI_CATALOG_REQUEST_TIMEOUT_MS = 120_000; const PASEO_PI_TREE_EXTENSION_COMMAND = "paseo_tree"; const PASEO_PI_CAPTURE_EXTENSION_COMMAND = "paseo_capture_entries"; const PASEO_PI_ENTRY_CAPTURE_MARKER = "PASEO_ENTRY_CAPTURE"; +const PASEO_PI_SUBMITTED_USER_ENTRY_MARKER = "PASEO_SUBMITTED_USER_ENTRY"; const PASEO_PI_COMMAND_RESULT_MARKER = "PASEO_COMMAND_RESULT"; const DEFAULT_PI_EXTENSION_RESULT_TIMEOUT_MS = 30_000; const QUESTION_RESPONSE_HEADER = "Response"; @@ -256,12 +257,6 @@ interface PiCapturedEntry extends PiCapturedUserMessageEntry { parentId: string | null; } -interface PendingPiUserMessage { - text: string; - turnId: string | undefined; - clientMessageId?: string; -} - interface PendingExtensionResult { resolve: (value: unknown) => void; reject: (error: Error) => void; @@ -653,11 +648,15 @@ function createPiPaseoExtensionFile(systemPrompt?: string): PiTempFile { return ctx.sessionManager .getEntries() .filter((entry) => entry.type === "message" && entry.message?.role === "user") - .map((entry) => ({ + .map(toCapturedUserEntry); + } + + function toCapturedUserEntry(entry) { + return { id: entry.id, parentId: entry.parentId ?? null, text: readTextContent(entry.message.content), - })); + }; } function emitEntryCapture(ctx, reason, requestId) { @@ -676,6 +675,30 @@ function createPiPaseoExtensionFile(systemPrompt?: string): PiTempFile { } export default function paseoIntegration(pi) { + const submittedUserMessages = []; + + function emitSubmittedUserEntries(ctx) { + const entries = ctx.sessionManager.getEntries(); + for (let index = 0; index < submittedUserMessages.length; index += 1) { + const message = submittedUserMessages[index]; + // Pi assigns the entry ID after message_end, then persists this same message object. + // Reference equality preserves the exact association even when another extension edits it. + const entry = entries.find( + (candidate) => candidate.type === "message" && candidate.message === message, + ); + if (!entry) { + continue; + } + submittedUserMessages.splice(index, 1); + index -= 1; + ctx.ui.notify( + "${PASEO_PI_SUBMITTED_USER_ENTRY_MARKER} " + + JSON.stringify({ entry: toCapturedUserEntry(entry) }), + "info", + ); + } + } + ${ systemPrompt ? `pi.on("before_agent_start", async (event) => ({ @@ -688,7 +711,20 @@ function createPiPaseoExtensionFile(systemPrompt?: string): PiTempFile { emitEntryCapture(ctx, "session_start"); }); + pi.on("message_end", async (event) => { + if (event.message?.role === "user") { + submittedUserMessages.push(event.message); + } + }); + + pi.on("message_start", async (event, ctx) => { + if (event.message?.role === "assistant") { + emitSubmittedUserEntries(ctx); + } + }); + pi.on("turn_end", async (_event, ctx) => { + emitSubmittedUserEntries(ctx); emitEntryCapture(ctx, "turn_end"); }); @@ -1197,8 +1233,6 @@ export class PiRpcAgentSession implements AgentSession { currentLeafOverrideId: string | null | undefined; private readonly capturedUserEntries: PiCapturedEntry[] = []; private readonly capturedUserEntriesById = new Map(); - private readonly seenUserEntryIds = new Set(); - private readonly pendingUserMessages: PendingPiUserMessage[] = []; private readonly pendingExtensionResults = new Map(); private outOfBandCompactionEmit: ((event: AgentStreamEvent) => void) | null = null; private outOfBandCompactionStarted = false; @@ -1776,42 +1810,34 @@ export class PiRpcAgentSession implements AgentSession { } private recordCapturedUserEntries(entries: PiCapturedEntry[]): void { - const previouslySeenEntryIds = new Set(this.seenUserEntryIds); this.capturedUserEntries.splice(0, this.capturedUserEntries.length, ...entries); this.capturedUserEntriesById.clear(); for (const entry of entries) { this.capturedUserEntriesById.set(entry.id, entry); } - this.flushPendingUserMessages(previouslySeenEntryIds); - for (const entry of entries) { - this.seenUserEntryIds.add(entry.id); - } } - private flushPendingUserMessages(previouslySeenEntryIds: Set): void { - for (let index = 0; index < this.pendingUserMessages.length; index += 1) { - const pending = this.pendingUserMessages[index]!; - const entry = this.capturedUserEntries.find( - (candidate) => !previouslySeenEntryIds.has(candidate.id), - ); - if (!entry) { - continue; - } - previouslySeenEntryIds.add(entry.id); - this.pendingUserMessages.splice(index, 1); - index -= 1; - this.emit({ - type: "timeline", - provider: this.provider, - turnId: pending.turnId, - item: { - type: "user_message", - text: pending.text, - messageId: entry.id, - ...(pending.clientMessageId ? { clientMessageId: pending.clientMessageId } : {}), - }, - }); + private handleSubmittedUserEntryMarker(message: string): boolean { + const payload = parseExtensionMarkerPayload(message, PASEO_PI_SUBMITTED_USER_ENTRY_MARKER); + if (!payload) { + return false; } + const [entry] = parseCapturedEntries([payload.entry]); + if (!entry) { + return true; + } + this.emit({ + type: "timeline", + provider: this.provider, + turnId: this.currentTurnIdForEvent(), + item: { + type: "user_message", + text: entry.text, + messageId: entry.id, + ...(this.activeClientMessageId ? { clientMessageId: this.activeClientMessageId } : {}), + }, + }); + return true; } private handleEntryCaptureMarker(message: string): boolean { @@ -1849,7 +1875,11 @@ export class PiRpcAgentSession implements AgentSession { ): void { const message = optionalString(event.message); if (event.method === "notify" && message) { - if (this.handleEntryCaptureMarker(message) || this.handleCommandResultMarker(message)) { + if ( + this.handleSubmittedUserEntryMarker(message) || + this.handleEntryCaptureMarker(message) || + this.handleCommandResultMarker(message) + ) { return; } this.bufferNoTurnOutput(message); @@ -2176,28 +2206,6 @@ export class PiRpcAgentSession implements AgentSession { this.completeTurn(turnId, []); return; } - - if (event.message.role !== "user") { - return; - } - const text = getUserMessageText(event.message.content); - if (!text) { - return; - } - this.pendingUserMessages.push({ - text, - turnId, - ...(this.activeClientMessageId ? { clientMessageId: this.activeClientMessageId } : {}), - }); - void this.requestEntryCapture("message_end").catch((error: unknown) => { - const message = error instanceof Error ? error.message : String(error); - this.emit({ - type: "turn_failed", - provider: this.provider, - turnId, - error: message, - }); - }); } private emitToolCallEvent( diff --git a/packages/server/src/server/agent/providers/pi/test-utils/fake-pi.ts b/packages/server/src/server/agent/providers/pi/test-utils/fake-pi.ts index 8093baeb2..bf0f5907d 100644 --- a/packages/server/src/server/agent/providers/pi/test-utils/fake-pi.ts +++ b/packages/server/src/server/agent/providers/pi/test-utils/fake-pi.ts @@ -45,6 +45,12 @@ export interface FakePiSubagentMessagesResult { messages: PiAgentMessage[]; } +interface FakePiUserEntry { + id: string; + parentId: string | null; + text: string; +} + export class FakePi implements PiRuntime { readonly recordedLaunches: PiRuntimeLaunch[] = []; private readonly sessions: FakePiSession[] = []; @@ -349,6 +355,19 @@ export class FakePiSession implements PiRuntimeSession { this.emit({ type: "agent_end", messages: this.messages }); } + finishSubmittedUserMessage(entry: FakePiUserEntry): void { + this.emit({ + type: "message_end", + message: { role: "user", content: entry.text }, + }); + this.emit({ + type: "extension_ui_request", + id: `submitted-user-${entry.id}`, + method: "notify", + message: `PASEO_SUBMITTED_USER_ENTRY ${JSON.stringify({ entry })}`, + }); + } + private handleTreeNavigationCommand(message: string): void { const prefix = "/paseo_tree "; if (!message.startsWith(prefix)) { 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 d992638ba..ef7d43a0f 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 @@ -12,7 +12,7 @@ import type { AgentTimelineItem, } from "../agent/agent-sdk-types.js"; import { DaemonClient } from "../test-utils/daemon-client.js"; -import { createTestPaseoDaemon } from "../test-utils/paseo-daemon.js"; +import { createTestPaseoDaemon, type TestPaseoDaemon } from "../test-utils/paseo-daemon.js"; import { canRunRealProvider, createRealProviderClient, @@ -107,7 +107,7 @@ async function waitForTimelineItem( } async function withConnectedPiDaemon( - run: (context: { client: DaemonClient }) => Promise, + run: (context: { client: DaemonClient; daemon: TestPaseoDaemon }) => Promise, ): Promise { const daemon = await createPiToolDaemon(); const client = new DaemonClient({ @@ -120,7 +120,7 @@ async function withConnectedPiDaemon( await client.fetchAgents({ subscribe: { subscriptionId: `pi-real-${randomUUID()}` }, }); - await run({ client }); + await run({ client, daemon }); } finally { await client.close().catch(() => undefined); await daemon.close().catch(() => undefined); @@ -645,6 +645,70 @@ test( PI_TEST_TIMEOUT_MS, ); +test( + "resumed Pi prompts retain their exact native entry ids after idle collection", + async () => { + const cwd = tmpCwd("pi-resumed-entry-id-"); + const firstPrompt = "PASEO_PI_ENTRY_ID_FIRST. Reply exactly: first-ok"; + const secondPrompt = "PASEO_PI_ENTRY_ID_SECOND. Reply exactly: second-ok"; + + try { + await withConnectedPiDaemon(async ({ client, daemon }) => { + const agent = await client.createAgent({ + cwd, + title: "pi-resumed-entry-id", + provider: "pi", + model: PI_REAL_TEST_MODEL, + }); + + await client.sendMessage(agent.id, firstPrompt); + const firstFinish = await client.waitForFinish(agent.id, PI_TEST_TIMEOUT_MS); + expect(firstFinish.status).toBe("idle"); + + const collection = await daemon.daemon.agentManager.collectIdleAgents({ + cutoff: new Date(Date.now() + 1_000), + protectedAgentIds: new Set(), + }); + expect(collection.failures).toEqual([]); + expect(collection.collected.map((entry) => entry.agentId)).toContain(agent.id); + + await client.sendMessage(agent.id, secondPrompt); + const secondFinish = await client.waitForFinish(agent.id, PI_TEST_TIMEOUT_MS); + expect(secondFinish.status).toBe("idle"); + + const nativeHandle = secondFinish.final?.persistence?.nativeHandle; + if (typeof nativeHandle !== "string") { + throw new Error("Real Pi run did not return a native session file"); + } + const nativeUserEntryIds = readFileSync(nativeHandle, "utf8") + .trim() + .split("\n") + .map((line) => JSON.parse(line) as Record) + .filter( + (entry) => + entry.type === "message" && + typeof entry.id === "string" && + typeof entry.message === "object" && + entry.message !== null && + (entry.message as { role?: unknown }).role === "user", + ) + .map((entry) => entry.id as string); + const userMessages = (await fetchCanonicalTimeline(client, agent.id)).filter( + (item): item is Extract => + item.type === "user_message", + ); + + expect(userMessages.map((item) => item.text)).toEqual([firstPrompt, secondPrompt]); + expect(userMessages.map((item) => item.messageId)).toEqual(nativeUserEntryIds); + expect(new Set(nativeUserEntryIds).size).toBe(2); + }); + } finally { + rmSync(cwd, { recursive: true, force: true }); + } + }, + PI_TEST_TIMEOUT_MS * 2, +); + test( "streamHistory replays user and assistant timeline after resume", async () => {