agent/pi
1
0
Fork 0
mirror of https://github.com/earendil-works/pi.git synced 2026-06-18 15:54:04 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(coding-agent): defer extension reload requests safely

Browse Source
This commit is contained in:
Armin Ronacher
2026-06-14 22:30:11 +02:00
Unverified
parent d683a581b7
commit 7d3bcb6d13
11 changed files with 680 additions and 145 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/coding-agent/docs/extensions.md
+43 -57
View File
@@ -994,6 +994,48 @@ ctx.compact({
});
```
### ctx.reload()
Request the same runtime reload flow as `/reload`. This is available in `ExtensionContext`, so tools, event handlers, shortcuts, and command handlers can call it.
`ctx.reload()` records a reload request and returns a promise that resolves once the request is accepted. The reload runs when the session reaches a safe boundary. In command handlers this may happen immediately; in tools and event handlers it is deferred so reload does not re-enter while hook results are still being applied.
Reload behavior:
- Reload emits `session_shutdown` for the current extension runtime
- It then reloads resources and emits `session_start` with `reason: "reload"` and `resources_discover` with reason `"reload"`
- The currently running handler continues in the old call frame
- Code after `await ctx.reload()` must not assume old in-memory extension state will remain valid
- After the handler returns and reload completes, future commands/events/tool calls use the new extension version
- A reload request made while a reload is already running queues one follow-up reload when the session is idle
For predictable command behavior, treat reload as terminal for that handler (`await ctx.reload(); return;`).
Example tool the LLM can call to request reload:
```typescript
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { Type } from "typebox";
export default function (pi: ExtensionAPI) {
pi.registerTool({
name: "reload_runtime",
label: "Reload Runtime",
description: "Reload extensions, skills, prompts, and themes",
parameters: Type.Object({}),
async execute(_toolCallId, _params, _signal, _onUpdate, ctx) {
await ctx.reload();
return {
content: [{ type: "text", text: "Reload requested. It will run after this turn finishes." }],
details: {},
terminate: true,
};
},
});
}
```
Return `terminate: true` when reload should happen as soon as the current tool batch finishes; otherwise the agent may continue with another LLM turn before it becomes idle.
### ctx.getSystemPrompt()
Returns Pi's current system prompt string.
@@ -1204,62 +1246,6 @@ pi.registerCommand("handoff", {
});
```
### ctx.reload()
Run the same reload flow as `/reload`.
```typescript
pi.registerCommand("reload-runtime", {
description: "Reload extensions, skills, prompts, and themes",
handler: async (_args, ctx) => {
await ctx.reload();
return;
},
});
```
Important behavior:
- `await ctx.reload()` emits `session_shutdown` for the current extension runtime
- It then reloads resources and emits `session_start` with `reason: "reload"` and `resources_discover` with reason `"reload"`
- The currently running command handler still continues in the old call frame
- Code after `await ctx.reload()` still runs from the pre-reload version
- Code after `await ctx.reload()` must not assume old in-memory extension state is still valid
- After the handler returns, future commands/events/tool calls use the new extension version
For predictable behavior, treat reload as terminal for that handler (`await ctx.reload(); return;`).
Tools run with `ExtensionContext`, so they cannot call `ctx.reload()` directly. Use a command as the reload entrypoint, then expose a tool that queues that command as a follow-up user message.
Example tool the LLM can call to trigger reload:
```typescript
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { Type } from "typebox";
export default function (pi: ExtensionAPI) {
pi.registerCommand("reload-runtime", {
description: "Reload extensions, skills, prompts, and themes",
handler: async (_args, ctx) => {
await ctx.reload();
return;
},
});
pi.registerTool({
name: "reload_runtime",
label: "Reload Runtime",
description: "Reload extensions, skills, prompts, and themes",
parameters: Type.Object({}),
async execute() {
pi.sendUserMessage("/reload-runtime", { deliverAs: "followUp" });
return {
content: [{ type: "text", text: "Queued /reload-runtime as a follow-up command." }],
};
},
});
}
```
## ExtensionAPI Methods
### pi.on(event, handler)
@@ -2597,7 +2583,7 @@ All examples in [examples/extensions/](../examples/extensions/).
| `handoff.ts` | Cross-provider model handoff | `registerCommand`, `ui.editor`, `ui.custom` |
| `qna.ts` | Q&A with custom UI | `registerCommand`, `ui.custom`, `setEditorText` |
| `send-user-message.ts` | Inject user messages | `registerCommand`, `sendUserMessage` |
| `reload-runtime.ts` | Reload command and LLM tool handoff | `registerCommand`, `ctx.reload()`, `sendUserMessage` |
| `reload-runtime.ts` | LLM tool reload request | `registerTool`, `ctx.reload()` |
| `shutdown-command.ts` | Graceful shutdown command | `registerCommand`, `shutdown()` |
| **Events & Gates** |||
| `permission-gate.ts` | Block dangerous commands | `on("tool_call")`, `ui.confirm` |
packages/coding-agent/examples/extensions/README.md
+1 -1
View File
@@ -74,7 +74,7 @@ cp permission-gate.ts ~/.pi/agent/extensions/
| `overlay-qa-tests.ts` | Comprehensive overlay QA tests: anchors, margins, stacking, overflow, animation |
| `doom-overlay/` | DOOM game running as an overlay at 35 FPS (demonstrates real-time game rendering) |
| `shutdown-command.ts` | Adds `/quit` command demonstrating `ctx.shutdown()` |
| `reload-runtime.ts` | Adds `/reload-runtime` and `reload_runtime` tool showing safe reload flow |
| `reload-runtime.ts` | Adds a `reload_runtime` tool showing safe direct reload requests |
| `interactive-shell.ts` | Run interactive commands (vim, htop) with full terminal via `user_bash` hook |
| `inline-bash.ts` | Expands `!{command}` patterns in prompts via `input` event transformation |
| `input-transform-streaming.ts` | Skips expensive input preprocessing for mid-stream steering via `streamingBehavior` |
packages/coding-agent/examples/extensions/reload-runtime.ts
+8 -17
View File
@@ -1,36 +1,27 @@
/**
* Reload Runtime Extension
*
* Demonstrates ctx.reload() from ExtensionCommandContext and an LLM-callable
* tool that queues a follow-up command to trigger reload.
* Demonstrates ctx.reload() from an LLM-callable tool. Tool-triggered
* reloads are deferred until the agent turn is idle.
*/
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { Type } from "typebox";
export default function (pi: ExtensionAPI) {
// Command entrypoint for reload.
// Treat reload as terminal for this handler.
pi.registerCommand("reload-runtime", {
description: "Reload extensions, skills, prompts, and themes",
handler: async (_args, ctx) => {
await ctx.reload();
return;
},
});
// LLM-callable tool. Tools get ExtensionContext, so they cannot call ctx.reload() directly.
// Instead, queue a follow-up user command that executes the command above.
// During a tool call, ctx.reload() requests a reload that
// runs after the current agent turn is fully idle.
pi.registerTool({
name: "reload_runtime",
label: "Reload Runtime",
description: "Reload extensions, skills, prompts, and themes",
parameters: Type.Object({}),
async execute() {
pi.sendUserMessage("/reload-runtime", { deliverAs: "followUp" });
async execute(_toolCallId, _params, _signal, _onUpdate, ctx) {
await ctx.reload();
return {
content: [{ type: "text", text: "Queued /reload-runtime as a follow-up command." }],
content: [{ type: "text", text: "Reload requested. It will run after this turn finishes." }],
details: {},
terminate: true,
};
},
});
packages/coding-agent/src/core/agent-session-runtime.ts
+24 -21
View File
@@ -139,11 +139,13 @@ export class AgentSessionRuntime {
return { cancelled: false };
}
const result = await runner.emit({
type: "session_before_switch",
reason,
targetSessionFile,
});
const result = await this.session.withReloadDeferred(() =>
runner.emit({
type: "session_before_switch",
reason,
targetSessionFile,
}),
);
return { cancelled: result?.cancel === true };
}
@@ -156,20 +158,26 @@ export class AgentSessionRuntime {
return { cancelled: false };
}
const result = await runner.emit({
type: "session_before_fork",
entryId,
...options,
});
const result = await this.session.withReloadDeferred(() =>
runner.emit({
type: "session_before_fork",
entryId,
...options,
}),
);
return { cancelled: result?.cancel === true };
}
private async teardownCurrent(reason: SessionShutdownEvent["reason"], targetSessionFile?: string): Promise<void> {
await emitSessionShutdownEvent(this.session.extensionRunner, {
type: "session_shutdown",
reason,
targetSessionFile,
});
await this.session.withReloadDeferred(
() =>
emitSessionShutdownEvent(this.session.extensionRunner, {
type: "session_shutdown",
reason,
targetSessionFile,
}),
{ flush: false },
);
this.beforeSessionInvalidate?.();
this.session.dispose();
}
@@ -388,12 +396,7 @@ export class AgentSessionRuntime {
}
async dispose(): Promise<void> {
await emitSessionShutdownEvent(this.session.extensionRunner, {
type: "session_shutdown",
reason: "quit",
});
this.beforeSessionInvalidate?.();
this.session.dispose();
await this.teardownCurrent("quit");
}
}
packages/coding-agent/src/core/agent-session.ts
+161 -27
View File
@@ -76,6 +76,8 @@ import {
type TreePreparation,
type TurnEndEvent,
type TurnStartEvent,
type UserBashEvent,
type UserBashEventResult,
wrapRegisteredTools,
} from "./extensions/index.ts";
import { emitSessionShutdownEvent } from "./extensions/runner.ts";
@@ -190,6 +192,7 @@ export interface ExtensionBindings {
uiContext?: ExtensionUIContext;
mode?: ExtensionMode;
commandContextActions?: ExtensionCommandContextActions;
reloadHandler?: RuntimeReloadHandler;
abortHandler?: () => void;
shutdownHandler?: ShutdownHandler;
onError?: ExtensionErrorListener;
@@ -242,6 +245,8 @@ interface ToolDefinitionEntry {
sourceInfo: SourceInfo;
}
type RuntimeReloadHandler = (reloadCore: () => Promise<void>) => Promise<void>;
// ============================================================================
// Constants
// ============================================================================
@@ -304,6 +309,8 @@ export class AgentSession {
private _extensionUIContext?: ExtensionUIContext;
private _extensionMode: ExtensionMode = "print";
private _extensionCommandContextActions?: ExtensionCommandContextActions;
private _runtimeReloadHandler?: RuntimeReloadHandler;
private _extensionsBound = false;
private _extensionAbortHandler?: () => void;
private _extensionShutdownHandler?: ShutdownHandler;
private _extensionErrorListener?: ExtensionErrorListener;
@@ -322,6 +329,11 @@ export class AgentSession {
private _baseSystemPrompt = "";
private _baseSystemPromptOptions!: BuildSystemPromptOptions;
// Deferred runtime reload state
private _reloadDeferralDepth = 0;
private _reloadRequested = false;
private _reloadInProgress = false;
constructor(config: AgentSessionConfig) {
this.agent = config.agent;
this.sessionManager = config.sessionManager;
@@ -725,6 +737,7 @@ export class AgentSession {
// Dispose must succeed even if an abort hook throws.
}
this._reloadRequested = false;
this._extensionRunner.invalidate(
"This extension ctx is stale after session replacement or reload. Do not use a captured pi or command ctx after ctx.newSession(), ctx.fork(), ctx.switchSession(), or ctx.reload(). For newSession, fork, and switchSession, move post-replacement work into withSession and use the ctx passed to withSession. For reload, do not use the old ctx after await ctx.reload().",
);
@@ -934,14 +947,16 @@ export class AgentSession {
// =========================================================================
private async _runAgentPrompt(messages: AgentMessage | AgentMessage[]): Promise<void> {
try {
await this.agent.prompt(messages);
while (await this._handlePostAgentRun()) {
await this.agent.continue();
await this.withReloadDeferred(async () => {
try {
await this.agent.prompt(messages);
while (await this._handlePostAgentRun()) {
await this.agent.continue();
}
} finally {
this._flushPendingBashMessages();
}
} finally {
this._flushPendingBashMessages();
}
});
}
private async _handlePostAgentRun(): Promise<boolean> {
@@ -986,6 +1001,8 @@ export class AgentSession {
async prompt(text: string, options?: PromptOptions): Promise<void> {
const expandPromptTemplates = options?.expandPromptTemplates ?? true;
const preflightResult = options?.preflightResult;
let preflightAccepted = false;
let reloadDeferred = false;
let messages: AgentMessage[] | undefined;
try {
@@ -996,10 +1013,14 @@ export class AgentSession {
if (handled) {
// Extension command executed, no prompt to send
preflightResult?.(true);
preflightAccepted = true;
return;
}
}
this._reloadDeferralDepth++;
reloadDeferred = true;
// Emit input event for extension interception (before skill/template expansion)
let currentText = text;
let currentImages = options?.images;
@@ -1012,6 +1033,7 @@ export class AgentSession {
);
if (inputResult.action === "handled") {
preflightResult?.(true);
preflightAccepted = true;
return;
}
if (inputResult.action === "transform") {
@@ -1040,6 +1062,7 @@ export class AgentSession {
await this._queueSteer(expandedText, currentImages);
}
preflightResult?.(true);
preflightAccepted = true;
return;
}
@@ -1123,17 +1146,27 @@ export class AgentSession {
// Ensure we're using the base prompt (in case previous turn had modifications)
this.agent.state.systemPrompt = this._baseSystemPrompt;
}
if (!messages) {
return;
}
preflightResult?.(true);
preflightAccepted = true;
await this._runAgentPrompt(messages);
} catch (error) {
preflightResult?.(false);
if (!preflightAccepted) {
preflightResult?.(false);
}
throw error;
} finally {
if (reloadDeferred) {
this._reloadDeferralDepth--;
if (this._reloadDeferralDepth === 0) {
await this._flushRequestedReload();
}
}
}
if (!messages) {
return;
}
preflightResult?.(true);
await this._runAgentPrompt(messages);
}
/**
@@ -1416,6 +1449,71 @@ export class AgentSession {
await this.agent.waitForIdle();
}
async withReloadDeferred<T>(callback: () => Promise<T>, options: { flush?: boolean } = {}): Promise<T> {
this._reloadDeferralDepth++;
try {
return await callback();
} finally {
this._reloadDeferralDepth--;
if (this._reloadDeferralDepth === 0 && options.flush !== false) {
await this._flushRequestedReload();
}
}
}
private _shouldDeferReload(): boolean {
return this._reloadDeferralDepth > 0 || this.isStreaming || this.isCompacting;
}
private async _runReloadHandler(): Promise<void> {
if (this._reloadInProgress) {
this._reloadRequested = true;
return;
}
this._reloadInProgress = true;
try {
do {
this._reloadRequested = false;
const reloadCore = () => this._reloadCore();
await (this._runtimeReloadHandler ? this._runtimeReloadHandler(reloadCore) : reloadCore());
} while (this._reloadRequested && !this._shouldDeferReload());
} finally {
this._reloadInProgress = false;
}
}
private async _flushRequestedReload(): Promise<void> {
if (!this._reloadRequested || this._shouldDeferReload()) {
return;
}
try {
await this._runReloadHandler();
} catch (error) {
this._extensionRunner.emitError({
extensionPath: "<runtime>",
event: "reload",
error: error instanceof Error ? error.message : String(error),
stack: error instanceof Error ? error.stack : undefined,
});
throw error;
}
}
/**
* Request a runtime reload. If the agent or session is busy, reload is deferred
* until the current prompt, compaction, or branch summary operation finishes.
*/
async requestReload(): Promise<void> {
this._reloadRequested = true;
if (this._shouldDeferReload() || this._reloadInProgress) {
return;
}
await this._flushRequestedReload();
}
// =========================================================================
// Model Management
// =========================================================================
@@ -1426,11 +1524,13 @@ export class AgentSession {
source: "set" | "cycle" | "restore",
): Promise<void> {
if (modelsAreEqual(previousModel, nextModel)) return;
await this._extensionRunner.emit({
type: "model_select",
model: nextModel,
previousModel,
source,
await this.withReloadDeferred(async () => {
await this._extensionRunner.emit({
type: "model_select",
model: nextModel,
previousModel,
source,
});
});
}
@@ -1548,10 +1648,15 @@ export class AgentSession {
this.settingsManager.setDefaultThinkingLevel(effectiveLevel);
}
this._emit({ type: "thinking_level_changed", level: effectiveLevel });
void this._extensionRunner.emit({
type: "thinking_level_select",
level: effectiveLevel,
previousLevel,
void this.withReloadDeferred(async () => {
await this._extensionRunner.emit({
type: "thinking_level_select",
level: effectiveLevel,
previousLevel,
});
}).catch(() => {
// _flushRequestedReload already emitted the extension error. setThinkingLevel
// is synchronous, so there is no caller to reject.
});
}
}
@@ -1766,6 +1871,7 @@ export class AgentSession {
} finally {
this._compactionAbortController = undefined;
this._reconnectToAgent();
await this._flushRequestedReload();
}
}
@@ -2053,6 +2159,7 @@ export class AgentSession {
return false;
} finally {
this._autoCompactionAbortController = undefined;
await this._flushRequestedReload();
}
}
@@ -2069,6 +2176,7 @@ export class AgentSession {
}
async bindExtensions(bindings: ExtensionBindings): Promise<void> {
this._extensionsBound = true;
if (bindings.uiContext !== undefined) {
this._extensionUIContext = bindings.uiContext;
}
@@ -2078,6 +2186,9 @@ export class AgentSession {
if (bindings.commandContextActions !== undefined) {
this._extensionCommandContextActions = bindings.commandContextActions;
}
if (bindings.reloadHandler !== undefined) {
this._runtimeReloadHandler = bindings.reloadHandler;
}
if (bindings.abortHandler !== undefined) {
this._extensionAbortHandler = bindings.abortHandler;
}
@@ -2089,8 +2200,10 @@ export class AgentSession {
}
this._applyExtensionBindings(this._extensionRunner);
await this._extensionRunner.emit(this._sessionStartEvent);
await this.extendResourcesFromExtensions(this._sessionStartEvent.reason === "reload" ? "reload" : "startup");
await this.withReloadDeferred(async () => {
await this._extensionRunner.emit(this._sessionStartEvent);
await this.extendResourcesFromExtensions(this._sessionStartEvent.reason === "reload" ? "reload" : "startup");
});
}
private async extendResourcesFromExtensions(reason: "startup" | "reload"): Promise<void> {
@@ -2148,7 +2261,14 @@ export class AgentSession {
private _applyExtensionBindings(runner: ExtensionRunner): void {
runner.setUIContext(this._extensionUIContext, this._extensionMode);
runner.bindCommandContext(this._extensionCommandContextActions);
runner.bindCommandContext({
waitForIdle: this._extensionCommandContextActions?.waitForIdle ?? (() => this.agent.waitForIdle()),
newSession: this._extensionCommandContextActions?.newSession ?? (async () => ({ cancelled: false })),
fork: this._extensionCommandContextActions?.fork ?? (async () => ({ cancelled: false })),
navigateTree: this._extensionCommandContextActions?.navigateTree ?? (async () => ({ cancelled: false })),
switchSession: this._extensionCommandContextActions?.switchSession ?? (async () => ({ cancelled: false })),
reload: this._extensionCommandContextActions?.reload ?? (() => this.requestReload()),
});
this._extensionErrorUnsubscriber?.();
this._extensionErrorUnsubscriber = this._extensionErrorListener
@@ -2269,6 +2389,7 @@ export class AgentSession {
}
})();
},
reload: () => this.requestReload(),
getSystemPrompt: () => this.systemPrompt,
getSystemPromptOptions: () => this._baseSystemPromptOptions,
},
@@ -2433,6 +2554,10 @@ export class AgentSession {
}
async reload(): Promise<void> {
await this._runReloadHandler();
}
private async _reloadCore(): Promise<void> {
const previousFlagValues = this._extensionRunner.getFlagValues();
await emitSessionShutdownEvent(this._extensionRunner, { type: "session_shutdown", reason: "reload" });
await this.settingsManager.reload();
@@ -2446,6 +2571,7 @@ export class AgentSession {
});
const hasBindings =
this._extensionsBound ||
this._extensionUIContext ||
this._extensionCommandContextActions ||
this._extensionShutdownHandler ||
@@ -2569,6 +2695,13 @@ export class AgentSession {
// Bash Execution
// =========================================================================
async emitUserBash(event: UserBashEvent): Promise<UserBashEventResult | undefined> {
if (!this._extensionRunner.hasHandlers("user_bash")) {
return undefined;
}
return this.withReloadDeferred(() => this._extensionRunner.emitUserBash(event));
}
/**
* Execute a bash command.
* Adds result to agent context and session.
@@ -2883,6 +3016,7 @@ export class AgentSession {
return { editorText, cancelled: false, summaryEntry };
} finally {
this._branchSummaryAbortController = undefined;
await this._flushRequestedReload();
}
}
packages/coding-agent/src/core/extensions/runner.ts
+9 -3
View File
@@ -284,6 +284,7 @@ export class ExtensionRunner {
private navigateTreeHandler: NavigateTreeHandler = async () => ({ cancelled: false });
private switchSessionHandler: SwitchSessionHandler = async () => ({ cancelled: false });
private reloadHandler: ReloadHandler = async () => {};
private commandReloadHandler: ReloadHandler | undefined;
private shutdownHandler: ShutdownHandler = () => {};
private shortcutDiagnostics: ResourceDiagnostic[] = [];
private commandDiagnostics: ResourceDiagnostic[] = [];
@@ -338,6 +339,7 @@ export class ExtensionRunner {
this.shutdownHandler = contextActions.shutdown;
this.getContextUsageFn = contextActions.getContextUsage;
this.compactFn = contextActions.compact;
this.reloadHandler = contextActions.reload ?? (async () => {});
this.getSystemPromptFn = contextActions.getSystemPrompt;
this.getSystemPromptOptionsFn = contextActions.getSystemPromptOptions ?? (() => ({ cwd: this.cwd }));
@@ -385,7 +387,7 @@ export class ExtensionRunner {
this.forkHandler = actions.fork;
this.navigateTreeHandler = actions.navigateTree;
this.switchSessionHandler = actions.switchSession;
this.reloadHandler = actions.reload;
this.commandReloadHandler = actions.reload;
return;
}
@@ -394,7 +396,7 @@ export class ExtensionRunner {
this.forkHandler = async () => ({ cancelled: false });
this.navigateTreeHandler = async () => ({ cancelled: false });
this.switchSessionHandler = async () => ({ cancelled: false });
this.reloadHandler = async () => {};
this.commandReloadHandler = undefined;
}
setUIContext(uiContext?: ExtensionUIContext, mode: ExtensionMode = "print"): void {
@@ -678,6 +680,10 @@ export class ExtensionRunner {
runner.assertActive();
runner.compactFn(options);
},
reload: () => {
runner.assertActive();
return runner.reloadHandler();
},
getSystemPrompt: () => {
runner.assertActive();
return runner.getSystemPromptFn();
@@ -719,7 +725,7 @@ export class ExtensionRunner {
};
context.reload = () => {
this.assertActive();
return this.reloadHandler();
return (this.commandReloadHandler ?? this.reloadHandler)();
};
return context;
}
packages/coding-agent/src/core/extensions/types.ts
+7 -1
View File
@@ -328,6 +328,11 @@ export interface ExtensionContext {
getContextUsage(): ContextUsage | undefined;
/** Trigger compaction without awaiting completion. */
compact(options?: CompactOptions): void;
/**
* Request a runtime reload. Resolves after the request is accepted; reload may
* run later when the session reaches a safe boundary.
*/
reload(): Promise<void>;
/** Get the current effective system prompt. */
getSystemPrompt(): string;
}
@@ -1537,6 +1542,7 @@ export interface ExtensionContextActions {
shutdown: () => void;
getContextUsage: () => ContextUsage | undefined;
compact: (options?: CompactOptions) => void;
reload?: () => Promise<void>;
getSystemPrompt: () => string;
getSystemPromptOptions?: () => BuildSystemPromptOptions;
}
@@ -1564,7 +1570,7 @@ export interface ExtensionCommandContextActions {
sessionPath: string,
options?: { withSession?: (ctx: ReplacedSessionContext) => Promise<void> },
) => Promise<{ cancelled: boolean }>;
reload: () => Promise<void>;
reload?: () => Promise<void>;
}
/**
packages/coding-agent/src/modes/interactive/interactive-mode.ts
+27 -18
View File
@@ -1602,10 +1602,8 @@ export class InteractiveMode {
switchSession: async (sessionPath, options) => {
return this.handleResumeSession(sessionPath, options);
},
reload: async () => {
await this.handleReloadCommand();
},
},
reloadHandler: (reloadCore) => this.performReload(reloadCore),
shutdownHandler: () => {
this.shutdownRequested = true;
if (!this.session.isStreaming) {
@@ -1718,6 +1716,7 @@ export class InteractiveMode {
}
})();
},
reload: () => this.session.requestReload(),
getSystemPrompt: () => this.session.systemPrompt,
});
@@ -5058,16 +5057,7 @@ export class InteractiveMode {
// Command handlers
// =========================================================================
private async handleReloadCommand(): Promise<void> {
if (this.session.isStreaming) {
this.showWarning("Wait for the current response to finish before reloading.");
return;
}
if (this.session.isCompacting) {
this.showWarning("Wait for compaction to finish before reloading.");
return;
}
private async performReload(reloadCore: () => Promise<void>): Promise<void> {
this.resetExtensionUI();
const reloadBox = new Container();
@@ -5095,7 +5085,8 @@ export class InteractiveMode {
};
try {
await this.session.reload();
let themeWarning: string | undefined;
await reloadCore();
configureHttpDispatcher(this.settingsManager.getHttpIdleTimeoutMs());
this.keybindings.reload();
const activeHeader = this.customHeader ?? this.builtInHeader;
@@ -5107,7 +5098,7 @@ export class InteractiveMode {
const themeName = this.settingsManager.getTheme();
const themeResult = themeName ? setTheme(themeName, true) : { success: true };
if (!themeResult.success) {
this.showError(`Failed to load theme "${themeName}": ${themeResult.error}\nFell back to dark theme.`);
themeWarning = `Failed to load theme "${themeName}": ${themeResult.error}\nFell back to dark theme.`;
}
const editorPaddingX = this.settingsManager.getEditorPaddingX();
const autocompleteMaxVisible = this.settingsManager.getAutocompleteMaxVisible();
@@ -5124,6 +5115,9 @@ export class InteractiveMode {
this.setupExtensionShortcuts(runner);
this.rebuildChatFromMessages();
dismissReloadBox(this.editor as Component);
if (themeWarning) {
this.showWarning(themeWarning);
}
this.showLoadedResources({
force: false,
showDiagnosticsWhenQuiet: true,
@@ -5140,6 +5134,23 @@ export class InteractiveMode {
);
} catch (error) {
dismissReloadBox(previousEditor as Component);
throw error;
}
}
private async handleReloadCommand(): Promise<void> {
if (this.session.isStreaming) {
this.showWarning("Wait for the current response to finish before reloading.");
return;
}
if (this.session.isCompacting) {
this.showWarning("Wait for compaction to finish before reloading.");
return;
}
try {
await this.session.reload();
} catch (error) {
this.showError(`Reload failed: ${error instanceof Error ? error.message : String(error)}`);
}
}
@@ -5633,10 +5644,8 @@ export class InteractiveMode {
}
private async handleBashCommand(command: string, excludeFromContext = false): Promise<void> {
const extensionRunner = this.session.extensionRunner;
// Emit user_bash event to let extensions intercept
const eventResult = await extensionRunner.emitUserBash({
const eventResult = await this.session.emitUserBash({
type: "user_bash",
command,
excludeFromContext,
packages/coding-agent/test/agent-session-reload-request.test.ts
+387
View File
@@ -0,0 +1,387 @@
import { existsSync, mkdirSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { fauxAssistantMessage, fauxToolCall, registerFauxProvider } from "@earendil-works/pi-ai";
import { Type } from "typebox";
import { afterEach, beforeEach, describe, expect, it } from "vitest";
import { AuthStorage } from "../src/core/auth-storage.ts";
import type { ExtensionCommandContextActions } from "../src/core/extensions/index.ts";
import { DefaultResourceLoader } from "../src/core/resource-loader.ts";
import { createAgentSession } from "../src/core/sdk.ts";
import { SessionManager } from "../src/core/session-manager.ts";
import { SettingsManager } from "../src/core/settings-manager.ts";
describe("AgentSession reload requests", () => {
let tempDir: string;
let agentDir: string;
const cleanups: Array<() => void> = [];
beforeEach(() => {
tempDir = join(tmpdir(), `pi-reload-request-test-${Date.now()}-${Math.random().toString(36).slice(2)}`);
agentDir = join(tempDir, "agent");
mkdirSync(agentDir, { recursive: true });
});
afterEach(() => {
while (cleanups.length > 0) {
cleanups.pop()?.();
}
if (existsSync(tempDir)) {
rmSync(tempDir, { recursive: true, force: true });
}
});
it("defers ctx.reload from an extension tool until the agent turn is idle", async () => {
const events: string[] = [];
const faux = registerFauxProvider();
cleanups.push(() => faux.unregister());
faux.setResponses([
fauxAssistantMessage(fauxToolCall("reload_runtime", {}, { id: "reload-1" }), { stopReason: "toolUse" }),
]);
const authStorage = AuthStorage.inMemory();
authStorage.setRuntimeApiKey(faux.getModel().provider, "faux-key");
const settingsManager = SettingsManager.create(tempDir, agentDir);
const sessionManager = SessionManager.inMemory();
const resourceLoader = new DefaultResourceLoader({
cwd: tempDir,
agentDir,
settingsManager,
extensionFactories: [
(pi) => {
pi.registerTool({
name: "reload_runtime",
label: "Reload Runtime",
description: "Reload extensions, skills, prompts, and themes",
parameters: Type.Object({}),
async execute(_toolCallId, _params, _signal, _onUpdate, ctx) {
events.push("tool_execute");
await ctx.reload();
events.push("reload_requested");
return {
content: [{ type: "text", text: "reload requested" }],
details: {},
terminate: true,
};
},
});
pi.on("agent_end", () => {
events.push("agent_end");
});
},
],
});
await resourceLoader.reload();
const { session } = await createAgentSession({
cwd: tempDir,
agentDir,
authStorage,
model: faux.getModel(),
settingsManager,
sessionManager,
resourceLoader,
});
cleanups.push(() => session.dispose());
const commandContextActions: ExtensionCommandContextActions = {
waitForIdle: () => session.agent.waitForIdle(),
newSession: async () => ({ cancelled: false }),
fork: async () => ({ cancelled: false }),
navigateTree: async () => ({ cancelled: false }),
switchSession: async () => ({ cancelled: false }),
};
await session.bindExtensions({
commandContextActions,
reloadHandler: async (reloadCore) => {
events.push("reload");
await reloadCore();
},
});
await session.prompt("reload now");
expect(events).toEqual(["tool_execute", "reload_requested", "agent_end", "reload"]);
});
it("defers ctx.reload from before_agent_start until the prompt finishes", async () => {
const events: string[] = [];
const faux = registerFauxProvider();
cleanups.push(() => faux.unregister());
faux.setResponses([fauxAssistantMessage("ok")]);
const authStorage = AuthStorage.inMemory();
authStorage.setRuntimeApiKey(faux.getModel().provider, "faux-key");
const settingsManager = SettingsManager.create(tempDir, agentDir);
const sessionManager = SessionManager.inMemory();
const resourceLoader = new DefaultResourceLoader({
cwd: tempDir,
agentDir,
settingsManager,
extensionFactories: [
(pi) => {
pi.on("before_agent_start", async (_event, ctx) => {
events.push("before_agent_start");
await ctx.reload();
events.push("reload_requested");
return { systemPrompt: "modified" };
});
pi.on("agent_end", () => {
events.push("agent_end");
});
},
],
});
await resourceLoader.reload();
const { session } = await createAgentSession({
cwd: tempDir,
agentDir,
authStorage,
model: faux.getModel(),
settingsManager,
sessionManager,
resourceLoader,
});
cleanups.push(() => session.dispose());
await session.bindExtensions({
reloadHandler: async (reloadCore) => {
events.push("reload");
await reloadCore();
},
});
await session.prompt("reload before start");
expect(events).toEqual(["before_agent_start", "reload_requested", "agent_end", "reload"]);
});
it("flushes ctx.reload requested from idle user_bash handlers", async () => {
const events: string[] = [];
const faux = registerFauxProvider();
cleanups.push(() => faux.unregister());
const authStorage = AuthStorage.inMemory();
authStorage.setRuntimeApiKey(faux.getModel().provider, "faux-key");
const settingsManager = SettingsManager.create(tempDir, agentDir);
const sessionManager = SessionManager.inMemory();
const resourceLoader = new DefaultResourceLoader({
cwd: tempDir,
agentDir,
settingsManager,
extensionFactories: [
(pi) => {
pi.on("user_bash", async (_event, ctx) => {
events.push("user_bash");
await ctx.reload();
events.push("reload_requested");
});
},
],
});
await resourceLoader.reload();
const { session } = await createAgentSession({
cwd: tempDir,
agentDir,
authStorage,
model: faux.getModel(),
settingsManager,
sessionManager,
resourceLoader,
});
cleanups.push(() => session.dispose());
await session.bindExtensions({
reloadHandler: async (reloadCore) => {
events.push("reload");
await reloadCore();
},
});
await session.emitUserBash({ type: "user_bash", command: "echo hi", excludeFromContext: false, cwd: tempDir });
expect(events).toEqual(["user_bash", "reload_requested", "reload"]);
});
it("flushes ctx.reload requested from idle selection events", async () => {
const events: string[] = [];
const faux = registerFauxProvider({ models: [{ id: "faux-thinker", reasoning: true }] });
cleanups.push(() => faux.unregister());
const model = faux.getModel("faux-thinker");
if (!model) throw new Error("faux-thinker model not found");
const authStorage = AuthStorage.inMemory();
authStorage.setRuntimeApiKey(model.provider, "faux-key");
const settingsManager = SettingsManager.create(tempDir, agentDir);
const sessionManager = SessionManager.inMemory();
const resourceLoader = new DefaultResourceLoader({
cwd: tempDir,
agentDir,
settingsManager,
extensionFactories: [
(pi) => {
pi.on("thinking_level_select", async (_event, ctx) => {
events.push("thinking_level_select");
await ctx.reload();
events.push("reload_requested");
});
},
],
});
await resourceLoader.reload();
const { session } = await createAgentSession({
cwd: tempDir,
agentDir,
authStorage,
model,
settingsManager,
sessionManager,
resourceLoader,
});
cleanups.push(() => session.dispose());
await session.bindExtensions({
reloadHandler: async (reloadCore) => {
events.push("reload");
await reloadCore();
},
});
session.setThinkingLevel("low");
await new Promise((resolve) => setImmediate(resolve));
expect(events).toEqual(["thinking_level_select", "reload_requested", "reload"]);
});
it("runs a follow-up reload when direct reload is called during reload", async () => {
const events: string[] = [];
const faux = registerFauxProvider();
cleanups.push(() => faux.unregister());
const authStorage = AuthStorage.inMemory();
authStorage.setRuntimeApiKey(faux.getModel().provider, "faux-key");
const settingsManager = SettingsManager.create(tempDir, agentDir);
const sessionManager = SessionManager.inMemory();
const resourceLoader = new DefaultResourceLoader({
cwd: tempDir,
agentDir,
settingsManager,
});
await resourceLoader.reload();
const { session } = await createAgentSession({
cwd: tempDir,
agentDir,
authStorage,
model: faux.getModel(),
settingsManager,
sessionManager,
resourceLoader,
});
cleanups.push(() => session.dispose());
let reloads = 0;
const commandContextActions: ExtensionCommandContextActions = {
waitForIdle: () => session.agent.waitForIdle(),
newSession: async () => ({ cancelled: false }),
fork: async () => ({ cancelled: false }),
navigateTree: async () => ({ cancelled: false }),
switchSession: async () => ({ cancelled: false }),
};
await session.bindExtensions({
commandContextActions,
reloadHandler: async (reloadCore) => {
reloads++;
events.push(`reload_${reloads}`);
if (reloads === 1) {
await session.reload();
}
await reloadCore();
},
});
await session.requestReload();
expect(events).toEqual(["reload_1", "reload_2"]);
});
it("rejects the prompt when a deferred reload fails", async () => {
const events: string[] = [];
const extensionErrors: string[] = [];
const faux = registerFauxProvider();
cleanups.push(() => faux.unregister());
faux.setResponses([
fauxAssistantMessage(fauxToolCall("reload_runtime", {}, { id: "reload-fail-1" }), { stopReason: "toolUse" }),
]);
const authStorage = AuthStorage.inMemory();
authStorage.setRuntimeApiKey(faux.getModel().provider, "faux-key");
const settingsManager = SettingsManager.create(tempDir, agentDir);
const sessionManager = SessionManager.inMemory();
const resourceLoader = new DefaultResourceLoader({
cwd: tempDir,
agentDir,
settingsManager,
extensionFactories: [
(pi) => {
pi.registerTool({
name: "reload_runtime",
label: "Reload Runtime",
description: "Reload extensions, skills, prompts, and themes",
parameters: Type.Object({}),
async execute(_toolCallId, _params, _signal, _onUpdate, ctx) {
events.push("tool_execute");
await ctx.reload();
events.push("reload_requested");
return {
content: [{ type: "text", text: "reload requested" }],
details: {},
terminate: true,
};
},
});
pi.on("agent_end", () => {
events.push("agent_end");
});
},
],
});
await resourceLoader.reload();
const { session } = await createAgentSession({
cwd: tempDir,
agentDir,
authStorage,
model: faux.getModel(),
settingsManager,
sessionManager,
resourceLoader,
});
cleanups.push(() => session.dispose());
const commandContextActions: ExtensionCommandContextActions = {
waitForIdle: () => session.agent.waitForIdle(),
newSession: async () => ({ cancelled: false }),
fork: async () => ({ cancelled: false }),
navigateTree: async () => ({ cancelled: false }),
switchSession: async () => ({ cancelled: false }),
};
await session.bindExtensions({
commandContextActions,
reloadHandler: async () => {
events.push("reload");
throw new Error("reload failed");
},
onError: (error) => {
extensionErrors.push(error.error);
},
});
await expect(session.prompt("reload now")).rejects.toThrow("reload failed");
expect(events).toEqual(["tool_execute", "reload_requested", "agent_end", "reload"]);
expect(extensionErrors).toContain("reload failed");
});
});
packages/coding-agent/test/extensions-runner.test.ts
+12
View File
@@ -873,6 +873,18 @@ describe("ExtensionRunner", () => {
await commandContext.fork("entry-2", { position: "at" });
expect(fork).toHaveBeenLastCalledWith("entry-2", { position: "at" });
});
it("keeps reload awaitable on command contexts", async () => {
const runtime = createExtensionRuntime();
const runner = new ExtensionRunner([], runtime, tempDir, sessionManager, modelRegistry);
const reload = vi.fn(async () => {});
runner.bindCore(extensionActions, { ...extensionContextActions, reload });
const commandContext = runner.createCommandContext();
await commandContext.reload();
expect(reload).toHaveBeenCalledOnce();
});
});
describe("hasHandlers", () => {
packages/coding-agent/test/trigger-compact-extension.test.ts
+1
View File
@@ -19,6 +19,7 @@ function createContext(tokens: number | null, compact = vi.fn()): ExtensionConte
shutdown: vi.fn(),
getContextUsage: () => ({ tokens, contextWindow: 200_000, percent: tokens === null ? null : tokens / 2000 }),
compact,
reload: vi.fn(),
getSystemPrompt: () => "",
};
}