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

fix(coding-agent): add replacement-session callbacks closes #2860

Browse Source
This commit is contained in:
Mario Zechner
2026-04-22 12:13:54 +02:00
parent a900d25119
commit 1cc303d053
11 changed files with 579 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -8,7 +8,7 @@
### Breaking Changes
- Session-replacement commands now invalidate captured pre-replacement session-bound extension objects after `ctx.newSession()`, `ctx.fork()`, and `ctx.switchSession()`. Old `pi` and command `ctx` references now throw instead of silently targeting the replaced session. Migration: move post-switch work into `withSession`, and use only the `ReplacedSessionContext` passed to that callback for session-bound operations such as `sendUserMessage()`, `sendMessage()`, and `sessionManager` access.
- Session-replacement commands now invalidate captured pre-replacement session-bound extension objects after `ctx.newSession()`, `ctx.fork()`, and `ctx.switchSession()`. Old `pi` and command `ctx` references now throw instead of silently targeting the replaced session. Migration: if code needs to keep working in the replacement session after one of those calls, pass `withSession` to that same method and do the post-switch work there. In practice, move post-switch `pi.sendUserMessage()`, `pi.sendMessage()`, and command-ctx/session-manager access into `withSession`, and use only the `ReplacedSessionContext` passed to that callback for session-bound operations. Footguns: `withSession` runs after the old extension instance has already received `session_shutdown`, old cleanup may already have invalidated captured state, captured old `pi` / old command `ctx` are stale, and previously extracted raw objects such as `const sm = ctx.sessionManager` remain the caller's responsibility and must not be reused after the switch.
### Fixed

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

