From e18cfb7639e76c4be52f891228fd98583cce8946 Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Sun, 12 Jul 2026 15:33:51 +0200 Subject: [PATCH] Keep Pi chats usable after canceling extension commands (#2019) * Fix Pi sessions stuck after interrupt * Guard Pi retries from stale prompt failures --- .../server/agent/providers/pi/agent.test.ts | 39 +++++++++++++++++++ .../src/server/agent/providers/pi/agent.ts | 13 +++++-- .../agent/providers/pi/test-utils/fake-pi.ts | 35 +++++++++++++++++ 3 files changed, 84 insertions(+), 3 deletions(-) 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 4372f0b0a..9a38af3e8 100644 --- a/packages/server/src/server/agent/providers/pi/agent.test.ts +++ b/packages/server/src/server/agent/providers/pi/agent.test.ts @@ -138,6 +138,13 @@ class SessionEvents { ); } + nextTurnCancellation(): Promise> { + return this.nextEvent( + (event): event is Extract => + event.type === "turn_canceled", + ); + } + nextPermissionRequest(): Promise> { return this.nextEvent( (event): event is Extract => @@ -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(); diff --git a/packages/server/src/server/agent/providers/pi/agent.ts b/packages/server/src/server/agent/providers/pi/agent.ts index 9d0bece60..9f88008a7 100644 --- a/packages/server/src/server/agent/providers/pi/agent.ts +++ b/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 { + 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 { 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 76e0f4b7f..2fc4072e1 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 @@ -76,6 +76,9 @@ export class FakePiSession implements PiRuntimeSession { state: PiSessionState; private readonly subscribers = new Set<(event: PiRuntimeEvent) => void>(); + private nextHeldPrompt: { promise: Promise; reject: (error: Error) => void } | null = null; + private activeHeldPrompt: { promise: Promise; 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 { 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((_resolve, rejectPromise) => { + reject = rejectPromise; + }); + this.nextHeldPrompt = { promise, reject }; + } + + async failHeldPrompt(error: Error): Promise { + const heldPrompt = this.activeHeldPrompt ?? this.nextHeldPrompt; + if (!heldPrompt) { + throw new Error("FakePi has no held prompt"); + } + heldPrompt.reject(error); + await new Promise((resolve) => setImmediate(resolve)); + } + async compact(customInstructions?: string): Promise { this.compactRequests.push(customInstructions === undefined ? {} : { customInstructions }); this.emit({ type: "compaction_start", reason: "manual" });