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

packages/coding-agent/CHANGELOG.md

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

packages/coding-agent/docs/extensions.md

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

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

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

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

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

packages/coding-agent/src/core/extensions/index.ts

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

packages/coding-agent/src/core/extensions/loader.ts

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

packages/coding-agent/src/core/extensions/runner.ts

@@ -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,
+			get ui() {
-			hasUI: this.hasUI(),
+				runner.assertActive();
-			cwd: this.cwd,
+				return runner.uiContext;
-			sessionManager: this.sessionManager,
+			},
-			modelRegistry: this.modelRegistry,
+			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(),
+			isIdle: () => {
-			signal: this.getSignalFn(),
+				runner.assertActive();
-			abort: () => this.abortFn(),
+				return runner.isIdleFn();
-			hasPendingMessages: () => this.hasPendingMessagesFn(),
+			},
-			shutdown: () => this.shutdownHandler(),
+			get signal() {
-			getContextUsage: () => this.getContextUsageFn(),
+				runner.assertActive();
-			compact: (options) => this.compactFn(options),
+				return runner.getSignalFn();
-			getSystemPrompt: () => this.getSystemPromptFn(),
+			},
+			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 {
+		// Use property descriptors instead of object spread so the guarded getters from
-			...this.createContext(),
+		// createContext() stay lazy. A spread would eagerly read them once and freeze the
-			waitForIdle: () => this.waitForIdleFn(),
+		// old values into the returned object, bypassing stale-instance checks.
-			newSession: (options) => this.newSessionHandler(options),
+		const context = Object.defineProperties(
-			fork: (entryId, options) => this.forkHandler(entryId, options),
+			{},
-			navigateTree: (targetId, options) => this.navigateTreeHandler(targetId, options),
+			Object.getOwnPropertyDescriptors(this.createContext()),
-			switchSession: (sessionPath) => this.switchSessionHandler(sessionPath),
+		) as ExtensionCommandContext;
-			reload: () => this.reloadHandler(),
+		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 {

packages/coding-agent/src/core/extensions/types.ts

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

packages/coding-agent/src/modes/print-mode.ts

@@ -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) => {
+				newSession: async (newSessionOptions) => runtimeHost.newSession(newSessionOptions),
-					const result = await runtimeHost.newSession(newSessionOptions);
-					if (!result.cancelled) {
-						await rebindSession();
-					}
-					return result;
-				},
 				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) => {
+				switchSession: async (sessionPath, switchOptions) => {
-					const result = await runtimeHost.switchSession(sessionPath);
+					return runtimeHost.switchSession(sessionPath, switchOptions);
-					if (!result.cancelled) {
-						await rebindSession();
-					}
-					return result;
 				},
 				reload: async () => {
 					await session.reload();

packages/coding-agent/src/modes/rpc/rpc-mode.ts

@@ -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) => {
+				newSession: async (options) => runtimeHost.newSession(options),
-					const result = await runtimeHost.newSession(options);
-					if (!result.cancelled) {
-						await rebindSession();
-					}
-					return result;
-				},
 				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) => {
+				switchSession: async (sessionPath, options) => {
-					const result = await runtimeHost.switchSession(sessionPath);
+					return runtimeHost.switchSession(sessionPath, options);
-					if (!result.cancelled) {
-						await rebindSession();
-					}
-					return result;
 				},
 				reload: async () => {
 					await session.reload();

packages/coding-agent/test/suite/regressions/2860-replaced-session-context.test.ts

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