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 378b556ff..dafe769be 100644 --- a/packages/server/src/server/agent/providers/pi/agent.test.ts +++ b/packages/server/src/server/agent/providers/pi/agent.test.ts @@ -215,6 +215,20 @@ class SessionEvents { return this.events.map((event) => event.type); } + turnLifecycleEvents() { + return this.events.flatMap((event) => { + if ( + event.type === "turn_started" || + event.type === "turn_completed" || + event.type === "turn_failed" || + event.type === "turn_canceled" + ) { + return [{ type: event.type, turnId: event.turnId }]; + } + return []; + }); + } + turnCompletedEvents() { return this.events.filter( (event): event is Extract => @@ -944,6 +958,132 @@ describe("PiRpcAgentSession", () => { }); }); + test("shows retry activity while keeping the original Pi turn active through recovery", async () => { + const { pi, session, events } = await createSession(); + const fakeSession = pi.latestSession(); + const { turnId } = await session.startTurn("hello"); + + fakeSession.emit({ type: "turn_start" }); + fakeSession.finishAgentRun({ + message: { + role: "assistant", + provider: "test-provider", + model: "test-model", + stopReason: "error", + errorMessage: "Request timed out.", + content: [], + }, + willRetry: true, + }); + fakeSession.emit({ + type: "auto_retry_start", + attempt: 1, + maxAttempts: 3, + delayMs: 2000, + errorMessage: "Request timed out.", + }); + + expect(events.timelineItems()).toContainEqual({ + type: "error", + message: "Provider retry (attempt 1): Request timed out.", + }); + expect(events.turnLifecycleEvents()).toEqual([{ type: "turn_started", turnId }]); + + fakeSession.emit({ type: "turn_start" }); + fakeSession.finishAgentRun({ + message: { + role: "assistant", + provider: "test-provider", + model: "test-model", + stopReason: "stop", + content: [{ type: "text", text: "Recovered response" }], + }, + willRetry: false, + }); + fakeSession.emit({ type: "auto_retry_end", success: true, attempt: 1 }); + fakeSession.settleTurn(); + + expect(events.turnLifecycleEvents()).toEqual([ + { type: "turn_started", turnId }, + { type: "turn_completed", turnId }, + ]); + }); + + test("fails an exhausted Pi recovery only after settlement", async () => { + const { pi, session, events } = await createSession(); + const fakeSession = pi.latestSession(); + const { turnId } = await session.startTurn("hello"); + + fakeSession.emit({ type: "turn_start" }); + fakeSession.finishAgentRun({ + message: { + role: "assistant", + provider: "test-provider", + model: "test-model", + stopReason: "error", + errorMessage: "Request timed out.", + content: [], + }, + willRetry: true, + }); + fakeSession.emit({ + type: "auto_retry_start", + attempt: 1, + maxAttempts: 1, + delayMs: 2000, + errorMessage: "Request timed out.", + }); + fakeSession.emit({ type: "turn_start" }); + fakeSession.finishAgentRun({ + message: { + role: "assistant", + provider: "test-provider", + model: "test-model", + stopReason: "error", + errorMessage: "Insufficient quota.", + content: [], + }, + willRetry: false, + }); + fakeSession.emit({ + type: "auto_retry_end", + success: false, + attempt: 1, + finalError: "Insufficient quota.", + }); + + expect(events.turnLifecycleEvents()).toEqual([{ type: "turn_started", turnId }]); + expect(events.timelineItems()).toContainEqual({ + type: "error", + message: "Provider retry (attempt 1): Request timed out.", + }); + + fakeSession.settleTurn(); + + expect(events.turnLifecycleEvents()).toEqual([ + { type: "turn_started", turnId }, + { type: "turn_failed", turnId }, + ]); + }); + + test("completes legacy Pi turns that have no settlement metadata", async () => { + const { pi, session, events } = await createSession(); + const fakeSession = pi.latestSession(); + const { turnId } = await session.startTurn("hello"); + + fakeSession.emit({ type: "turn_start" }); + fakeSession.finishLegacyTurn({ + role: "assistant", + stopReason: "stop", + content: [{ type: "text", text: "Legacy response" }], + }); + + expect(events.turnLifecycleEvents()).toEqual([ + { type: "turn_started", turnId }, + { type: "turn_completed", turnId }, + ]); + }); + test("resumes by launching Pi with the persisted session file and cwd metadata", 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 a30e46c8b..9251c682d 100644 --- a/packages/server/src/server/agent/providers/pi/agent.ts +++ b/packages/server/src/server/agent/providers/pi/agent.ts @@ -1003,6 +1003,8 @@ function isPiAgentSessionEvent(event: PiRuntimeEvent): event is PiAgentSessionEv case "compaction_start": case "compaction_end": case "agent_end": + case "agent_settled": + case "auto_retry_start": return true; default: return false; @@ -1204,6 +1206,8 @@ export class PiRpcAgentSession implements AgentSession { private activeClientMessageId: string | null = null; private activeAssistantMessageId: string | null = null; private activeTurnStarted = false; + private activeTurnStartedEmitted = false; + private pendingSettledMessages: PiAgentMessage[] | null = null; private activeNoTurnPromptText: string | null = null; private readonly pendingNoTurnOutputs: Array<{ turnId: string; message: string }> = []; private activePromptRequestId: string | null = null; @@ -1297,6 +1301,8 @@ export class PiRpcAgentSession implements AgentSession { this.activeClientMessageId = options?.clientMessageId ?? null; this.activeAssistantMessageId = null; this.activeTurnStarted = false; + this.activeTurnStartedEmitted = false; + this.pendingSettledMessages = null; this.activePromptRequestId = null; this.clearNoTurnBuffers(); this.activeNoTurnPromptText = payload.text; @@ -1328,6 +1334,8 @@ export class PiRpcAgentSession implements AgentSession { this.activeTurnId = null; this.activeClientMessageId = null; this.activeTurnStarted = false; + this.activeTurnStartedEmitted = false; + this.pendingSettledMessages = null; this.activeAssistantMessageId = null; this.clearNoTurnBuffers(); if (isPiRequestAbortError(error)) { @@ -1456,6 +1464,8 @@ export class PiRpcAgentSession implements AgentSession { this.activeTurnId = null; this.activeClientMessageId = null; this.activeTurnStarted = false; + this.activeTurnStartedEmitted = false; + this.pendingSettledMessages = null; this.activeAssistantMessageId = null; this.clearNoTurnBuffers(); this.emit({ @@ -1472,6 +1482,8 @@ export class PiRpcAgentSession implements AgentSession { this.activeTurnId = null; this.activeClientMessageId = null; this.activeTurnStarted = false; + this.activeTurnStartedEmitted = false; + this.pendingSettledMessages = null; this.activeAssistantMessageId = null; this.clearNoTurnBuffers(); this.emit({ @@ -2047,6 +2059,8 @@ export class PiRpcAgentSession implements AgentSession { this.activeTurnId = null; this.activeClientMessageId = null; this.activeTurnStarted = false; + this.activeTurnStartedEmitted = false; + this.pendingSettledMessages = null; this.clearNoTurnBuffers(); this.emit({ type: "turn_failed", @@ -2058,6 +2072,10 @@ export class PiRpcAgentSession implements AgentSession { private handleSessionEvent(event: PiAgentSessionEvent): void { const turnId = this.currentTurnIdForEvent(); + if (event.type === "agent_end" || event.type === "agent_settled") { + this.handleTurnBoundaryEvent({ event, turnId }); + return; + } switch (event.type) { case "agent_start": @@ -2072,6 +2090,10 @@ export class PiRpcAgentSession implements AgentSession { case "turn_start": this.activeTurnStarted = true; this.clearNoTurnBuffers(); + if (this.activeTurnStartedEmitted) { + return; + } + this.activeTurnStartedEmitted = true; this.emit({ type: "turn_started", provider: this.provider, @@ -2128,14 +2150,46 @@ export class PiRpcAgentSession implements AgentSession { }, }); return; - case "agent_end": - this.completeTurn(turnId, event.messages ?? []); + case "auto_retry_start": + this.emit({ + type: "timeline", + provider: this.provider, + turnId, + item: { + type: "error", + message: `Provider retry (attempt ${event.attempt}): ${event.errorMessage}`, + }, + }); return; default: return; } } + private handleTurnBoundaryEvent({ + event, + turnId, + }: { + event: Extract; + turnId: string | undefined; + }): void { + if (event.type === "agent_end") { + // COMPAT(piAgentSettled): added in v0.5.0, remove after 2027-02-21 once the Pi + // floor emits agent_settled and willRetry. + if (event.willRetry === undefined) { + this.completeTurn(turnId, event.messages ?? []); + return; + } + if (this.activeTurnId) { + this.pendingSettledMessages = event.messages ?? []; + } + return; + } + if (this.activeTurnId) { + this.completeTurn(turnId, this.pendingSettledMessages ?? []); + } + } + private handleToolExecutionEnd( event: Extract, ): void { @@ -2305,6 +2359,8 @@ export class PiRpcAgentSession implements AgentSession { this.activeClientMessageId = null; this.activeAssistantMessageId = null; this.activeTurnStarted = false; + this.activeTurnStartedEmitted = false; + this.pendingSettledMessages = null; this.clearNoTurnBuffers(); const errorMessage = latestPiErrorMessage(messages); if (typeof errorMessage === "string" && errorMessage.length > 0) { diff --git a/packages/server/src/server/agent/providers/pi/rpc-types.ts b/packages/server/src/server/agent/providers/pi/rpc-types.ts index a84b42876..ce3241a19 100644 --- a/packages/server/src/server/agent/providers/pi/rpc-types.ts +++ b/packages/server/src/server/agent/providers/pi/rpc-types.ts @@ -187,7 +187,16 @@ export type PiAgentSessionEvent = } | { type: "compaction_start"; reason?: "manual" | "threshold" | "overflow" | string } | { type: "compaction_end"; reason?: string; errorMessage?: string; aborted?: boolean } - | { type: "agent_end"; messages?: PiAgentMessage[] }; + | { type: "agent_end"; messages?: PiAgentMessage[]; willRetry?: boolean } + | { type: "agent_settled" } + | { + type: "auto_retry_start"; + attempt: number; + maxAttempts: number; + delayMs: number; + errorMessage: string; + } + | { type: "auto_retry_end"; success: boolean; attempt: number; finalError?: string }; export type PiRuntimeEvent = | PiAgentSessionEvent 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 c931a853e..0d23486e8 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 @@ -354,7 +354,25 @@ export class FakePiSession implements PiRuntimeSession { } } + finishAgentRun({ message, willRetry }: { message: PiAgentMessage; willRetry: boolean }): void { + this.messages = [...this.messages, message]; + this.emit({ + type: "agent_end", + messages: this.messages, + willRetry, + }); + } + + settleTurn(): void { + this.emit({ type: "agent_settled" }); + } + finishTurn(message: PiAgentMessage = { role: "assistant", content: [] }): void { + this.finishAgentRun({ message, willRetry: false }); + this.settleTurn(); + } + + finishLegacyTurn(message: PiAgentMessage = { role: "assistant", content: [] }): void { this.messages = [...this.messages, message]; this.emit({ type: "agent_end", messages: this.messages }); }