From 12d71618848744f029fd3f74c1cd783b5b0245d6 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Mon, 20 Apr 2026 23:10:49 +0200 Subject: [PATCH] fix(coding-agent): add session_shutdown reasons closes #2863 --- packages/coding-agent/CHANGELOG.md | 1 + packages/coding-agent/docs/extensions.md | 6 +++-- .../src/core/agent-session-runtime.ts | 27 ++++++++++++------- .../coding-agent/src/core/agent-session.ts | 3 ++- .../src/core/extensions/runner.ts | 10 ++++--- .../coding-agent/src/core/extensions/types.ts | 5 +++- .../test/agent-session-runtime-events.test.ts | 18 +++++++++++-- packages/coding-agent/test/print-mode.test.ts | 11 ++++---- .../agent-session-model-extension.test.ts | 6 ++--- .../test/suite/agent-session-runtime.test.ts | 18 +++++++++++-- 10 files changed, 75 insertions(+), 30 deletions(-) diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 567c3bbf..26ca6bd5 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -6,6 +6,7 @@ - Added extension support for customizing the interactive streaming working indicator via `ctx.ui.setWorkingIndicator()`, including custom animated frames, static indicators, hidden indicators, a new `working-indicator.ts` example extension, and updated extension/TUI/RPC docs ([#3413](https://github.com/badlogic/pi-mono/issues/3413)) - Added `/clone` to duplicate the current active branch into a new session, while keeping `/fork` focused on forking from a previous user message ([#2962](https://github.com/badlogic/pi-mono/issues/2962)) +- Added `reason` and `targetSessionFile` metadata to `session_shutdown` extension events, so extensions can distinguish quit, reload, new-session, resume, and fork teardown paths ([#2863](https://github.com/badlogic/pi-mono/issues/2863)) ### Changed diff --git a/packages/coding-agent/docs/extensions.md b/packages/coding-agent/docs/extensions.md index eb546f87..94aeb7f7 100644 --- a/packages/coding-agent/docs/extensions.md +++ b/packages/coding-agent/docs/extensions.md @@ -446,10 +446,12 @@ pi.on("session_tree", async (event, ctx) => { #### session_shutdown -Fired on exit (Ctrl+C, Ctrl+D, SIGHUP, SIGTERM). +Fired before an extension runtime is torn down. ```typescript -pi.on("session_shutdown", async (_event, ctx) => { +pi.on("session_shutdown", async (event, ctx) => { + // event.reason - "quit" | "reload" | "new" | "resume" | "fork" + // event.targetSessionFile - destination session for session replacement flows // Cleanup, save state, etc. }); ``` diff --git a/packages/coding-agent/src/core/agent-session-runtime.ts b/packages/coding-agent/src/core/agent-session-runtime.ts index 5f3e8d83..74b8e430 100644 --- a/packages/coding-agent/src/core/agent-session-runtime.ts +++ b/packages/coding-agent/src/core/agent-session-runtime.ts @@ -2,7 +2,7 @@ import { copyFileSync, existsSync, mkdirSync } from "node:fs"; import { basename, join, resolve } from "node:path"; import type { AgentSession } from "./agent-session.js"; import type { AgentSessionRuntimeDiagnostic, AgentSessionServices } from "./agent-session-services.js"; -import type { SessionStartEvent } from "./extensions/index.js"; +import type { SessionShutdownEvent, SessionStartEvent } from "./extensions/index.js"; import { emitSessionShutdownEvent } from "./extensions/runner.js"; import type { CreateAgentSessionResult } from "./sdk.js"; import { assertSessionCwdExists } from "./session-cwd.js"; @@ -127,8 +127,12 @@ export class AgentSessionRuntime { return { cancelled: result?.cancel === true }; } - private async teardownCurrent(): Promise { - await emitSessionShutdownEvent(this.session.extensionRunner); + private async teardownCurrent(reason: SessionShutdownEvent["reason"], targetSessionFile?: string): Promise { + await emitSessionShutdownEvent(this.session.extensionRunner, { + type: "session_shutdown", + reason, + targetSessionFile, + }); this.session.dispose(); } @@ -148,7 +152,7 @@ export class AgentSessionRuntime { const previousSessionFile = this.session.sessionFile; const sessionManager = SessionManager.open(sessionPath, undefined, cwdOverride); assertSessionCwdExists(sessionManager, this.cwd); - await this.teardownCurrent(); + await this.teardownCurrent("resume", sessionManager.getSessionFile()); this.apply( await this.createRuntime({ cwd: sessionManager.getCwd(), @@ -176,7 +180,7 @@ export class AgentSessionRuntime { sessionManager.newSession({ parentSession: options.parentSession }); } - await this.teardownCurrent(); + await this.teardownCurrent("new", sessionManager.getSessionFile()); this.apply( await this.createRuntime({ cwd: this.cwd, @@ -229,7 +233,7 @@ export class AgentSessionRuntime { if (!targetLeafId) { const sessionManager = SessionManager.create(this.cwd, sessionDir); sessionManager.newSession({ parentSession: currentSessionFile }); - await this.teardownCurrent(); + await this.teardownCurrent("fork", sessionManager.getSessionFile()); this.apply( await this.createRuntime({ cwd: this.cwd, @@ -247,7 +251,7 @@ export class AgentSessionRuntime { throw new Error("Failed to create forked session"); } const sessionManager = SessionManager.open(forkedSessionPath, sessionDir); - await this.teardownCurrent(); + await this.teardownCurrent("fork", sessionManager.getSessionFile()); this.apply( await this.createRuntime({ cwd: sessionManager.getCwd(), @@ -265,7 +269,7 @@ export class AgentSessionRuntime { } else { sessionManager.createBranchedSession(targetLeafId); } - await this.teardownCurrent(); + await this.teardownCurrent("fork", sessionManager.getSessionFile()); this.apply( await this.createRuntime({ cwd: this.cwd, @@ -308,7 +312,7 @@ export class AgentSessionRuntime { const sessionManager = SessionManager.open(destinationPath, sessionDir, cwdOverride); assertSessionCwdExists(sessionManager, this.cwd); - await this.teardownCurrent(); + await this.teardownCurrent("resume", sessionManager.getSessionFile()); this.apply( await this.createRuntime({ cwd: sessionManager.getCwd(), @@ -321,7 +325,10 @@ export class AgentSessionRuntime { } async dispose(): Promise { - await emitSessionShutdownEvent(this.session.extensionRunner); + await emitSessionShutdownEvent(this.session.extensionRunner, { + type: "session_shutdown", + reason: "quit", + }); this.session.dispose(); } } diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index ae47371d..77dacd5b 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -67,6 +67,7 @@ import { type TurnStartEvent, wrapRegisteredTools, } from "./extensions/index.js"; +import { emitSessionShutdownEvent } from "./extensions/runner.js"; import type { BashExecutionMessage, CustomMessage } from "./messages.js"; import type { ModelRegistry } from "./model-registry.js"; import { expandPromptTemplate, type PromptTemplate } from "./prompt-templates.js"; @@ -2372,7 +2373,7 @@ export class AgentSession { async reload(): Promise { const previousFlagValues = this._extensionRunner.getFlagValues(); - await this._extensionRunner.emit({ type: "session_shutdown" }); + await emitSessionShutdownEvent(this._extensionRunner, { type: "session_shutdown", reason: "reload" }); await this.settingsManager.reload(); resetApiProviders(); await this._resourceLoader.reload(); diff --git a/packages/coding-agent/src/core/extensions/runner.ts b/packages/coding-agent/src/core/extensions/runner.ts index 2022bbea..39d9a9a5 100644 --- a/packages/coding-agent/src/core/extensions/runner.ts +++ b/packages/coding-agent/src/core/extensions/runner.ts @@ -45,6 +45,7 @@ import type { SessionBeforeForkResult, SessionBeforeSwitchResult, SessionBeforeTreeResult, + SessionShutdownEvent, ToolCallEvent, ToolCallEventResult, ToolResultEvent, @@ -168,11 +169,12 @@ export type ShutdownHandler = () => void; * Helper function to emit session_shutdown event to extensions. * Returns true if the event was emitted, false if there were no handlers. */ -export async function emitSessionShutdownEvent(extensionRunner: ExtensionRunner): Promise { +export async function emitSessionShutdownEvent( + extensionRunner: ExtensionRunner, + event: SessionShutdownEvent, +): Promise { if (extensionRunner.hasHandlers("session_shutdown")) { - await extensionRunner.emit({ - type: "session_shutdown", - }); + await extensionRunner.emit(event); return true; } return false; diff --git a/packages/coding-agent/src/core/extensions/types.ts b/packages/coding-agent/src/core/extensions/types.ts index 31e10ee9..2517a34c 100644 --- a/packages/coding-agent/src/core/extensions/types.ts +++ b/packages/coding-agent/src/core/extensions/types.ts @@ -512,9 +512,12 @@ export interface SessionCompactEvent { fromExtension: boolean; } -/** Fired on graceful process shutdown paths such as Ctrl+C, Ctrl+D, SIGHUP, and SIGTERM. */ +/** Fired before an extension runtime is torn down due to quit, reload, or session replacement. */ export interface SessionShutdownEvent { type: "session_shutdown"; + reason: "quit" | "reload" | "new" | "resume" | "fork"; + /** Destination session file when shutting down due to session replacement. */ + targetSessionFile?: string; } /** Preparation data for tree navigation */ diff --git a/packages/coding-agent/test/agent-session-runtime-events.test.ts b/packages/coding-agent/test/agent-session-runtime-events.test.ts index 06a33d49..f21f7997 100644 --- a/packages/coding-agent/test/agent-session-runtime-events.test.ts +++ b/packages/coding-agent/test/agent-session-runtime-events.test.ts @@ -15,10 +15,15 @@ import type { ExtensionFactory, SessionBeforeForkEvent, SessionBeforeSwitchEvent, + SessionShutdownEvent, SessionStartEvent, } from "../src/index.js"; -type RecordedSessionEvent = SessionBeforeSwitchEvent | SessionBeforeForkEvent | SessionStartEvent; +type RecordedSessionEvent = + | SessionBeforeSwitchEvent + | SessionBeforeForkEvent + | SessionShutdownEvent + | SessionStartEvent; describe("AgentSessionRuntime session lifecycle events", () => { const cleanups: Array<() => Promise | void> = []; @@ -90,6 +95,9 @@ describe("AgentSessionRuntime session lifecycle events", () => { pi.on("session_before_switch", (event) => { events.push(event); }); + pi.on("session_shutdown", (event) => { + events.push(event); + }); pi.on("session_start", (event) => { events.push(event); }); @@ -105,13 +113,14 @@ describe("AgentSessionRuntime session lifecycle events", () => { const newSessionResult = await runtimeHost.newSession(); expect(newSessionResult.cancelled).toBe(false); await runtimeHost.session.bindExtensions({}); + const secondSessionFile = runtimeHost.session.sessionFile; expect(events).toEqual([ { type: "session_before_switch", reason: "new", targetSessionFile: undefined }, + { type: "session_shutdown", reason: "new", targetSessionFile: secondSessionFile }, { type: "session_start", reason: "new", previousSessionFile: originalSessionFile }, ]); events.length = 0; - const secondSessionFile = runtimeHost.session.sessionFile; expect(secondSessionFile).toBeTruthy(); const switchResult = await runtimeHost.switchSession(originalSessionFile!); @@ -119,6 +128,7 @@ describe("AgentSessionRuntime session lifecycle events", () => { await runtimeHost.session.bindExtensions({}); expect(events).toEqual([ { type: "session_before_switch", reason: "resume", targetSessionFile: originalSessionFile }, + { type: "session_shutdown", reason: "resume", targetSessionFile: originalSessionFile }, { type: "session_start", reason: "resume", previousSessionFile: secondSessionFile }, ]); }); @@ -158,6 +168,9 @@ describe("AgentSessionRuntime session lifecycle events", () => { return { cancel: true }; } }); + pi.on("session_shutdown", (event) => { + events.push(event); + }); pi.on("session_start", (event) => { events.push(event); }); @@ -176,6 +189,7 @@ describe("AgentSessionRuntime session lifecycle events", () => { await runtimeHost.session.bindExtensions({}); expect(events).toEqual([ { type: "session_before_fork", entryId: userMessage.entryId, position: "before" }, + { type: "session_shutdown", reason: "fork", targetSessionFile: runtimeHost.session.sessionFile }, { type: "session_start", reason: "fork", previousSessionFile }, ]); diff --git a/packages/coding-agent/test/print-mode.test.ts b/packages/coding-agent/test/print-mode.test.ts index 047816a3..80d0accb 100644 --- a/packages/coding-agent/test/print-mode.test.ts +++ b/packages/coding-agent/test/print-mode.test.ts @@ -1,8 +1,9 @@ import type { AssistantMessage, ImageContent } from "@mariozechner/pi-ai"; import { afterEach, describe, expect, it, vi } from "vitest"; +import type { SessionShutdownEvent } from "../src/index.js"; import { runPrintMode } from "../src/modes/print-mode.js"; -type EmitEvent = { type: string }; +type EmitEvent = SessionShutdownEvent; type FakeExtensionRunner = { hasHandlers: (eventType: string) => boolean; @@ -78,7 +79,7 @@ function createRuntimeHost(assistantMessage: AssistantMessage): FakeRuntimeHost fork: vi.fn(async () => ({ selectedText: "" })), switchSession: vi.fn(async () => undefined), dispose: vi.fn(async () => { - await session.extensionRunner.emit({ type: "session_shutdown" }); + await session.extensionRunner.emit({ type: "session_shutdown", reason: "quit" }); }), }; } @@ -102,7 +103,7 @@ describe("runPrintMode", () => { 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" }); + expect(session.extensionRunner.emit).toHaveBeenCalledWith({ type: "session_shutdown", reason: "quit" }); }); it("emits session_shutdown in json mode", async () => { @@ -117,7 +118,7 @@ describe("runPrintMode", () => { expect(exitCode).toBe(0); expect(session.prompt).toHaveBeenCalledWith("hello"); expect(session.extensionRunner.emit).toHaveBeenCalledTimes(1); - expect(session.extensionRunner.emit).toHaveBeenCalledWith({ type: "session_shutdown" }); + expect(session.extensionRunner.emit).toHaveBeenCalledWith({ type: "session_shutdown", reason: "quit" }); }); it("emits session_shutdown and returns non-zero on assistant error", async () => { @@ -134,6 +135,6 @@ describe("runPrintMode", () => { expect(exitCode).toBe(1); expect(errorSpy).toHaveBeenCalledWith("provider failure"); expect(session.extensionRunner.emit).toHaveBeenCalledTimes(1); - expect(session.extensionRunner.emit).toHaveBeenCalledWith({ type: "session_shutdown" }); + expect(session.extensionRunner.emit).toHaveBeenCalledWith({ type: "session_shutdown", reason: "quit" }); }); }); diff --git a/packages/coding-agent/test/suite/agent-session-model-extension.test.ts b/packages/coding-agent/test/suite/agent-session-model-extension.test.ts index 1b207959..652093aa 100644 --- a/packages/coding-agent/test/suite/agent-session-model-extension.test.ts +++ b/packages/coding-agent/test/suite/agent-session-model-extension.test.ts @@ -309,8 +309,8 @@ describe("AgentSession model and extension characterization", () => { pi.on("session_start", async (event) => { lifecycleEvents.push(`start:${event.reason}`); }); - pi.on("session_shutdown", async () => { - lifecycleEvents.push("shutdown"); + pi.on("session_shutdown", async (event) => { + lifecycleEvents.push(`shutdown:${event.reason}`); }); }, ], @@ -320,6 +320,6 @@ describe("AgentSession model and extension characterization", () => { await harness.session.bindExtensions({ shutdownHandler: () => {} }); await harness.session.reload(); - expect(lifecycleEvents).toEqual(["start:startup", "shutdown", "start:reload"]); + expect(lifecycleEvents).toEqual(["start:startup", "shutdown:reload", "start:reload"]); }); }); diff --git a/packages/coding-agent/test/suite/agent-session-runtime.test.ts b/packages/coding-agent/test/suite/agent-session-runtime.test.ts index ed88c7d9..220abfef 100644 --- a/packages/coding-agent/test/suite/agent-session-runtime.test.ts +++ b/packages/coding-agent/test/suite/agent-session-runtime.test.ts @@ -16,10 +16,15 @@ import type { ExtensionFactory, SessionBeforeForkEvent, SessionBeforeSwitchEvent, + SessionShutdownEvent, SessionStartEvent, } from "../../src/index.js"; -type RecordedSessionEvent = SessionBeforeSwitchEvent | SessionBeforeForkEvent | SessionStartEvent; +type RecordedSessionEvent = + | SessionBeforeSwitchEvent + | SessionBeforeForkEvent + | SessionShutdownEvent + | SessionStartEvent; describe("AgentSessionRuntime characterization", () => { const cleanups: Array<() => Promise | void> = []; @@ -121,6 +126,9 @@ describe("AgentSessionRuntime characterization", () => { pi.on("session_before_switch", (event) => { events.push(event); }); + pi.on("session_shutdown", (event) => { + events.push(event); + }); pi.on("session_start", (event) => { events.push(event); }); @@ -138,19 +146,21 @@ describe("AgentSessionRuntime characterization", () => { await runtime.session.bindExtensions({}); expect(runtime.session).not.toBe(originalSession); expect(runtime.session.messages).toEqual([]); + const secondSessionFile = runtime.session.sessionFile; expect(events).toEqual([ { type: "session_before_switch", reason: "new", targetSessionFile: undefined }, + { type: "session_shutdown", reason: "new", targetSessionFile: secondSessionFile }, { type: "session_start", reason: "new", previousSessionFile: originalSessionFile }, ]); events.length = 0; - const secondSessionFile = runtime.session.sessionFile; const switchResult = await runtime.switchSession(originalSessionFile!); expect(switchResult.cancelled).toBe(false); await runtime.session.bindExtensions({}); expect(events).toEqual([ { type: "session_before_switch", reason: "resume", targetSessionFile: originalSessionFile }, + { type: "session_shutdown", reason: "resume", targetSessionFile: originalSessionFile }, { type: "session_start", reason: "resume", previousSessionFile: secondSessionFile }, ]); }); @@ -201,6 +211,9 @@ describe("AgentSessionRuntime characterization", () => { return { cancel: true }; } }); + pi.on("session_shutdown", (event) => { + events.push(event); + }); pi.on("session_start", (event) => { events.push(event); }); @@ -217,6 +230,7 @@ describe("AgentSessionRuntime characterization", () => { await runtime.session.bindExtensions({}); expect(events).toEqual([ { type: "session_before_fork", entryId: userMessage.entryId, position: "before" }, + { type: "session_shutdown", reason: "fork", targetSessionFile: runtime.session.sessionFile }, { type: "session_start", reason: "fork", previousSessionFile }, ]);