shumengya/sproutclaw
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(coding-agent): add session_shutdown reasons closes #2863

Browse Source
This commit is contained in:
Mario Zechner
2026-04-20 23:10:49 +02:00
parent 1891b9ac01
commit 12d7161884
10 changed files with 75 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
packages/coding-agent/CHANGELOG.md
View File

@@ -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

6
packages/coding-agent/docs/extensions.md
View File

@@ -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.
});
```

27
packages/coding-agent/src/core/agent-session-runtime.ts
View File

@@ -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<void> {
await emitSessionShutdownEvent(this.session.extensionRunner);
private async teardownCurrent(reason: SessionShutdownEvent["reason"], targetSessionFile?: string): Promise<void> {
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<void> {
await emitSessionShutdownEvent(this.session.extensionRunner);
await emitSessionShutdownEvent(this.session.extensionRunner, {
type: "session_shutdown",
reason: "quit",
});
this.session.dispose();
}
}

3
packages/coding-agent/src/core/agent-session.ts
View File

@@ -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<void> {
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();

10
packages/coding-agent/src/core/extensions/runner.ts
View File

@@ -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<boolean> {
export async function emitSessionShutdownEvent(
extensionRunner: ExtensionRunner,
event: SessionShutdownEvent,
): Promise<boolean> {
if (extensionRunner.hasHandlers("session_shutdown")) {
await extensionRunner.emit({
type: "session_shutdown",
});
await extensionRunner.emit(event);
return true;
}
return false;

5
packages/coding-agent/src/core/extensions/types.ts
View File

@@ -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 */

18
packages/coding-agent/test/agent-session-runtime-events.test.ts
View File

@@ -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> | 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 },
]);

11
packages/coding-agent/test/print-mode.test.ts
View File

@@ -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" });
});
});

6
packages/coding-agent/test/suite/agent-session-model-extension.test.ts
View File

@@ -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"]);
});
});

18
packages/coding-agent/test/suite/agent-session-runtime.test.ts
View File

@@ -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> | 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 },
]);