fix(coding-agent): reject reload while running or compacting (#9222)

Fixes #9221.

An extension tool can lose its context when an RPC command or SDK caller
reloads the session during tool execution: the tool awaits work, reload
invalidates its runner, then `ctx.getSystemPrompt()` throws a
stale-context error. The failed tool result is stored in session
history.

Reject reload when `isStreaming` or `isCompacting` is already set,
before shutdown or runner invalidation. This gives RPC and direct
callers the same entry checks as the TUI. The production change remains
seven lines.

#9548 removed the wrapper access behind the original reproduction. The
regression exercises the tool's own context access after an await,
through both an RPC-style extension command and direct
`AgentSession.reload()`. Rejection preserves the successful result in
session history and the next provider request. Compaction, tree
navigation, and idle reload are also covered.

Scope: this guards operations already marked running or compacting. It
does not serialize reload with prompt preflight or prevent a new prompt
from starting during an ongoing reload.

Related: #5735 proposes a broader deferred-reload API. This patch only
adds the existing TUI entry checks to the core reload method.

Validation on main `f1b2e77f5`:

- Without the guard, four regression cases fail and idle reload passes.
With it, all five pass.
- `npm run check`, offline build, and `./test.sh` pass.

packages/coding-agent/src/core/agent-session.ts

@@ -3763,6 +3763,13 @@ export class AgentSession {
 	}
 
 	async reload(options?: { beforeSessionStart?: () => void | Promise<void> }): Promise<void> {
+		if (this.isStreaming) {
+			throw new Error("Wait for the current response to finish before reloading.");
+		}
+		if (this.isCompacting) {
+			throw new Error("Wait for compaction to finish before reloading.");
+		}
+
 		const oldRunner = this._extensionRunner;
 		const previousFlagValues = oldRunner.getFlagValues();
 		await emitSessionShutdownEvent(oldRunner, { type: "session_shutdown", reason: "reload" });

packages/coding-agent/test/suite/regressions/9221-reload-during-tool.test.ts

@@ -0,0 +1,186 @@
+import { fauxAssistantMessage, fauxToolCall } from "@earendil-works/pi-ai";
+import type { ToolResultMessage } from "@earendil-works/pi-ai/compat";
+import { Type } from "typebox";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { assistantMsg, userMsg } from "../../utilities.ts";
+import { createHarness, getMessageText, type Harness } from "../harness.ts";
+
+function deferred(): { promise: Promise<void>; resolve: () => void } {
+	let resolve!: () => void;
+	const promise = new Promise<void>((done) => {
+		resolve = done;
+	});
+	return { promise, resolve };
+}
+
+// Regression for https://github.com/earendil-works/pi/issues/9221.
+describe("reload during an active session operation", () => {
+	let harness: Harness | undefined;
+
+	afterEach(() => {
+		harness?.cleanup();
+		vi.restoreAllMocks();
+	});
+
+	it.each(["command", "direct"] as const)(
+		"rejects %s reload without invalidating an active tool's context",
+		async (source) => {
+			const started = deferred();
+			const released = deferred();
+			const shutdown = vi.fn();
+			const errors: string[] = [];
+			let providerResults: ToolResultMessage[] = [];
+			const h = await createHarness({
+				extensionFactories: [
+					(pi) => {
+						pi.on("session_shutdown", shutdown);
+						pi.registerCommand("reload-runtime", {
+							description: "Reload extensions",
+							handler: async (_args, ctx) => ctx.reload(),
+						});
+						pi.registerTool({
+							name: "wait_then_succeed",
+							label: "Wait then succeed",
+							description: "Use the tool context after waiting",
+							parameters: Type.Object({}),
+							execute: async (_id, _params, _signal, _update, ctx) => {
+								started.resolve();
+								await released.promise;
+								// #9221: the tool itself can still use ctx after an async operation.
+								ctx.getSystemPrompt();
+								return { content: [{ type: "text", text: "tool succeeded" }], details: {} };
+							},
+						});
+					},
+				],
+			});
+			harness = h;
+			await h.session.bindExtensions({
+				mode: "rpc",
+				onError: (error) => errors.push(error.error),
+				commandContextActions: {
+					waitForIdle: () => h.session.waitForIdle(),
+					newSession: async () => ({ cancelled: true }),
+					fork: async () => ({ cancelled: true }),
+					navigateTree: (targetId, options) => h.session.navigateTree(targetId, options),
+					switchSession: async () => ({ cancelled: true }),
+					// Match RPC's ctx.reload binding to the core reload entry point.
+					reload: () => h.session.reload(),
+				},
+			});
+			const runner = h.session.extensionRunner;
+			h.setResponses([
+				fauxAssistantMessage(fauxToolCall("wait_then_succeed", {}), { stopReason: "toolUse" }),
+				(context) => {
+					providerResults = context.messages.filter((message) => message.role === "toolResult");
+					return fauxAssistantMessage("done");
+				},
+			]);
+
+			const prompt = h.session.prompt("run the tool");
+			await started.promise;
+			try {
+				if (source === "command") {
+					await h.session.prompt("/reload-runtime", { source: "rpc" });
+				} else {
+					await expect(h.session.reload()).rejects.toThrow(
+						"Wait for the current response to finish before reloading.",
+					);
+				}
+			} finally {
+				released.resolve();
+				await prompt;
+			}
+
+			const persistedResults = h.sessionManager
+				.getEntries()
+				.flatMap((entry) =>
+					entry.type === "message" && entry.message.role === "toolResult" ? [entry.message] : [],
+				);
+			expect(persistedResults).toHaveLength(1);
+			expect(persistedResults[0]).toMatchObject({ isError: false });
+			expect(getMessageText(persistedResults[0])).toBe("tool succeeded");
+			expect(providerResults).toEqual(persistedResults);
+			expect(errors).toEqual(
+				source === "command" ? ["Wait for the current response to finish before reloading."] : [],
+			);
+			expect(shutdown).not.toHaveBeenCalled();
+			expect(h.session.extensionRunner).toBe(runner);
+			expect(runner.getActiveTools()).toContain("wait_then_succeed");
+		},
+	);
+
+	it.each(["compaction", "tree navigation"] as const)("rejects reload during %s", async (operation) => {
+		const started = deferred();
+		const released = deferred();
+		const shutdown = vi.fn();
+		const h = await createHarness({
+			settings: { compaction: { keepRecentTokens: 1 } },
+			extensionFactories: [
+				(pi) => {
+					pi.on("session_shutdown", shutdown);
+					pi.on("session_before_compact", async (event) => {
+						started.resolve();
+						await released.promise;
+						return {
+							compaction: {
+								summary: "summary",
+								firstKeptEntryId: event.preparation.firstKeptEntryId,
+								tokensBefore: event.preparation.tokensBefore,
+							},
+						};
+					});
+					pi.on("session_before_tree", async () => {
+						started.resolve();
+						await released.promise;
+					});
+				},
+			],
+		});
+		harness = h;
+		h.sessionManager.appendMessage(userMsg("first user"));
+		const targetId = h.sessionManager.appendMessage(assistantMsg("first assistant"));
+		h.sessionManager.appendMessage(userMsg("second user"));
+		h.sessionManager.appendMessage(assistantMsg("second assistant"));
+		h.session.agent.state.messages = h.sessionManager.buildSessionContext().messages;
+		const runner = h.session.extensionRunner;
+		const pending = operation === "compaction" ? h.session.compact() : h.session.navigateTree(targetId);
+		await started.promise;
+		try {
+			expect(h.session.isStreaming).toBe(false);
+			await expect(h.session.reload()).rejects.toThrow("Wait for compaction to finish before reloading.");
+			expect(shutdown).not.toHaveBeenCalled();
+			expect(h.session.extensionRunner).toBe(runner);
+		} finally {
+			released.resolve();
+			await pending;
+		}
+		expect(h.session.isIdle).toBe(true);
+	});
+
+	it("still reloads and replaces the runner when idle", async () => {
+		const shutdown = vi.fn();
+		const start = vi.fn();
+		const h = await createHarness({
+			extensionFactories: [
+				(pi) => {
+					pi.on("session_shutdown", shutdown);
+					pi.on("session_start", start);
+				},
+			],
+		});
+		harness = h;
+		await h.session.bindExtensions({ onError: () => {} });
+		start.mockClear();
+		const runner = h.session.extensionRunner;
+		const reload = vi.spyOn(h.session.resourceLoader, "reload");
+
+		await h.session.reload();
+
+		expect(shutdown).toHaveBeenCalledOnce();
+		expect(reload).toHaveBeenCalledOnce();
+		expect(start).toHaveBeenCalledWith(expect.objectContaining({ reason: "reload" }), expect.anything());
+		expect(h.session.extensionRunner).not.toBe(runner);
+		expect(() => runner.getActiveTools()).toThrow("stale");
+	});
+});