diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 8247666ad..d051c8226 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -10,6 +10,8 @@ ### Fixed - A failed token refresh of an MCP connection no longer registers an OAuth client or replaces the PKCE verifier of a sign-in in progress, which made that sign-in fail. MCP sign-in now refreshes the stored tokens with the stored client even when the loopback callback port changed, instead of registering a new client and opening the browser. Ctrl+C in `/mcp` and the `--timeout` of `pi mcp login` now also stop a sign-in while it waits up to 25 seconds for another pi process's token refresh, and a refresh lock that fails for another reason than that wait now fails the sign-in instead of opening the browser. +- Fixed RPC mode ignoring an extension's `ctx.shutdown()` while idle, for example from a timer: the process kept running until the client sent another command. pi now shuts down right away when idle, and a shutdown requested by an extension command waits until the command's `prompt` response is written ([#11599](https://github.com/earendil-works/pi/issues/11599)) + ## [1.5.0] - 2026-11-02 diff --git a/packages/coding-agent/docs/extensions.md b/packages/coding-agent/docs/extensions.md index 6fea5a841..7d214fcf0 100644 --- a/packages/coding-agent/docs/extensions.md +++ b/packages/coding-agent/docs/extensions.md @@ -1261,7 +1261,7 @@ Control flow helpers. `ctx.isIdle()` is false while Pi is processing an agent ru Request a graceful shutdown of pi. - **Interactive mode:** Deferred until the agent becomes idle (after processing all queued steering and follow-up messages). -- **RPC mode:** Deferred until the next idle state (after completing the current command response, when waiting for the next command). +- **RPC mode:** Runs right away when the session is idle and no command is being handled. Otherwise deferred until the current command finishes (including its response) or the active run settles. (A request made while an extension-started compaction is running waits for the next command or run to finish.) - **Print mode:** No-op. The process exits automatically when all prompts are processed. Emits `session_shutdown` event to all extensions before exiting. Available in all contexts (event handlers, tools, commands, shortcuts). diff --git a/packages/coding-agent/docs/rpc.md b/packages/coding-agent/docs/rpc.md index a39cc8da2..d1ae4d8d6 100644 --- a/packages/coding-agent/docs/rpc.md +++ b/packages/coding-agent/docs/rpc.md @@ -92,7 +92,7 @@ Clients must also handle child-process startup failures, unexpected exits, stder Close the child's stdin to request an orderly shutdown. Pi disposes the active runtime before exiting. Clients should still handle process signals and unexpected exits. -An extension can also request shutdown through its extension context. Pi completes shutdown after the current command or after the active run emits `agent_settled`. +An extension can also request shutdown with `ctx.shutdown()`. If the session is idle and no command is being handled, Pi shuts down right away. Otherwise it shuts down once the current command finishes (for `prompt`, once the prompt settles) or the active run emits `agent_settled`. (A request made while an extension-started compaction is running waits for the next command or run to finish.) ## Minimal client diff --git a/packages/coding-agent/src/modes/rpc/rpc-mode.ts b/packages/coding-agent/src/modes/rpc/rpc-mode.ts index 1c0995d36..24ffaac7b 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-mode.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-mode.ts @@ -82,9 +82,12 @@ export async function runRpcMode(runtimeHost: AgentSessionRuntime): Promise void; reject: (error: Error) => void } >(); - // Shutdown request flag + // Extension shutdown requests (ctx.shutdown()) wait until input is attached, no command is + // being handled, and the session is idle. See checkShutdownRequested(). let shutdownRequested = false; let shuttingDown = false; + let inputReady = false; + let activeCommands = 0; const signalCleanupHandlers: Array<() => void> = []; /** Helper for dialog methods with signal/timeout support */ @@ -344,6 +347,9 @@ export async function runRpcMode(runtimeHost: AgentSessionRuntime): Promise { shutdownRequested = true; + // An idle session has no command response or agent_settled left to check the flag. + // Check on the next tick so the requesting handler finishes first. + setImmediate(checkShutdownRequested); }, onError: (err) => { output({ type: "extension_error", extensionPath: err.extensionPath, event: err.event, error: err.error }); @@ -355,7 +361,7 @@ export async function runRpcMode(runtimeHost: AgentSessionRuntime): Promise { output(toJsonEvent(event)); if (event.type === "agent_settled") { - void checkShutdownRequested(); + checkShutdownRequested(); } }); unsubscribeBackpressure = session.agent.subscribe(async () => { @@ -395,6 +401,8 @@ export async function runRpcMode(runtimeHost: AgentSessionRuntime): Promise { + activeCommands--; + checkShutdownRequested(); }); return undefined; } @@ -717,10 +729,6 @@ export async function runRpcMode(runtimeHost: AgentSessionRuntime): Promise {}; async function shutdown(exitCode = 0, signal?: NodeJS.Signals): Promise { @@ -742,9 +750,15 @@ export async function runRpcMode(runtimeHost: AgentSessionRuntime): Promise { - if (!shutdownRequested) return; - await shutdown(); + /** + * Shut down if an extension requested it and nothing still needs the process: input is attached, + * no command (including a prompt that has not settled) is being handled, and the session is idle. + * Each of those ending calls this again. A shutdown already in progress is left alone, because a + * second shutdown() call exits immediately. + */ + function checkShutdownRequested(): void { + if (!shutdownRequested || shuttingDown || !inputReady || activeCommands > 0 || !session.isIdle) return; + void shutdown(); } const handleInputLine = async (line: string) => { @@ -780,13 +794,19 @@ export async function runRpcMode(runtimeHost: AgentSessionRuntime): Promise {}); diff --git a/packages/coding-agent/test/rpc-prompt-response-semantics.test.ts b/packages/coding-agent/test/rpc-prompt-response-semantics.test.ts index 18c40646c..ab6d745d0 100644 --- a/packages/coding-agent/test/rpc-prompt-response-semantics.test.ts +++ b/packages/coding-agent/test/rpc-prompt-response-semantics.test.ts @@ -450,3 +450,103 @@ describe("RPC prompt response semantics", () => { } }); }); + +describe("RPC extension shutdown requests", () => { + afterEach(() => { + rpcIo.outputLines = []; + rpcIo.lineHandler = undefined; + vi.restoreAllMocks(); + }); + + // #11599: an async ctx.shutdown() while idle must not wait for another stdin command. + it("shuts down an idle session when an extension requests shutdown asynchronously", async () => { + const exitSpy = vi.spyOn(process, "exit").mockImplementation((() => undefined) as never); + let requestShutdown: (() => void) | undefined; + const { cleanup } = await startRpcMode({ + withAuth: false, + responseDelayMs: 0, + extensionsResult: await createTestExtensionsResult([ + (pi) => { + pi.on("session_start", (_event, ctx) => { + requestShutdown = () => ctx.shutdown(); + }); + }, + ]), + }); + + try { + await vi.waitFor(() => expect(requestShutdown).toBeDefined()); + expect(exitSpy).not.toHaveBeenCalled(); + requestShutdown!(); + await vi.waitFor(() => expect(exitSpy).toHaveBeenCalledWith(0)); + expect(exitSpy).toHaveBeenCalledTimes(1); + } finally { + await cleanup(); + } + }); + + it("writes the command response before shutting down for a shutdown requested by an extension command", async () => { + const exitSpy = vi.spyOn(process, "exit").mockImplementation((() => { + rpcIo.outputLines.push(`${JSON.stringify({ type: "exit" })}\n`); + }) as never); + const { lineHandler, cleanup } = await startRpcMode({ + withAuth: false, + responseDelayMs: 0, + extensionsResult: await createTestExtensionsResult([ + (pi) => { + pi.registerCommand("quit-now", { + handler: async (_args, ctx) => { + ctx.shutdown(); + await sleep(10); + }, + }); + }, + ]), + }); + + try { + lineHandler(JSON.stringify({ id: "q", type: "prompt", message: "/quit-now" })); + await vi.waitFor(() => expect(exitSpy).toHaveBeenCalledTimes(1)); + const types = parseOutputLines(rpcIo.outputLines).map((line) => (line.id === "q" ? "response" : line.type)); + expect(types.indexOf("response")).toBeGreaterThanOrEqual(0); + expect(types.indexOf("response")).toBeLessThan(types.indexOf("exit")); + } finally { + await cleanup(); + } + }); + + it("waits for a command that arrives before an idle shutdown request is processed", async () => { + const exitSpy = vi.spyOn(process, "exit").mockImplementation((() => { + rpcIo.outputLines.push(`${JSON.stringify({ type: "exit" })}\n`); + }) as never); + let requestShutdown: (() => void) | undefined; + const { lineHandler, cleanup } = await startRpcMode({ + withAuth: false, + responseDelayMs: 0, + extensionsResult: await createTestExtensionsResult([ + (pi) => { + pi.on("session_start", (_event, ctx) => { + requestShutdown = () => ctx.shutdown(); + }); + pi.registerCommand("slow", { + handler: async () => { + await sleep(50); + }, + }); + }, + ]), + }); + + try { + await vi.waitFor(() => expect(requestShutdown).toBeDefined()); + // The idle check is scheduled for the next tick; the command starts before it runs. + requestShutdown!(); + lineHandler(JSON.stringify({ id: "s", type: "prompt", message: "/slow" })); + await vi.waitFor(() => expect(exitSpy).toHaveBeenCalledTimes(1)); + const types = parseOutputLines(rpcIo.outputLines).map((line) => (line.id === "s" ? "response" : line.type)); + expect(types.filter((type) => type === "response" || type === "exit")).toEqual(["response", "exit"]); + } finally { + await cleanup(); + } + }); +});