From ffb36b644d6eefbdfc1aa3458e82359d9b41a7b4 Mon Sep 17 00:00:00 2001 From: dinq Date: Mon, 5 Oct 2026 05:24:33 -0400 Subject: [PATCH] fix: treat extension commands handled by Pi as non-turn prompts --- .../features/chat/runtime/chat-queue.test.ts | 11 +++++++++++ apps/app/src/features/chat/runtime/chat.ts | 3 +++ packages/server/src/harness/pi/process.ts | 7 ++++++- packages/server/src/harness/pi/rpc/rpc-mode.ts | 5 +++-- packages/server/src/harness/pi/rpc/rpc-types.ts | 2 +- packages/server/test/harness/pi/agent.test.ts | 17 +++++++++++++++++ 6 files changed, 41 insertions(+), 4 deletions(-) diff --git a/apps/app/src/features/chat/runtime/chat-queue.test.ts b/apps/app/src/features/chat/runtime/chat-queue.test.ts index b0626e981..8c3457f4f 100644 --- a/apps/app/src/features/chat/runtime/chat-queue.test.ts +++ b/apps/app/src/features/chat/runtime/chat-queue.test.ts @@ -81,6 +81,17 @@ describe("Chat pending prompt", () => { expect(chat.store.getState().messages).toEqual([]); }); + it("returns to ready when Pi handles a command without starting a turn", async () => { + const { chat, transport, attach } = makeChat(); + await attach({}); + transport.promptStarted = false; + + await chat.prompt("/mcp"); + + expect(chat.store.getState().messages).toEqual([]); + expect(chat.store.getState().status).toBe("ready"); + }); + it("keeps the queue empty when a follow-up races to a real prompt", async () => { const { chat, transport, attach, live } = makeChat(); await attach({}); diff --git a/apps/app/src/features/chat/runtime/chat.ts b/apps/app/src/features/chat/runtime/chat.ts index 0e20a1aff..a0186bba0 100644 --- a/apps/app/src/features/chat/runtime/chat.ts +++ b/apps/app/src/features/chat/runtime/chat.ts @@ -635,6 +635,9 @@ export class Chat { }); if (!receipt.started) { this.#state.messages = this.#state.messages.filter((message) => message.id !== messageId); + // No turn started for this prompt (queued into a running turn, or an + // extension command handled it). A running turn re-stamps the phase. + if (this.#state.status === "submitted") this.#setStatus("ready"); } } catch (promptError) { this.#state.error = diff --git a/packages/server/src/harness/pi/process.ts b/packages/server/src/harness/pi/process.ts index 3fb1b3e27..552a97aae 100644 --- a/packages/server/src/harness/pi/process.ts +++ b/packages/server/src/harness/pi/process.ts @@ -535,7 +535,7 @@ export const makePiProcessWithDependencies = ( if (before._tag === "Idle") yield* drainQueue(session.chunks); // Stock Pi before 0.99 omits `data`. 0.99 returns `{ disposition }`. - // pie-pi-process still returns `{ started }`. + // pie-pi-process returns `{ started, disposition }`. const admission = yield* restore( session.transport.command< { readonly started?: boolean; readonly disposition?: string } | undefined @@ -559,6 +559,11 @@ export const makePiProcessWithDependencies = ( output: Stream.empty, }; } + // An extension command (e.g. `/mcp`) consumed the input + // without starting a turn. + if (admission?.disposition === "handled") { + return { turnId: uuid(), started: false, output: Stream.empty }; + } return yield* new AgentOperationError({ sessionId: input.sessionId, operation: "prompt-admission-state", diff --git a/packages/server/src/harness/pi/rpc/rpc-mode.ts b/packages/server/src/harness/pi/rpc/rpc-mode.ts index 2d0222226..835e2c299 100644 --- a/packages/server/src/harness/pi/rpc/rpc-mode.ts +++ b/packages/server/src/harness/pi/rpc/rpc-mode.ts @@ -444,8 +444,9 @@ export async function runRpcMode(runtimeHost: AgentSessionRuntime): Promise { diff --git a/packages/server/src/harness/pi/rpc/rpc-types.ts b/packages/server/src/harness/pi/rpc/rpc-types.ts index e22de8388..e0c3997c3 100644 --- a/packages/server/src/harness/pi/rpc/rpc-types.ts +++ b/packages/server/src/harness/pi/rpc/rpc-types.ts @@ -144,7 +144,7 @@ export type RpcResponse = type: "response"; command: "prompt"; success: true; - data: { started: boolean }; + data: { started: boolean; disposition: "started" | "queued" | "handled" }; } | { id?: string; type: "response"; command: "steer"; success: true } | { id?: string; type: "response"; command: "follow_up"; success: true } diff --git a/packages/server/test/harness/pi/agent.test.ts b/packages/server/test/harness/pi/agent.test.ts index 2c6a128ee..2fea863df 100644 --- a/packages/server/test/harness/pi/agent.test.ts +++ b/packages/server/test/harness/pi/agent.test.ts @@ -94,6 +94,7 @@ rl.on("line", (line) => { if (msg.type !== "prompt") return; const text = msg.message; if (text === "fail") { send({ id: msg.id, type: "response", command: "prompt", success: false, error: "cannot prompt" }); return; } + if (text === "/cmd") { send({ id: msg.id, type: "response", command: "prompt", success: true, data: { started: false, disposition: "handled" } }); return; } if (holding && !msg.streamingBehavior) { send({ id: msg.id, type: "response", command: "prompt", success: false, error: "Agent is already processing. Specify streamingBehavior ('steer' or 'followUp') to queue the message." }); return; @@ -514,6 +515,22 @@ layer(NodeServices.layer)("PiAgent", (it) => { }), ); + it.effect("an extension command handled without a turn leaves the session promptable", () => + Effect.gen(function* () { + const agent = yield* makePiProcess({ executable: { command: makeFake(), prefixArgs: [] } }); + const { sessionId } = yield* agent.session.create({ cwd: "/tmp" }); + const handled = yield* agent.session.prompt({ sessionId, text: "/cmd" }); + assert.equal(handled.started, false); + assert.equal(Array.from(yield* Stream.runCollect(handled.output)).length, 0); + + const prompt = yield* agent.session.prompt({ sessionId, text: "ping" }); + assert.equal(prompt.started, true); + const chunks = yield* Stream.runCollect(prompt.output); + assert.equal(Array.from(chunks).at(-1)?.type, "finish"); + yield* agent.session.abort(sessionId); + }), + ); + it.effect("a child crash evicts only that session and surfaces an error chunk", () => Effect.gen(function* () { const agent = yield* makePiProcess({ executable: { command: makeFake(), prefixArgs: [] } });