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.
This commit is contained in:
Gangqiang Yao 2026-11-04 17:18:59 +08:00 • committed by GitHub
parent dca1f55527
commit c3b938e3e9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 193 additions and 0 deletions

View file

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