@@ -949,8 +949,11 @@ pi.registerCommand("my-cmd", {
Create a new session:
```typescript
const parentSession = ctx.sessionManager.getSessionFile();
const kickoff = "Continue in the replacement session";
const result = await ctx.newSession({
parentSession: ctx.sessionManager.getSessionFile(),
parentSession,
setup: async (sm) => {
sm.appendMessage({
role: "user",
@@ -958,6 +961,10 @@ const result = await ctx.newSession({
timestamp: Date.now(),
});
},
withSession: async (ctx) => {
// Use only the replacement-session ctx here.
await ctx.sendUserMessage(kickoff);
},
});
if (result.cancelled) {
@@ -965,6 +972,11 @@ if (result.cancelled) {
}
```
Options:
- `parentSession`: parent session file to record in the new session header
- `setup`: mutate the new session's `SessionManager` before `withSession` runs
- `withSession`: run post-switch work against a fresh replacement-session context. Do not use captured old `pi` / command `ctx`; see [Session replacement lifecycle and footguns](#session-replacement-lifecycle-and-footguns).
### ctx.fork(entryId, options?)
Fork from a specific entry, creating a new session file:
@@ -984,6 +996,7 @@ if (!cloneResult.cancelled) {
Options:
- `position`: `"before"` (default) forks before the selected user message, restoring that prompt into the editor
- `position`: `"at"` duplicates the active path through the selected entry without restoring editor text
- `withSession`: run post-switch work against a fresh replacement-session context. Do not use captured old `pi` / command `ctx`; see [Session replacement lifecycle and footguns](#session-replacement-lifecycle-and-footguns).
### ctx.navigateTree(targetId, options?)
@@ -1004,17 +1017,24 @@ Options:
- `replaceInstructions`: If true, `customInstructions` replaces the default prompt instead of being appended
- `label`: Label to attach to the branch summary entry (or target entry if not summarizing)
### ctx.switchSession(sessionPath)
### ctx.switchSession(sessionPath, options?)
Switch to a different session file:
```typescript
const result = await ctx.switchSession("/path/to/session.jsonl");
const result = await ctx.switchSession("/path/to/session.jsonl", {
withSession: async (ctx) => {
await ctx.sendUserMessage("Resume work in the replacement session");
},
});
if (result.cancelled) {
// An extension cancelled the switch via session_before_switch
}
```
Options:
- `withSession`: run post-switch work against a fresh replacement-session context. Do not use captured old `pi` / command `ctx`; see [Session replacement lifecycle and footguns](#session-replacement-lifecycle-and-footguns).
To discover available sessions, use the static `SessionManager.list()` or `SessionManager.listAll()` methods:
```typescript
@@ -1036,6 +1056,49 @@ pi.registerCommand("switch", {
});
```
### Session replacement lifecycle and footguns
`withSession` receives a fresh `ReplacedSessionContext`, which extends `ExtensionCommandContext` with async `sendMessage()` and `sendUserMessage()` helpers bound to the replacement session.
Lifecycle and footguns:
- `withSession` runs only after the old session has emitted `session_shutdown`, the old runtime has been torn down, the replacement session has been rebound, and the new extension instance has already received `session_start`.
- The callback still executes in the original closure, not inside the new extension instance. That means your old extension instance may already have run its shutdown cleanup before `withSession` starts.
- Captured old `pi` / old command `ctx` session-bound objects are stale after replacement and will throw if used. Use only the `ctx` passed to `withSession` for session-bound work.
- Previously extracted raw objects are still your responsibility. For example, if you capture `const sm = ctx.sessionManager` before replacement, `sm` is still the old `SessionManager` object. Do not reuse it after replacement.
- Code in `withSession` should assume any state invalidated by your `session_shutdown` handler is already gone. Only capture plain data that survives shutdown cleanly, such as strings, ids, and serialized config.
Safe pattern:
```typescript
pi.registerCommand("handoff", {
handler: async (_args, ctx) => {
const kickoff = "Continue from the replacement session";
await ctx.newSession({
withSession: async (ctx) => {
await ctx.sendUserMessage(kickoff);
},
});
},
});
```
Unsafe pattern:
```typescript
pi.registerCommand("handoff", {
handler: async (_args, ctx) => {
const oldSessionManager = ctx.sessionManager;
await ctx.newSession({
withSession: async (_ctx) => {
// stale old objects: do not do this
oldSessionManager.getSessionFile();
pi.sendUserMessage("wrong");
},
});
},
});
```
### ctx.reload()
Run the same reload flow as `/reload`.

33
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 { SessionShutdownEvent, SessionStartEvent } from "./extensions/index.js";
import type { ReplacedSessionContext, SessionShutdownEvent, SessionStartEvent } from "./extensions/index.js";
import { emitSessionShutdownEvent } from "./extensions/runner.js";
import type { CreateAgentSessionResult } from "./sdk.js";
import { assertSessionCwdExists } from "./session-cwd.js";
@@ -65,6 +65,8 @@ function extractUserMessageText(content: string | Array<{ type: string; text?: s
* caller. The caller is responsible for user-facing error handling.
*/
export class AgentSessionRuntime {
private rebindSession?: (session: AgentSession) => Promise<void>;
constructor(
private _session: AgentSession,
private _services: AgentSessionServices,
@@ -93,6 +95,10 @@ export class AgentSessionRuntime {
return this._modelFallbackMessage;
}
setRebindSession(rebindSession?: (session: AgentSession) => Promise<void>): void {
this.rebindSession = rebindSession;
}
private async emitBeforeSwitch(
reason: "new" | "resume",
targetSessionFile?: string,
@@ -143,14 +149,26 @@ export class AgentSessionRuntime {
this._modelFallbackMessage = result.modelFallbackMessage;
}
async switchSession(sessionPath: string, cwdOverride?: string): Promise<{ cancelled: boolean }> {
private async finishSessionReplacement(withSession?: (ctx: ReplacedSessionContext) => Promise<void>): Promise<void> {
if (this.rebindSession) {
await this.rebindSession(this.session);
}
if (withSession) {
await withSession(this.session.createReplacedSessionContext());
}
}
async switchSession(
sessionPath: string,
options?: { cwdOverride?: string; withSession?: (ctx: ReplacedSessionContext) => Promise<void> },
): Promise<{ cancelled: boolean }> {
const beforeResult = await this.emitBeforeSwitch("resume", sessionPath);
if (beforeResult.cancelled) {
return beforeResult;
}
const previousSessionFile = this.session.sessionFile;
const sessionManager = SessionManager.open(sessionPath, undefined, cwdOverride);
const sessionManager = SessionManager.open(sessionPath, undefined, options?.cwdOverride);
assertSessionCwdExists(sessionManager, this.cwd);
await this.teardownCurrent("resume", sessionManager.getSessionFile());
this.apply(
@@ -161,12 +179,14 @@ export class AgentSessionRuntime {
sessionStartEvent: { type: "session_start", reason: "resume", previousSessionFile },
}),
);
await this.finishSessionReplacement(options?.withSession);
return { cancelled: false };
}
async newSession(options?: {
parentSession?: string;
setup?: (sessionManager: SessionManager) => Promise<void>;
withSession?: (ctx: ReplacedSessionContext) => Promise<void>;
}): Promise<{ cancelled: boolean }> {
const beforeResult = await this.emitBeforeSwitch("new");
if (beforeResult.cancelled) {
@@ -193,12 +213,13 @@ export class AgentSessionRuntime {
await options.setup(this.session.sessionManager);
this.session.agent.state.messages = this.session.sessionManager.buildSessionContext().messages;
}
await this.finishSessionReplacement(options?.withSession);
return { cancelled: false };
}
async fork(
entryId: string,
options?: { position?: "before" | "at" },
options?: { position?: "before" | "at"; withSession?: (ctx: ReplacedSessionContext) => Promise<void> },
): Promise<{ cancelled: boolean; selectedText?: string }> {
const position = options?.position ?? "before";
const beforeResult = await this.emitBeforeFork(entryId, { position });
@@ -242,6 +263,7 @@ export class AgentSessionRuntime {
sessionStartEvent: { type: "session_start", reason: "fork", previousSessionFile },
}),
);
await this.finishSessionReplacement(options?.withSession);
return { cancelled: false, selectedText };
}
@@ -260,6 +282,7 @@ export class AgentSessionRuntime {
sessionStartEvent: { type: "session_start", reason: "fork", previousSessionFile },
}),
);
await this.finishSessionReplacement(options?.withSession);
return { cancelled: false, selectedText };
}
@@ -278,6 +301,7 @@ export class AgentSessionRuntime {
sessionStartEvent: { type: "session_start", reason: "fork", previousSessionFile },
}),
);
await this.finishSessionReplacement(options?.withSession);
return { cancelled: false, selectedText };
}
@@ -321,6 +345,7 @@ export class AgentSessionRuntime {
sessionStartEvent: { type: "session_start", reason: "resume", previousSessionFile },
}),
);
await this.finishSessionReplacement();
return { cancelled: false };
}

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

@@ -53,6 +53,7 @@ import {
type MessageEndEvent,
type MessageStartEvent,
type MessageUpdateEvent,
type ReplacedSessionContext,
type SessionBeforeCompactResult,
type SessionBeforeTreeResult,
type SessionStartEvent,
@@ -722,6 +723,9 @@ export class AgentSession {
* Call this when completely done with the session.
*/
dispose(): void {
this._extensionRunner.invalidate(
"This extension instance is stale after session replacement or reload. Use the provided replacement-session context instead.",
);
this._disconnectFromAgent();
this._eventListeners = [];
}
@@ -3066,6 +3070,16 @@ export class AgentSession {
// Extension System
// =========================================================================
createReplacedSessionContext(): ReplacedSessionContext {
const context = Object.defineProperties(
{},
Object.getOwnPropertyDescriptors(this._extensionRunner.createCommandContext()),
) as ReplacedSessionContext;
context.sendMessage = (message, options) => this.sendCustomMessage(message, options);
context.sendUserMessage = (content, options) => this.sendUserMessage(content, options);
return context;
}
/**
* Check if extensions have handlers for a specific event type.
*/

1
packages/coding-agent/src/core/extensions/index.ts
View File

@@ -103,6 +103,7 @@ export type {
// Commands
RegisteredCommand,
RegisteredTool,
ReplacedSessionContext,
ResolvedCommand,
// Events - Resources
ResourcesDiscoverEvent,

35
packages/coding-agent/src/core/extensions/loader.ts
View File

@@ -121,6 +121,12 @@ export function createExtensionRuntime(): ExtensionRuntime {
const notInitialized = () => {
throw new Error("Extension runtime not initialized. Action methods cannot be called during extension loading.");
};
const state: { staleMessage?: string } = {};
const assertActive = () => {
if (state.staleMessage) {
throw new Error(state.staleMessage);
}
};
const runtime: ExtensionRuntime = {
sendMessage: notInitialized,
@@ -140,6 +146,12 @@ export function createExtensionRuntime(): ExtensionRuntime {
setThinkingLevel: notInitialized,
flagValues: new Map(),
pendingProviderRegistrations: [],
assertActive,
invalidate: (message) => {
state.staleMessage ??=
message ??
"This extension instance is stale after session replacement or reload. Use the provided replacement-session context instead.";
},
// Pre-bind: queue registrations so bindCore() can flush them once the
// model registry is available. bindCore() replaces both with direct calls.
registerProvider: (name, config, extensionPath = "<unknown>") => {
@@ -167,12 +179,14 @@ function createExtensionAPI(
const api = {
// Registration methods - write to extension
on(event: string, handler: HandlerFn): void {
runtime.assertActive();
const list = extension.handlers.get(event) ?? [];
list.push(handler);
extension.handlers.set(event, list);
},
registerTool(tool: ToolDefinition): void {
runtime.assertActive();
extension.tools.set(tool.name, {
definition: tool,
sourceInfo: extension.sourceInfo,
@@ -181,6 +195,7 @@ function createExtensionAPI(
},
registerCommand(name: string, options: Omit<RegisteredCommand, "name" | "sourceInfo">): void {
runtime.assertActive();
extension.commands.set(name, {
name,
sourceInfo: extension.sourceInfo,
@@ -195,6 +210,7 @@ function createExtensionAPI(
handler: (ctx: import("./types.js").ExtensionContext) => Promise<void> | void;
},
): void {
runtime.assertActive();
extension.shortcuts.set(shortcut, { shortcut, extensionPath: extension.path, ...options });
},
@@ -202,6 +218,7 @@ function createExtensionAPI(
name: string,
options: { description?: string; type: "boolean" | "string"; default?: boolean | string },
): void {
runtime.assertActive();
extension.flags.set(name, { name, extensionPath: extension.path, ...options });
if (options.default !== undefined && !runtime.flagValues.has(name)) {
runtime.flagValues.set(name, options.default);
@@ -209,77 +226,95 @@ function createExtensionAPI(
},
registerMessageRenderer<T>(customType: string, renderer: MessageRenderer<T>): void {
runtime.assertActive();
extension.messageRenderers.set(customType, renderer as MessageRenderer);
},
// Flag access - checks extension registered it, reads from runtime
getFlag(name: string): boolean | string | undefined {
runtime.assertActive();
if (!extension.flags.has(name)) return undefined;
return runtime.flagValues.get(name);
},
// Action methods - delegate to shared runtime
sendMessage(message, options): void {
runtime.assertActive();
runtime.sendMessage(message, options);
},
sendUserMessage(content, options): void {
runtime.assertActive();
runtime.sendUserMessage(content, options);
},
appendEntry(customType: string, data?: unknown): void {
runtime.assertActive();
runtime.appendEntry(customType, data);
},
setSessionName(name: string): void {
runtime.assertActive();
runtime.setSessionName(name);
},
getSessionName(): string | undefined {
runtime.assertActive();
return runtime.getSessionName();
},
setLabel(entryId: string, label: string | undefined): void {
runtime.assertActive();
runtime.setLabel(entryId, label);
},
exec(command: string, args: string[], options?: ExecOptions) {
runtime.assertActive();
return execCommand(command, args, options?.cwd ?? cwd, options);
},
getActiveTools(): string[] {
runtime.assertActive();
return runtime.getActiveTools();
},
getAllTools() {
runtime.assertActive();
return runtime.getAllTools();
},
setActiveTools(toolNames: string[]): void {
runtime.assertActive();
runtime.setActiveTools(toolNames);
},
getCommands() {
runtime.assertActive();
return runtime.getCommands();
},
setModel(model) {
runtime.assertActive();
return runtime.setModel(model);
},
getThinkingLevel() {
runtime.assertActive();
return runtime.getThinkingLevel();
},
setThinkingLevel(level) {
runtime.assertActive();
runtime.setThinkingLevel(level);
},
registerProvider(name: string, config: ProviderConfig) {
runtime.assertActive();
runtime.registerProvider(name, config, extension.path);
},
unregisterProvider(name: string) {
runtime.assertActive();
runtime.unregisterProvider(name, extension.path);
},

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

@@ -38,6 +38,7 @@ import type {
ProviderConfig,
RegisteredCommand,
RegisteredTool,
ReplacedSessionContext,
ResolvedCommand,
ResourcesDiscoverEvent,
ResourcesDiscoverResult,
@@ -147,11 +148,12 @@ export type ExtensionErrorListener = (error: ExtensionError) => void;
export type NewSessionHandler = (options?: {
parentSession?: string;
setup?: (sessionManager: SessionManager) => Promise<void>;
withSession?: (ctx: ReplacedSessionContext) => Promise<void>;
}) => Promise<{ cancelled: boolean }>;
export type ForkHandler = (
entryId: string,
options?: { position?: "before" | "at" },
options?: { position?: "before" | "at"; withSession?: (ctx: ReplacedSessionContext) => Promise<void> },
) => Promise<{ cancelled: boolean }>;
export type NavigateTreeHandler = (
@@ -159,7 +161,10 @@ export type NavigateTreeHandler = (
options?: { summarize?: boolean; customInstructions?: string; replaceInstructions?: boolean; label?: string },
) => Promise<{ cancelled: boolean }>;
export type SwitchSessionHandler = (sessionPath: string) => Promise<{ cancelled: boolean }>;
export type SwitchSessionHandler = (
sessionPath: string,
options?: { withSession?: (ctx: ReplacedSessionContext) => Promise<void> },
) => Promise<{ cancelled: boolean }>;
export type ReloadHandler = () => Promise<void>;
@@ -235,6 +240,7 @@ export class ExtensionRunner {
private shutdownHandler: ShutdownHandler = () => {};
private shortcutDiagnostics: ResourceDiagnostic[] = [];
private commandDiagnostics: ResourceDiagnostic[] = [];
private staleMessage: string | undefined;
constructor(
extensions: Extension[],
@@ -451,6 +457,19 @@ export class ExtensionRunner {
return this.shortcutDiagnostics;
}
invalidate(message = "This extension instance is stale after session replacement or reload."): void {
if (!this.staleMessage) {
this.staleMessage = message;
this.runtime.invalidate(message);
}
}
private assertActive(): void {
if (this.staleMessage) {
throw new Error(this.staleMessage);
}
}
onError(listener: ExtensionErrorListener): () => void {
this.errorListeners.add(listener);
return () => this.errorListeners.delete(listener);
@@ -544,37 +563,101 @@ export class ExtensionRunner {
* Context values are resolved at call time, so changes via bindCore/bindUI are reflected.
*/
createContext(): ExtensionContext {
const runner = this;
const getModel = this.getModel;
return {
ui: this.uiContext,
hasUI: this.hasUI(),
cwd: this.cwd,
sessionManager: this.sessionManager,
modelRegistry: this.modelRegistry,
get ui() {
runner.assertActive();
return runner.uiContext;
},
get hasUI() {
runner.assertActive();
return runner.hasUI();
},
get cwd() {
runner.assertActive();
return runner.cwd;
},
get sessionManager() {
runner.assertActive();
return runner.sessionManager;
},
get modelRegistry() {
runner.assertActive();
return runner.modelRegistry;
},
get model() {
runner.assertActive();
return getModel();
},
isIdle: () => this.isIdleFn(),
signal: this.getSignalFn(),
abort: () => this.abortFn(),
hasPendingMessages: () => this.hasPendingMessagesFn(),
shutdown: () => this.shutdownHandler(),
getContextUsage: () => this.getContextUsageFn(),
compact: (options) => this.compactFn(options),
getSystemPrompt: () => this.getSystemPromptFn(),
isIdle: () => {
runner.assertActive();
return runner.isIdleFn();
},
get signal() {
runner.assertActive();
return runner.getSignalFn();
},
abort: () => {
runner.assertActive();
runner.abortFn();
},
hasPendingMessages: () => {
runner.assertActive();
return runner.hasPendingMessagesFn();
},
shutdown: () => {
runner.assertActive();
runner.shutdownHandler();
},
getContextUsage: () => {
runner.assertActive();
return runner.getContextUsageFn();
},
compact: (options) => {
runner.assertActive();
runner.compactFn(options);
},
getSystemPrompt: () => {
runner.assertActive();
return runner.getSystemPromptFn();
},
};
}
createCommandContext(): ExtensionCommandContext {
return {
...this.createContext(),
waitForIdle: () => this.waitForIdleFn(),
newSession: (options) => this.newSessionHandler(options),
fork: (entryId, options) => this.forkHandler(entryId, options),
navigateTree: (targetId, options) => this.navigateTreeHandler(targetId, options),
switchSession: (sessionPath) => this.switchSessionHandler(sessionPath),
reload: () => this.reloadHandler(),
// Use property descriptors instead of object spread so the guarded getters from
// createContext() stay lazy. A spread would eagerly read them once and freeze the
// old values into the returned object, bypassing stale-instance checks.
const context = Object.defineProperties(
{},
Object.getOwnPropertyDescriptors(this.createContext()),
) as ExtensionCommandContext;
context.waitForIdle = () => {
this.assertActive();
return this.waitForIdleFn();
};
context.newSession = (options) => {
this.assertActive();
return this.newSessionHandler(options);
};
context.fork = (entryId, options) => {
this.assertActive();
return this.forkHandler(entryId, options);
};
context.navigateTree = (targetId, options) => {
this.assertActive();
return this.navigateTreeHandler(targetId, options);
};
context.switchSession = (sessionPath, options) => {
this.assertActive();
return this.switchSessionHandler(sessionPath, options);
};
context.reload = () => {
this.assertActive();
return this.reloadHandler();
};
return context;
}
private isSessionBeforeEvent(event: RunnerEmitEvent): event is SessionBeforeEvent {

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

@@ -326,10 +326,14 @@ export interface ExtensionCommandContext extends ExtensionContext {
newSession(options?: {
parentSession?: string;
setup?: (sessionManager: SessionManager) => Promise<void>;
withSession?: (ctx: ReplacedSessionContext) => Promise<void>;
}): Promise<{ cancelled: boolean }>;
/** Fork from a specific entry, creating a new session file. */
fork(entryId: string, options?: { position?: "before" | "at" }): Promise<{ cancelled: boolean }>;
fork(
entryId: string,
options?: { position?: "before" | "at"; withSession?: (ctx: ReplacedSessionContext) => Promise<void> },
): Promise<{ cancelled: boolean }>;
/** Navigate to a different point in the session tree. */
navigateTree(
@@ -338,12 +342,32 @@ export interface ExtensionCommandContext extends ExtensionContext {
): Promise<{ cancelled: boolean }>;
/** Switch to a different session file. */
switchSession(sessionPath: string): Promise<{ cancelled: boolean }>;
switchSession(
sessionPath: string,
options?: { withSession?: (ctx: ReplacedSessionContext) => Promise<void> },
): Promise<{ cancelled: boolean }>;
/** Reload extensions, skills, prompts, and themes. */
reload(): Promise<void>;
}
/**
* Fresh command-capable context bound to the replacement session after a session switch.
*
* This is passed to `withSession()` callbacks on `newSession()`, `fork()`, and `switchSession()`.
*/
export interface ReplacedSessionContext extends ExtensionCommandContext {
sendMessage<T = unknown>(
message: Pick<CustomMessage<T>, "customType" | "content" | "display" | "details">,
options?: { triggerTurn?: boolean; deliverAs?: "steer" | "followUp" | "nextTurn" },
): Promise<void>;
sendUserMessage(
content: string | (TextContent | ImageContent)[],
options?: { deliverAs?: "steer" | "followUp" },
): Promise<void>;
}
// ============================================================================
// Tool Types
// ============================================================================
@@ -1395,6 +1419,10 @@ export interface ExtensionRuntimeState {
flagValues: Map<string, boolean | string>;
/** Provider registrations queued during extension loading, processed when runner binds */
pendingProviderRegistrations: Array<{ name: string; config: ProviderConfig; extensionPath: string }>;
/** Throws when this extension instance is stale after runtime replacement. */
assertActive: () => void;
/** Marks this extension instance as stale after runtime replacement or reload. */
invalidate: (message?: string) => void;
/**
* Register or unregister a provider.
*
@@ -1451,13 +1479,20 @@ export interface ExtensionCommandContextActions {
newSession: (options?: {
parentSession?: string;
setup?: (sessionManager: SessionManager) => Promise<void>;
withSession?: (ctx: ReplacedSessionContext) => Promise<void>;
}) => Promise<{ cancelled: boolean }>;
fork: (entryId: string, options?: { position?: "before" | "at" }) => Promise<{ cancelled: boolean }>;
fork: (
entryId: string,
options?: { position?: "before" | "at"; withSession?: (ctx: ReplacedSessionContext) => Promise<void> },
) => Promise<{ cancelled: boolean }>;
navigateTree: (
targetId: string,
options?: { summarize?: boolean; customInstructions?: string; replaceInstructions?: boolean; label?: string },
) => Promise<{ cancelled: boolean }>;
switchSession: (sessionPath: string) => Promise<{ cancelled: boolean }>;
switchSession: (
sessionPath: string,
options?: { withSession?: (ctx: ReplacedSessionContext) => Promise<void> },
) => Promise<{ cancelled: boolean }>;
reload: () => Promise<void>;
}

23
packages/coding-agent/src/modes/print-mode.ts
View File

@@ -64,23 +64,18 @@ export async function runPrintMode(runtimeHost: AgentSessionRuntime, options: Pr
registerSignalHandlers();
runtimeHost.setRebindSession(async () => {
await rebindSession();
});
const rebindSession = async (): Promise<void> => {
session = runtimeHost.session;
await session.bindExtensions({
commandContextActions: {
waitForIdle: () => session.agent.waitForIdle(),
newSession: async (newSessionOptions) => {
const result = await runtimeHost.newSession(newSessionOptions);
if (!result.cancelled) {
await rebindSession();
}
return result;
},
newSession: async (newSessionOptions) => runtimeHost.newSession(newSessionOptions),
fork: async (entryId, forkOptions) => {
const result = await runtimeHost.fork(entryId, forkOptions);
if (!result.cancelled) {
await rebindSession();
}
return { cancelled: result.cancelled };
},
navigateTree: async (targetId, navigateOptions) => {
@@ -92,12 +87,8 @@ export async function runPrintMode(runtimeHost: AgentSessionRuntime, options: Pr
});
return { cancelled: result.cancelled };
},
switchSession: async (sessionPath) => {
const result = await runtimeHost.switchSession(sessionPath);
if (!result.cancelled) {
await rebindSession();
}
return result;
switchSession: async (sessionPath, switchOptions) => {
return runtimeHost.switchSession(sessionPath, switchOptions);
},
reload: async () => {
await session.reload();

23
packages/coding-agent/src/modes/rpc/rpc-mode.ts
View File

@@ -290,24 +290,19 @@ export async function runRpcMode(runtimeHost: AgentSessionRuntime): Promise<neve
},
});
runtimeHost.setRebindSession(async () => {
await rebindSession();
});
const rebindSession = async (): Promise<void> => {
session = runtimeHost.session;
await session.bindExtensions({
uiContext: createExtensionUIContext(),
commandContextActions: {
waitForIdle: () => session.agent.waitForIdle(),
newSession: async (options) => {
const result = await runtimeHost.newSession(options);
if (!result.cancelled) {
await rebindSession();
}
return result;
},
newSession: async (options) => runtimeHost.newSession(options),
fork: async (entryId, forkOptions) => {
const result = await runtimeHost.fork(entryId, forkOptions);
if (!result.cancelled) {
await rebindSession();
}
return { cancelled: result.cancelled };
},
navigateTree: async (targetId, options) => {
@@ -319,12 +314,8 @@ export async function runRpcMode(runtimeHost: AgentSessionRuntime): Promise<neve
});
return { cancelled: result.cancelled };
},
switchSession: async (sessionPath) => {
const result = await runtimeHost.switchSession(sessionPath);
if (!result.cancelled) {
await rebindSession();
}
return result;
switchSession: async (sessionPath, options) => {
return runtimeHost.switchSession(sessionPath, options);
},
reload: async () => {
await session.reload();

274
packages/coding-agent/test/suite/regressions/2860-replaced-session-context.test.ts Normal file
View File

@@ -0,0 +1,274 @@
import { existsSync, mkdirSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { fauxAssistantMessage, registerFauxProvider } from "@mariozechner/pi-ai";
import { afterEach, describe, expect, it } from "vitest";
import type { AgentSession } from "../../../src/core/agent-session.js";
import {
type CreateAgentSessionRuntimeFactory,
createAgentSessionFromServices,
createAgentSessionRuntime,
createAgentSessionServices,
} from "../../../src/core/agent-session-runtime.js";
import { AuthStorage } from "../../../src/core/auth-storage.js";
import { SessionManager } from "../../../src/core/session-manager.js";
import type { ExtensionAPI, ExtensionCommandContext, ExtensionFactory } from "../../../src/index.js";
function getText(message: AgentSession["messages"][number]): string {
if (!("content" in message)) {
return "";
}
return typeof message.content === "string"
? message.content
: message.content
.filter((part): part is { type: "text"; text: string } => part.type === "text")
.map((part) => part.text)
.join("");
}
describe("regression #2860: replaced session callbacks", () => {
const cleanups: Array<() => Promise<void> | void> = [];
afterEach(async () => {
while (cleanups.length > 0) {
await cleanups.pop()?.();
}
});
async function createRuntimeForTest(extensionFactory: ExtensionFactory, responses: string[]) {
const tempDir = join(tmpdir(), `pi-2860-${Date.now()}-${Math.random().toString(36).slice(2)}`);
mkdirSync(tempDir, { recursive: true });
const faux = registerFauxProvider({
models: [{ id: "faux-1", reasoning: false }],
});
faux.setResponses(responses.map((response) => fauxAssistantMessage(response)));
const authStorage = AuthStorage.inMemory();
authStorage.setRuntimeApiKey(faux.getModel().provider, "faux-key");
const createRuntime: CreateAgentSessionRuntimeFactory = async ({ cwd, sessionManager, sessionStartEvent }) => {
const services = await createAgentSessionServices({
cwd,
agentDir: tempDir,
authStorage,
resourceLoaderOptions: {
extensionFactories: [
(pi: ExtensionAPI) => {
pi.registerProvider(faux.getModel().provider, {
baseUrl: faux.getModel().baseUrl,
apiKey: "faux-key",
api: faux.api,
models: faux.models.map((registeredModel) => ({
id: registeredModel.id,
name: registeredModel.name,
api: registeredModel.api,
reasoning: registeredModel.reasoning,
input: registeredModel.input,
cost: registeredModel.cost,
contextWindow: registeredModel.contextWindow,
maxTokens: registeredModel.maxTokens,
})),
});
extensionFactory(pi);
},
],
noSkills: true,
noPromptTemplates: true,
noThemes: true,
},
});
return {
...(await createAgentSessionFromServices({
services,
sessionManager,
sessionStartEvent,
model: faux.getModel(),
})),
services,
diagnostics: services.diagnostics,
};
};
const runtime = await createAgentSessionRuntime(createRuntime, {
cwd: tempDir,
agentDir: tempDir,
sessionManager: SessionManager.create(tempDir),
});
const rebindSession = async (): Promise<void> => {
const session = runtime.session;
await session.bindExtensions({
commandContextActions: {
waitForIdle: () => session.agent.waitForIdle(),
newSession: async (options) => runtime.newSession(options),
fork: async (entryId, options) => {
const result = await runtime.fork(entryId, options);
return { cancelled: result.cancelled };
},
navigateTree: async (targetId, options) => {
const result = await session.navigateTree(targetId, {
summarize: options?.summarize,
customInstructions: options?.customInstructions,
replaceInstructions: options?.replaceInstructions,
label: options?.label,
});
return { cancelled: result.cancelled };
},
switchSession: async (sessionPath, options) => runtime.switchSession(sessionPath, options),
reload: async () => {
await session.reload();
},
},
});
};
runtime.setRebindSession(async () => {
await rebindSession();
});
await rebindSession();
cleanups.push(async () => {
await runtime.dispose();
faux.unregister();
if (existsSync(tempDir)) {
rmSync(tempDir, { recursive: true, force: true });
}
});
return { runtime, faux };
}
it("rebinds before withSession, targets the replacement session, and invalidates stale pi/ctx", async () => {
const events: string[] = [];
let oldCtx: ExtensionCommandContext | undefined;
let oldPi: ExtensionAPI | undefined;
let oldSessionFile: string | undefined;
let staleCtxThrows = false;
let stalePiThrows = false;
let replacementSessionFile: string | undefined;
let instanceId = 0;
const { runtime } = await createRuntimeForTest(
(pi) => {
const currentInstance = ++instanceId;
pi.on("session_start", () => {
events.push(`start:${currentInstance}`);
});
pi.on("session_shutdown", () => {
events.push(`shutdown:${currentInstance}`);
});
pi.registerCommand("repro", {
description: "repro",
handler: async (_args, ctx) => {
oldCtx = ctx;
oldPi = pi;
oldSessionFile = ctx.sessionManager.getSessionFile();
await ctx.newSession({
parentSession: oldSessionFile,
withSession: async (replacedCtx) => {
events.push(`with:${currentInstance}`);
replacementSessionFile = replacedCtx.sessionManager.getSessionFile();
try {
oldCtx?.sessionManager.getSessionFile();
} catch {
staleCtxThrows = true;
}
try {
oldPi?.sendUserMessage("stale message");
} catch {
stalePiThrows = true;
}
await replacedCtx.sendUserMessage("Hello from the new session!");
},
});
},
});
},
["hello reply"],
);
expect(events).toEqual(["start:1"]);
await runtime.session.prompt("/repro");
expect(events).toEqual(["start:1", "shutdown:1", "start:2", "with:1"]);
expect(replacementSessionFile).toBeDefined();
expect(replacementSessionFile).not.toBe(oldSessionFile);
expect(staleCtxThrows).toBe(true);
expect(stalePiThrows).toBe(true);
expect(runtime.session.messages.map((message) => `${message.role}:${getText(message)}`)).toEqual([
"user:Hello from the new session!",
"assistant:hello reply",
]);
});
it("supports withSession for fork", async () => {
const { runtime } = await createRuntimeForTest(
(pi) => {
pi.registerCommand("fork-it", {
description: "fork-it",
handler: async (_args, ctx) => {
const leafId = ctx.sessionManager.getLeafId();
if (!leafId) {
throw new Error("Missing leaf id");
}
await ctx.fork(leafId, {
position: "at",
withSession: async (replacedCtx) => {
await replacedCtx.sendUserMessage("fork callback message");
},
});
},
});
},
["seed reply", "fork reply"],
);
await runtime.session.prompt("seed");
await runtime.session.prompt("/fork-it");
expect(runtime.session.messages.map((message) => `${message.role}:${getText(message)}`)).toEqual([
"user:seed",
"assistant:seed reply",
"user:fork callback message",
"assistant:fork reply",
]);
});
it("supports withSession for switchSession", async () => {
let targetSessionPath = "";
const { runtime } = await createRuntimeForTest(
(pi) => {
pi.registerCommand("switch-it", {
description: "switch-it",
handler: async (_args, ctx) => {
await ctx.switchSession(targetSessionPath, {
withSession: async (replacedCtx) => {
await replacedCtx.sendUserMessage("switch callback message");
},
});
},
});
},
["root reply", "target reply", "switch reply"],
);
await runtime.session.prompt("root");
const originalSessionPath = runtime.session.sessionFile;
const newSessionResult = await runtime.newSession();
expect(newSessionResult.cancelled).toBe(false);
await runtime.session.prompt("target");
targetSessionPath = runtime.session.sessionFile!;
await runtime.switchSession(originalSessionPath!);
await runtime.session.prompt("/switch-it");
expect(runtime.session.sessionFile).toBe(targetSessionPath);
expect(runtime.session.messages.map((message) => `${message.role}:${getText(message)}`)).toEqual([
"user:target",
"assistant:target reply",
"user:switch callback message",
"assistant:switch reply",
]);
});
});