Keep Pi chats active through retry recovery (#3639)

* fix(server): preserve Pi turns through retries

Pi reports each provider attempt with agent_end, while agent_settled marks the final user-turn boundary. Keeping modern streams active until settlement prevents a recovered attempt from leaving the chat in an error state.

* fix(server): show Pi provider retries

Automatic retries keep the turn alive, but users still need to understand the delay. Surface Pi's retry-start signal through the existing non-terminal error activity row without changing terminal lifecycle handling.

* refactor(server): clarify Pi retry boundaries

Use the server's object-parameter convention for the retry boundary helpers so their event, turn, and retry semantics stay explicit as the lifecycle evolves.

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<AgentStreamEvent, { type: "turn_completed" }> =>
@@ -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);

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<PiAgentSessionEvent, { type: "agent_end" | "agent_settled" }>;
+    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<PiAgentSessionEvent, { type: "tool_execution_end" }>,
   ): 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) {

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

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 });
   }