From ebe437081e4bc093883cd7241111f6b1396a5141 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Tue, 24 Mar 2026 23:50:47 +0100 Subject: [PATCH] fix(coding-agent): emit session_shutdown in print mode closes #2576 --- packages/coding-agent/CHANGELOG.md | 1 + packages/coding-agent/src/main.ts | 5 +- packages/coding-agent/src/modes/print-mode.ts | 171 ++++++++++-------- packages/coding-agent/test/print-mode.test.ts | 124 +++++++++++++ 4 files changed, 220 insertions(+), 81 deletions(-) create mode 100644 packages/coding-agent/test/print-mode.test.ts diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index c91aa45b..d64edc1d 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -8,6 +8,7 @@ - Fixed interactive bash execution collapsed previews to recompute visual line wrapping at render time, so previews respect the current terminal width after resizes and split-pane width changes ([#2569](https://github.com/badlogic/pi-mono/issues/2569)) - Fixed RPC `get_session_stats` to expose `contextUsage`, so headless clients can read actual current context-window usage instead of deriving it from token totals ([#2550](https://github.com/badlogic/pi-mono/issues/2550)) - Fixed `pi update` for git packages to fetch only the tracked target branch with `--no-tags`, reducing unrelated branch and tag noise while preserving force-push-safe updates ([#2548](https://github.com/badlogic/pi-mono/issues/2548)) +- Fixed print and JSON modes to emit `session_shutdown` before exit, so extensions can release long-lived resources and non-interactive runs terminate cleanly ([#2576](https://github.com/badlogic/pi-mono/issues/2576)) ## [0.62.0] - 2026-03-23 diff --git a/packages/coding-agent/src/main.ts b/packages/coding-agent/src/main.ts index e69f7a91..482dcd8d 100644 --- a/packages/coding-agent/src/main.ts +++ b/packages/coding-agent/src/main.ts @@ -882,7 +882,7 @@ export async function main(args: string[]) { printTimings(); await interactiveMode.run(); } else { - await runPrintMode(session, { + const exitCode = await runPrintMode(session, { mode, messages: parsed.messages, initialMessage, @@ -890,6 +890,9 @@ export async function main(args: string[]) { }); stopThemeWatcher(); restoreStdout(); + if (exitCode !== 0) { + process.exitCode = exitCode; + } return; } } diff --git a/packages/coding-agent/src/modes/print-mode.ts b/packages/coding-agent/src/modes/print-mode.ts index 5bced90f..90ae6511 100644 --- a/packages/coding-agent/src/modes/print-mode.ts +++ b/packages/coding-agent/src/modes/print-mode.ts @@ -28,93 +28,104 @@ export interface PrintModeOptions { * Run in print (single-shot) mode. * Sends prompts to the agent and outputs the result. */ -export async function runPrintMode(session: AgentSession, options: PrintModeOptions): Promise { +export async function runPrintMode(session: AgentSession, options: PrintModeOptions): Promise { const { mode, messages = [], initialMessage, initialImages } = options; - if (mode === "json") { - const header = session.sessionManager.getHeader(); - if (header) { - writeRawStdout(`${JSON.stringify(header)}\n`); - } - } - // Set up extensions for print mode (no UI) - await session.bindExtensions({ - commandContextActions: { - waitForIdle: () => session.agent.waitForIdle(), - newSession: async (options) => { - const success = await session.newSession({ parentSession: options?.parentSession }); - if (success && options?.setup) { - await options.setup(session.sessionManager); - } - return { cancelled: !success }; - }, - fork: async (entryId) => { - const result = await session.fork(entryId); - return { cancelled: result.cancelled }; - }, - navigateTree: async (targetId, options) => { - const result = await session.navigateTree(targetId, { - summarize: options?.summarize, - customInstructions: options?.customInstructions, - replaceInstructions: options?.replaceInstructions, - label: options?.label, - }); - return { cancelled: result.cancelled }; - }, - switchSession: async (sessionPath) => { - const success = await session.switchSession(sessionPath); - return { cancelled: !success }; - }, - reload: async () => { - await session.reload(); - }, - }, - onError: (err) => { - console.error(`Extension error (${err.extensionPath}): ${err.error}`); - }, - }); + let exitCode = 0; - // Always subscribe to enable session persistence via _handleAgentEvent - session.subscribe((event) => { - // In JSON mode, output all events + try { if (mode === "json") { - writeRawStdout(`${JSON.stringify(event)}\n`); - } - }); - - // Send initial message with attachments - if (initialMessage) { - await session.prompt(initialMessage, { images: initialImages }); - } - - // Send remaining messages - for (const message of messages) { - await session.prompt(message); - } - - // In text mode, output final response - if (mode === "text") { - const state = session.state; - const lastMessage = state.messages[state.messages.length - 1]; - - if (lastMessage?.role === "assistant") { - const assistantMsg = lastMessage as AssistantMessage; - - // Check for error/aborted - if (assistantMsg.stopReason === "error" || assistantMsg.stopReason === "aborted") { - console.error(assistantMsg.errorMessage || `Request ${assistantMsg.stopReason}`); - process.exit(1); + const header = session.sessionManager.getHeader(); + if (header) { + writeRawStdout(`${JSON.stringify(header)}\n`); } + } + // Set up extensions for print mode (no UI) + await session.bindExtensions({ + commandContextActions: { + waitForIdle: () => session.agent.waitForIdle(), + newSession: async (options) => { + const success = await session.newSession({ parentSession: options?.parentSession }); + if (success && options?.setup) { + await options.setup(session.sessionManager); + } + return { cancelled: !success }; + }, + fork: async (entryId) => { + const result = await session.fork(entryId); + return { cancelled: result.cancelled }; + }, + navigateTree: async (targetId, options) => { + const result = await session.navigateTree(targetId, { + summarize: options?.summarize, + customInstructions: options?.customInstructions, + replaceInstructions: options?.replaceInstructions, + label: options?.label, + }); + return { cancelled: result.cancelled }; + }, + switchSession: async (sessionPath) => { + const success = await session.switchSession(sessionPath); + return { cancelled: !success }; + }, + reload: async () => { + await session.reload(); + }, + }, + onError: (err) => { + console.error(`Extension error (${err.extensionPath}): ${err.error}`); + }, + }); - // Output text content - for (const content of assistantMsg.content) { - if (content.type === "text") { - writeRawStdout(`${content.text}\n`); + // Always subscribe to enable session persistence via _handleAgentEvent + session.subscribe((event) => { + // In JSON mode, output all events + if (mode === "json") { + writeRawStdout(`${JSON.stringify(event)}\n`); + } + }); + + // Send initial message with attachments + if (initialMessage) { + await session.prompt(initialMessage, { images: initialImages }); + } + + // Send remaining messages + for (const message of messages) { + await session.prompt(message); + } + + // In text mode, output final response + if (mode === "text") { + const state = session.state; + const lastMessage = state.messages[state.messages.length - 1]; + + if (lastMessage?.role === "assistant") { + const assistantMsg = lastMessage as AssistantMessage; + + // Check for error/aborted + if (assistantMsg.stopReason === "error" || assistantMsg.stopReason === "aborted") { + console.error(assistantMsg.errorMessage || `Request ${assistantMsg.stopReason}`); + exitCode = 1; + } else { + // Output text content + for (const content of assistantMsg.content) { + if (content.type === "text") { + writeRawStdout(`${content.text}\n`); + } + } } } } - } - // Ensure stdout is fully flushed before returning - // This prevents race conditions where the process exits before all output is written - await flushRawStdout(); + return exitCode; + } finally { + const extensionRunner = session.extensionRunner; + if (extensionRunner?.hasHandlers("session_shutdown")) { + await extensionRunner.emit({ type: "session_shutdown" }); + } + + // Ensure stdout is fully flushed before returning + // This prevents race conditions where the process exits before all output is written + await flushRawStdout(); + } } diff --git a/packages/coding-agent/test/print-mode.test.ts b/packages/coding-agent/test/print-mode.test.ts new file mode 100644 index 00000000..03bdc6f4 --- /dev/null +++ b/packages/coding-agent/test/print-mode.test.ts @@ -0,0 +1,124 @@ +import type { AssistantMessage, ImageContent } from "@mariozechner/pi-ai"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { runPrintMode } from "../src/modes/print-mode.js"; + +type EmitEvent = { type: string }; + +type FakeExtensionRunner = { + hasHandlers: (eventType: string) => boolean; + emit: ReturnType Promise>>; +}; + +type FakeSession = { + sessionManager: { getHeader: () => object | undefined }; + agent: { waitForIdle: () => Promise }; + state: { messages: AssistantMessage[] }; + extensionRunner: FakeExtensionRunner; + bindExtensions: ReturnType; + subscribe: ReturnType; + prompt: ReturnType; + newSession: ReturnType; + fork: ReturnType; + navigateTree: ReturnType; + switchSession: ReturnType; + reload: ReturnType; +}; + +function createAssistantMessage(options?: { + text?: string; + stopReason?: AssistantMessage["stopReason"]; + errorMessage?: string; +}): AssistantMessage { + return { + role: "assistant", + content: options?.text ? [{ type: "text", text: options.text }] : [], + api: "openai-responses", + provider: "openai", + model: "gpt-4o-mini", + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: options?.stopReason ?? "stop", + errorMessage: options?.errorMessage, + timestamp: Date.now(), + }; +} + +function createSession(assistantMessage: AssistantMessage): FakeSession { + const extensionRunner: FakeExtensionRunner = { + hasHandlers: (eventType: string) => eventType === "session_shutdown", + emit: vi.fn(async () => {}), + }; + + const state = { messages: [assistantMessage] }; + + return { + sessionManager: { getHeader: () => undefined }, + agent: { waitForIdle: async () => {} }, + state, + extensionRunner, + bindExtensions: vi.fn(async () => {}), + subscribe: vi.fn(() => () => {}), + prompt: vi.fn(async () => {}), + newSession: vi.fn(async () => true), + fork: vi.fn(async () => ({ cancelled: false })), + navigateTree: vi.fn(async () => ({ cancelled: false })), + switchSession: vi.fn(async () => true), + reload: vi.fn(async () => {}), + }; +} + +afterEach(() => { + vi.restoreAllMocks(); +}); + +describe("runPrintMode", () => { + it("emits session_shutdown in text mode", async () => { + const session = createSession(createAssistantMessage({ text: "done" })); + const images: ImageContent[] = [{ type: "image", mimeType: "image/png", data: "abc" }]; + + const exitCode = await runPrintMode(session as unknown as Parameters[0], { + mode: "text", + initialMessage: "Say done", + initialImages: images, + }); + + expect(exitCode).toBe(0); + expect(session.prompt).toHaveBeenCalledWith("Say done", { images }); + expect(session.extensionRunner.emit).toHaveBeenCalledTimes(1); + expect(session.extensionRunner.emit).toHaveBeenCalledWith({ type: "session_shutdown" }); + }); + + it("emits session_shutdown in json mode", async () => { + const session = createSession(createAssistantMessage({ text: "done" })); + + const exitCode = await runPrintMode(session as unknown as Parameters[0], { + mode: "json", + messages: ["hello"], + }); + + expect(exitCode).toBe(0); + expect(session.prompt).toHaveBeenCalledWith("hello"); + expect(session.extensionRunner.emit).toHaveBeenCalledTimes(1); + expect(session.extensionRunner.emit).toHaveBeenCalledWith({ type: "session_shutdown" }); + }); + + it("emits session_shutdown and returns non-zero on assistant error", async () => { + const session = createSession(createAssistantMessage({ stopReason: "error", errorMessage: "provider failure" })); + const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {}); + + const exitCode = await runPrintMode(session as unknown as Parameters[0], { + mode: "text", + }); + + expect(exitCode).toBe(1); + expect(errorSpy).toHaveBeenCalledWith("provider failure"); + expect(session.extensionRunner.emit).toHaveBeenCalledTimes(1); + expect(session.extensionRunner.emit).toHaveBeenCalledWith({ type: "session_shutdown" }); + }); +});