Keep Pi chats usable after canceling extension commands (#2019)

* Fix Pi sessions stuck after interrupt

* Guard Pi retries from stale prompt failures

packages/server/src/server/agent/providers/pi/agent.test.ts

@@ -138,6 +138,13 @@ class SessionEvents {
     );
   }
 
+  nextTurnCancellation(): Promise<Extract<AgentStreamEvent, { type: "turn_canceled" }>> {
+    return this.nextEvent(
+      (event): event is Extract<AgentStreamEvent, { type: "turn_canceled" }> =>
+        event.type === "turn_canceled",
+    );
+  }
+
   nextPermissionRequest(): Promise<Extract<AgentStreamEvent, { type: "permission_requested" }>> {
     return this.nextEvent(
       (event): event is Extract<AgentStreamEvent, { type: "permission_requested" }> =>
@@ -508,6 +515,38 @@ describe("PiRpcAgentSession", () => {
     ]);
   });
 
+  test("canceling a silent Pi extension command leaves the session usable", async () => {
+    const { pi, session, events } = await createSession();
+    const fakeSession = pi.latestSession();
+
+    fakeSession.holdNextPrompt();
+    const firstTurn = await session.startTurn("/silent-search");
+    fakeSession.emit({
+      type: "extension_ui_request",
+      id: "notify-1",
+      method: "notify",
+      message: "Search finished",
+    });
+    await session.interrupt();
+    const cancellation = await events.nextTurnCancellation();
+    await session.startTurn("next request");
+    await fakeSession.failHeldPrompt(new Error("Canceled prompt timed out"));
+
+    expect(cancellation).toEqual({
+      type: "turn_canceled",
+      provider: "pi",
+      reason: "interrupted",
+      turnId: firstTurn.turnId,
+    });
+    expect(fakeSession.prompts).toEqual([
+      { message: "/silent-search", imageCount: 0 },
+      { message: "next request", imageCount: 0 },
+    ]);
+    await expect(session.startTurn("overlapping request")).rejects.toThrow(
+      "A Pi turn is already active",
+    );
+  });
+
   test("adds Pi assistant context to generic provider finish errors", async () => {
     const { pi, session, events } = await createSession();
 

packages/server/src/server/agent/providers/pi/agent.ts

@@ -1125,13 +1125,15 @@ export class PiRpcAgentSession implements AgentSession {
     this.activeTurnId = turnId;
 
     void this.runtimeSession.prompt(payload.text, payload.images).catch((error) => {
-      const failedTurnId = this.activeTurnId ?? turnId;
+      if (this.activeTurnId !== turnId) {
+        return;
+      }
       this.activeTurnId = null;
       if (isPiRequestAbortError(error)) {
         this.emit({
           type: "turn_canceled",
           provider: PI_PROVIDER,
-          turnId: failedTurnId,
+          turnId,
           reason: toDiagnosticErrorMessage(error),
         });
         return;
@@ -1139,7 +1141,7 @@ export class PiRpcAgentSession implements AgentSession {
       this.emit({
         type: "turn_failed",
         provider: PI_PROVIDER,
-        turnId: failedTurnId,
+        turnId,
         error: toDiagnosticErrorMessage(error),
       });
     });
@@ -1234,7 +1236,12 @@ export class PiRpcAgentSession implements AgentSession {
   }
 
   async interrupt(): Promise<void> {
+    const turnId = this.activeTurnId;
     await this.runtimeSession.abort();
+    if (turnId && this.activeTurnId === turnId) {
+      this.activeTurnId = null;
+      this.emit({ type: "turn_canceled", provider: PI_PROVIDER, reason: "interrupted", turnId });
+    }
   }
 
   async revertConversation(input: { messageId: string }): Promise<void> {

packages/server/src/server/agent/providers/pi/test-utils/fake-pi.ts

@@ -76,6 +76,9 @@ export class FakePiSession implements PiRuntimeSession {
   state: PiSessionState;
 
   private readonly subscribers = new Set<(event: PiRuntimeEvent) => void>();
+  private nextHeldPrompt: { promise: Promise<void>; reject: (error: Error) => void } | null = null;
+  private activeHeldPrompt: { promise: Promise<void>; reject: (error: Error) => void } | null =
+    null;
 
   constructor(launch: PiRuntimeLaunch) {
     this.state = {
@@ -103,10 +106,42 @@ export class FakePiSession implements PiRuntimeSession {
     images?: Array<{ type: "image"; data: string; mimeType: string }>,
   ): Promise<void> {
     this.prompts.push({ message, imageCount: images?.length ?? 0 });
+    const heldPrompt = this.nextHeldPrompt;
+    if (heldPrompt) {
+      this.nextHeldPrompt = null;
+      this.activeHeldPrompt = heldPrompt;
+      try {
+        await heldPrompt.promise;
+      } finally {
+        if (this.activeHeldPrompt === heldPrompt) {
+          this.activeHeldPrompt = null;
+        }
+      }
+    }
     this.handleTreeNavigationCommand(message);
     this.handleEntryCaptureCommand(message);
   }
 
+  holdNextPrompt(): void {
+    if (this.nextHeldPrompt || this.activeHeldPrompt) {
+      throw new Error("FakePi already has a held prompt");
+    }
+    let reject!: (error: Error) => void;
+    const promise = new Promise<void>((_resolve, rejectPromise) => {
+      reject = rejectPromise;
+    });
+    this.nextHeldPrompt = { promise, reject };
+  }
+
+  async failHeldPrompt(error: Error): Promise<void> {
+    const heldPrompt = this.activeHeldPrompt ?? this.nextHeldPrompt;
+    if (!heldPrompt) {
+      throw new Error("FakePi has no held prompt");
+    }
+    heldPrompt.reject(error);
+    await new Promise<void>((resolve) => setImmediate(resolve));
+  }
+
   async compact(customInstructions?: string): Promise<void> {
     this.compactRequests.push(customInstructions === undefined ? {} : { customInstructions });
     this.emit({ type: "compaction_start", reason: "manual" });