diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index f78a8bf5..37374a77 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/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 diff --git a/packages/coding-agent/docs/extensions.md b/packages/coding-agent/docs/extensions.md index de771026..b5cd5d9f 100644 --- a/packages/coding-agent/docs/extensions.md +++ b/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`. diff --git a/packages/coding-agent/src/core/agent-session-runtime.ts b/packages/coding-agent/src/core/agent-session-runtime.ts index 74b8e430..72d937d5 100644 --- a/packages/coding-agent/src/core/agent-session-runtime.ts +++ b/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; + constructor( private _session: AgentSession, private _services: AgentSessionServices, @@ -93,6 +95,10 @@ export class AgentSessionRuntime { return this._modelFallbackMessage; } + setRebindSession(rebindSession?: (session: AgentSession) => Promise): 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): Promise { + 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 }, + ): 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; + withSession?: (ctx: ReplacedSessionContext) => Promise; }): 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 }, ): 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 }; } diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index f90a4c8d..550b4858 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/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. */ diff --git a/packages/coding-agent/src/core/extensions/index.ts b/packages/coding-agent/src/core/extensions/index.ts index 339cbe9c..0820f47b 100644 --- a/packages/coding-agent/src/core/extensions/index.ts +++ b/packages/coding-agent/src/core/extensions/index.ts @@ -103,6 +103,7 @@ export type { // Commands RegisteredCommand, RegisteredTool, + ReplacedSessionContext, ResolvedCommand, // Events - Resources ResourcesDiscoverEvent, diff --git a/packages/coding-agent/src/core/extensions/loader.ts b/packages/coding-agent/src/core/extensions/loader.ts index 667a1d7d..5981aed0 100644 --- a/packages/coding-agent/src/core/extensions/loader.ts +++ b/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 = "") => { @@ -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): 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 { + 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(customType: string, renderer: MessageRenderer): 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); }, diff --git a/packages/coding-agent/src/core/extensions/runner.ts b/packages/coding-agent/src/core/extensions/runner.ts index 39d9a9a5..7ce0a2c4 100644 --- a/packages/coding-agent/src/core/extensions/runner.ts +++ b/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; + withSession?: (ctx: ReplacedSessionContext) => Promise; }) => Promise<{ cancelled: boolean }>; export type ForkHandler = ( entryId: string, - options?: { position?: "before" | "at" }, + options?: { position?: "before" | "at"; withSession?: (ctx: ReplacedSessionContext) => Promise }, ) => 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 }, +) => Promise<{ cancelled: boolean }>; export type ReloadHandler = () => Promise; @@ -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 { diff --git a/packages/coding-agent/src/core/extensions/types.ts b/packages/coding-agent/src/core/extensions/types.ts index 2517a34c..185195b8 100644 --- a/packages/coding-agent/src/core/extensions/types.ts +++ b/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; + withSession?: (ctx: ReplacedSessionContext) => Promise; }): 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 }, + ): 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 }, + ): Promise<{ cancelled: boolean }>; /** Reload extensions, skills, prompts, and themes. */ reload(): Promise; } +/** + * 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( + message: Pick, "customType" | "content" | "display" | "details">, + options?: { triggerTurn?: boolean; deliverAs?: "steer" | "followUp" | "nextTurn" }, + ): Promise; + + sendUserMessage( + content: string | (TextContent | ImageContent)[], + options?: { deliverAs?: "steer" | "followUp" }, + ): Promise; +} + // ============================================================================ // Tool Types // ============================================================================ @@ -1395,6 +1419,10 @@ export interface ExtensionRuntimeState { flagValues: Map; /** 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; + withSession?: (ctx: ReplacedSessionContext) => Promise; }) => Promise<{ cancelled: boolean }>; - fork: (entryId: string, options?: { position?: "before" | "at" }) => Promise<{ cancelled: boolean }>; + fork: ( + entryId: string, + options?: { position?: "before" | "at"; withSession?: (ctx: ReplacedSessionContext) => Promise }, + ) => 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 }, + ) => Promise<{ cancelled: boolean }>; reload: () => Promise; } diff --git a/packages/coding-agent/src/modes/print-mode.ts b/packages/coding-agent/src/modes/print-mode.ts index 0c79956c..bee517da 100644 --- a/packages/coding-agent/src/modes/print-mode.ts +++ b/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 => { 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(); diff --git a/packages/coding-agent/src/modes/rpc/rpc-mode.ts b/packages/coding-agent/src/modes/rpc/rpc-mode.ts index 7b5d5c90..bfe49703 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-mode.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-mode.ts @@ -290,24 +290,19 @@ export async function runRpcMode(runtimeHost: AgentSessionRuntime): Promise { + await rebindSession(); + }); + const rebindSession = async (): Promise => { 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 { - 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(); diff --git a/packages/coding-agent/test/suite/regressions/2860-replaced-session-context.test.ts b/packages/coding-agent/test/suite/regressions/2860-replaced-session-context.test.ts new file mode 100644 index 00000000..626365c6 --- /dev/null +++ b/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> = []; + + 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 => { + 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", + ]); + }); +});