From ef03c1328f6953ede55fcc4879c8d3bfae0b2a0f Mon Sep 17 00:00:00 2001 From: Armin Ronacher Date: Sun, 26 Jul 2026 15:44:38 +0200 Subject: [PATCH] fix(coding-agent): harden extension unload on reload --- .../src/core/agent-session-services.ts | 5 +- .../coding-agent/src/core/agent-session.ts | 77 +++++++++++++--- .../src/core/extensions/loader.ts | 28 +++++- .../coding-agent/src/core/extensions/types.ts | 2 + .../coding-agent/src/core/model-runtime.ts | 81 +++++++++++++++-- .../src/modes/interactive/interactive-mode.ts | 23 +++-- .../agent-session-dynamic-provider.test.ts | 87 +++++++++++++++++++ .../test/extensions-runner.test.ts | 62 ++++++++++++- .../agent-session-model-extension.test.ts | 69 ++++++++++++++- .../suite/session-loadout-interactive.test.ts | 17 ++++ 10 files changed, 418 insertions(+), 33 deletions(-) diff --git a/packages/coding-agent/src/core/agent-session-services.ts b/packages/coding-agent/src/core/agent-session-services.ts index 495154b29..6ed5e2805 100644 --- a/packages/coding-agent/src/core/agent-session-services.ts +++ b/packages/coding-agent/src/core/agent-session-services.ts @@ -153,9 +153,10 @@ export async function createAgentSessionServices( const diagnostics: AgentSessionRuntimeDiagnostic[] = []; const extensionsResult = resourceLoader.getExtensions(); + modelRuntime.resetSessionExtensionProviders(); for (const { name, config, extensionPath } of extensionsResult.runtime.pendingProviderRegistrations) { try { - modelRuntime.registerProvider(name, config); + modelRuntime.registerSessionExtensionProvider(name, config); } catch (error) { const message = error instanceof Error ? error.message : String(error); diagnostics.push({ @@ -167,7 +168,7 @@ export async function createAgentSessionServices( extensionsResult.runtime.pendingProviderRegistrations = []; for (const { provider, extensionPath } of extensionsResult.runtime.pendingNativeProviderRegistrations) { try { - modelRuntime.registerNativeProvider(provider); + modelRuntime.registerSessionNativeExtensionProvider(provider); } catch (error) { const message = error instanceof Error ? error.message : String(error); diagnostics.push({ diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index a688528ce..3f77abee4 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -256,6 +256,11 @@ export interface ModelCycleResult { isScoped: boolean; } +export interface ReloadResult { + /** User-facing diagnostic when reload removed the active model. */ + modelFallbackMessage?: string; +} + /** Session statistics for /session command */ export interface SessionStats { sessionFile: string | undefined; @@ -2316,16 +2321,45 @@ export class AgentSession { private _refreshCurrentModelFromRegistry(): void { const currentModel = this.model; - if (!currentModel) { - return; - } - + if (!currentModel) return; const refreshedModel = this._modelRuntime.getModel(currentModel.provider, currentModel.id); - if (!refreshedModel || refreshedModel === currentModel) { - return; + if (refreshedModel && refreshedModel !== currentModel) this.agent.state.model = refreshedModel; + } + + private _reconcileModelAfterReload(previousModel: Model | undefined): string | undefined { + this._scopedModels = this._scopedModels.flatMap((scoped) => { + const model = this._modelRuntime.getModel(scoped.model.provider, scoped.model.id); + return model ? [{ ...scoped, model }] : []; + }); + if (!previousModel) return undefined; + + const refreshedModel = this._modelRuntime.getModel(previousModel.provider, previousModel.id); + if (refreshedModel) { + this.agent.state.model = refreshedModel; + return undefined; } - this.agent.state.model = refreshedModel; + const fallback = this._modelRuntime.getAvailableSnapshot()[0]; + if (fallback) { + this.agent.state.model = fallback; + this.agent.state.thinkingLevel = clampThinkingLevel(fallback, this.thinkingLevel) as ThinkingLevel; + return `Active model "${previousModel.provider}/${previousModel.id}" is no longer available after reload. Switched to "${fallback.provider}/${fallback.id}".`; + } + + this.agent.state.model = { + ...previousModel, + id: "unknown", + name: "unknown", + api: "unknown", + provider: "unknown", + baseUrl: "", + reasoning: false, + input: [], + contextWindow: 0, + maxTokens: 0, + }; + this.agent.state.thinkingLevel = "off"; + return `Active model "${previousModel.provider}/${previousModel.id}" is no longer available after reload. No configured model is available; use /model or /login to select one.`; } private _bindExtensionCore(runner: ExtensionRunner): void { @@ -2436,15 +2470,15 @@ export class AgentSession { }, { registerProvider: (name, config) => { - this._modelRuntime.registerProvider(name, config); + this._modelRuntime.registerSessionExtensionProvider(name, config); this._refreshCurrentModelFromRegistry(); }, registerNativeProvider: (provider) => { - this._modelRuntime.registerNativeProvider(provider); + this._modelRuntime.registerSessionNativeExtensionProvider(provider); this._refreshCurrentModelFromRegistry(); }, unregisterProvider: (name) => { - this._modelRuntime.unregisterProvider(name); + this._modelRuntime.unregisterSessionExtensionProvider(name); this._refreshCurrentModelFromRegistry(); }, }, @@ -2598,15 +2632,28 @@ export class AgentSession { }); } - async reload(options?: { beforeSessionStart?: () => void | Promise }): Promise { - const previousFlagValues = this._extensionRunner.getFlagValues(); - await emitSessionShutdownEvent(this._extensionRunner, { type: "session_shutdown", reason: "reload" }); + async reload(options?: { + beforeExtensionInvalidate?: () => void; + beforeSessionStart?: () => void | Promise; + }): Promise { + const oldRunner = this._extensionRunner; + const previousFlagValues = oldRunner.getFlagValues(); + const previousActiveToolNames = this.getActiveToolNames(); + const previousModel = this.model; + await emitSessionShutdownEvent(oldRunner, { type: "session_shutdown", reason: "reload" }); + try { + options?.beforeExtensionInvalidate?.(); + } finally { + oldRunner.invalidate(); + } + await this.settingsManager.reload(); this.syncQueueModesFromSettings(); resetApiProviders(); + this._modelRuntime.resetSessionExtensionProviders(); await this._resourceLoader.reload(); this._buildRuntime({ - activeToolNames: this.getActiveToolNames(), + activeToolNames: previousActiveToolNames, flagValues: previousFlagValues, includeAllExtensionTools: true, }); @@ -2621,6 +2668,8 @@ export class AgentSession { await this._extensionRunner.emit({ type: "session_start", reason: "reload" }); await this.extendResourcesFromExtensions("reload"); } + await this._modelRuntime.refresh({ allowNetwork: false }); + return { modelFallbackMessage: this._reconcileModelAfterReload(previousModel) }; } // ========================================================================= diff --git a/packages/coding-agent/src/core/extensions/loader.ts b/packages/coding-agent/src/core/extensions/loader.ts index 89f59c03a..c12aab036 100644 --- a/packages/coding-agent/src/core/extensions/loader.ts +++ b/packages/coding-agent/src/core/extensions/loader.ts @@ -172,6 +172,7 @@ export function createExtensionRuntime(): ExtensionRuntime { throw new Error("Extension runtime not initialized. Action methods cannot be called during extension loading."); }; const state: { staleMessage?: string } = {}; + const eventBusUnsubscribers = new Set<() => void>(); const assertActive = () => { if (state.staleMessage) { throw new Error(state.staleMessage); @@ -199,9 +200,23 @@ export function createExtensionRuntime(): ExtensionRuntime { pendingNativeProviderRegistrations: [], assertActive, invalidate: (message) => { - state.staleMessage ??= + if (state.staleMessage) return; + state.staleMessage = message ?? "This extension ctx is stale after session replacement or reload. Do not use a captured pi or command ctx after ctx.newSession(), ctx.fork(), ctx.switchSession(), or ctx.reload(). For newSession, fork, and switchSession, move post-replacement work into withSession and use the ctx passed to withSession. For reload, do not use the old ctx after await ctx.reload()."; + for (const unsubscribe of eventBusUnsubscribers) unsubscribe(); + eventBusUnsubscribers.clear(); + }, + trackEventBusSubscription: (unsubscribe) => { + let active = true; + const trackedUnsubscribe = () => { + if (!active) return; + active = false; + eventBusUnsubscribers.delete(trackedUnsubscribe); + unsubscribe(); + }; + eventBusUnsubscribers.add(trackedUnsubscribe); + return trackedUnsubscribe; }, // Pre-bind: queue registrations so bindCore() can flush them once the // model registry is available. bindCore() replaces both with direct calls. @@ -386,7 +401,16 @@ function createExtensionAPI( runtime.unregisterProvider(name, extension.path); }, - events: eventBus, + events: { + emit(channel, data) { + runtime.assertActive(); + eventBus.emit(channel, data); + }, + on(channel, handler) { + runtime.assertActive(); + return runtime.trackEventBusSubscription(eventBus.on(channel, handler)); + }, + }, } as ExtensionAPI; return api; diff --git a/packages/coding-agent/src/core/extensions/types.ts b/packages/coding-agent/src/core/extensions/types.ts index c8989a59c..f87f040d7 100644 --- a/packages/coding-agent/src/core/extensions/types.ts +++ b/packages/coding-agent/src/core/extensions/types.ts @@ -1577,6 +1577,8 @@ export interface ExtensionRuntimeState { assertActive: () => void; /** Marks this extension instance as stale after runtime replacement or reload. */ invalidate: (message?: string) => void; + /** Track an event-bus subscription so invalidating this runtime removes it. */ + trackEventBusSubscription: (unsubscribe: () => void) => () => void; /** * Register or unregister a provider. * diff --git a/packages/coding-agent/src/core/model-runtime.ts b/packages/coding-agent/src/core/model-runtime.ts index e4575b07f..ba75abc90 100644 --- a/packages/coding-agent/src/core/model-runtime.ts +++ b/packages/coding-agent/src/core/model-runtime.ts @@ -98,6 +98,10 @@ export class ModelRuntime implements Models { private readonly builtins = new Map(); private readonly nativeExtensionProviders = new Map(); private readonly extensionProviders = new Map(); + private readonly sessionExtensionProviders = new Map< + string, + { nativeProvider?: Provider; config?: ProviderConfigInput } + >(); private readonly compositionErrors = new Map(); private readonly modelsPath: string | undefined; private readonly modelNetworkEnabled: boolean; @@ -194,12 +198,26 @@ export class ModelRuntime implements Models { ...this.nativeExtensionProviders.keys(), ...this.config.getProviderIds(), ...this.extensionProviders.keys(), + ...this.sessionExtensionProviders.keys(), ]); } + private getEffectiveRegisteredProvider(providerId: string): { + nativeProvider?: Provider; + config?: ProviderConfigInput; + } { + const sessionProvider = this.sessionExtensionProviders.get(providerId); + if (sessionProvider) return sessionProvider; + return { + nativeProvider: this.nativeExtensionProviders.get(providerId), + config: this.extensionProviders.get(providerId), + }; + } + private recomposeProvider(providerId: string): void { - const base = this.nativeExtensionProviders.get(providerId) ?? this.builtins.get(providerId); - const extension = this.extensionProviders.get(providerId); + const registered = this.getEffectiveRegisteredProvider(providerId); + const base = registered.nativeProvider ?? this.builtins.get(providerId); + const extension = registered.config; if (!base && !this.config.getProvider(providerId) && !extension) { this.models.deleteProvider(providerId); this.compositionErrors.delete(providerId); @@ -343,15 +361,18 @@ export class ModelRuntime implements Models { } getRegisteredProviderConfig(providerId: string): ProviderConfigInput | undefined { - return this.extensionProviders.get(providerId); + return this.getEffectiveRegisteredProvider(providerId).config; } getRegisteredProviderIds(): readonly string[] { - return [...new Set([...this.extensionProviders.keys(), ...this.nativeExtensionProviders.keys()])]; + return [...this.providerIds()].filter((providerId) => { + const registered = this.getEffectiveRegisteredProvider(providerId); + return registered.config !== undefined || registered.nativeProvider !== undefined; + }); } getRegisteredNativeProvider(providerId: string): Provider | undefined { - return this.nativeExtensionProviders.get(providerId); + return this.getEffectiveRegisteredProvider(providerId).nativeProvider; } /** @internal Compatibility fallback for ModelRegistry when provider auth is unconfigured. */ @@ -359,7 +380,7 @@ export class ModelRuntime implements Models { return resolveCompatibilityRequestConfig( model, this.config.getProvider(model.provider), - this.extensionProviders.get(model.provider), + this.getEffectiveRegisteredProvider(model.provider).config, ); } @@ -383,7 +404,7 @@ export class ModelRuntime implements Models { const configuredHeaders = resolveConfiguredModelHeaders( providerOrModel, this.config.getProvider(providerOrModel.provider), - this.extensionProviders.get(providerOrModel.provider), + this.getEffectiveRegisteredProvider(providerOrModel.provider).config, { ...(resolution.env ?? {}), ...(overrides.env ?? {}) }, ); return { @@ -428,7 +449,7 @@ export class ModelRuntime implements Models { if (this.snapshot.storedProviders.has(providerId)) return { configured: true, source: "stored" }; const configured = configuredRequestAuthStatus( this.config.getProvider(providerId), - this.extensionProviders.get(providerId), + this.getEffectiveRegisteredProvider(providerId).config, ); if (configured) return configured; const check = this.snapshot.auth.get(providerId); @@ -590,4 +611,48 @@ export class ModelRuntime implements Models { this.updateModelSnapshot(); void this.refresh({ allowNetwork: false }); } + + /** Clear provider registrations owned by the current extension runtime. */ + resetSessionExtensionProviders(): void { + const providerIds = [...this.sessionExtensionProviders.keys()]; + this.sessionExtensionProviders.clear(); + for (const providerId of providerIds) this.recomposeProvider(providerId); + this.updateModelSnapshot(); + } + + /** Register a provider in the current extension-runtime layer. */ + registerSessionExtensionProvider(providerId: string, config: ProviderConfigInput): void { + const previous = this.getEffectiveRegisteredProvider(providerId).config; + validateExtensionProvider( + providerId, + this.getEffectiveRegisteredProvider(providerId).nativeProvider ?? this.builtins.get(providerId), + this.config.getProvider(providerId), + config, + ); + const effective: ProviderConfigInput = { ...previous }; + for (const [key, value] of Object.entries(config)) { + if (value !== undefined) (effective as Record)[key] = value; + } + this.sessionExtensionProviders.set(providerId, { config: effective }); + this.recomposeProvider(providerId); + this.updateModelSnapshot(); + void this.refresh({ allowNetwork: false }); + } + + /** Register a native provider in the current extension-runtime layer. */ + registerSessionNativeExtensionProvider(provider: Provider): void { + if (!provider.id.trim()) throw new Error("Provider id must not be empty."); + this.sessionExtensionProviders.set(provider.id, { nativeProvider: provider }); + this.recomposeProvider(provider.id); + this.updateModelSnapshot(); + void this.refresh({ allowNetwork: false }); + } + + /** Unregister a provider from the current extension-runtime layer. */ + unregisterSessionExtensionProvider(providerId: string): void { + this.sessionExtensionProviders.set(providerId, {}); + this.recomposeProvider(providerId); + this.updateModelSnapshot(); + void this.refresh({ allowNetwork: false }); + } } diff --git a/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/packages/coding-agent/src/modes/interactive/interactive-mode.ts index ae4212d3a..ce7cd05ed 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -5427,8 +5427,10 @@ export class InteractiveMode { this.showWarning("Wait for compaction to finish before reloading."); return; } - - this.resetExtensionUI(); + if (this.session.isBashRunning) { + this.showWarning("Wait for the bash command to finish before reloading."); + return; + } const reloadBox = new Container(); const borderColor = (s: string) => theme.fg("border", s); @@ -5444,7 +5446,6 @@ export class InteractiveMode { reloadBox.addChild(new Spacer(1)); reloadBox.addChild(new DynamicBorder(borderColor)); - const previousEditor = this.editor; this.editorContainer.clear(); this.editorContainer.addChild(reloadBox); this.ui.setFocus(reloadBox); @@ -5471,7 +5472,16 @@ export class InteractiveMode { }; try { - await this.session.reload({ beforeSessionStart: restoreChatBeforeSessionStart }); + const reloadResult = await this.session.reload({ + beforeExtensionInvalidate: () => { + this.resetExtensionUI(); + this.editorContainer.clear(); + this.editorContainer.addChild(reloadBox); + this.ui.setFocus(reloadBox); + this.ui.requestRender(); + }, + beforeSessionStart: restoreChatBeforeSessionStart, + }); restoreChatBeforeSessionStart(); configureHttpDispatcher(this.settingsManager.getHttpIdleTimeoutMs()); this.keybindings.reload(); @@ -5507,6 +5517,9 @@ export class InteractiveMode { if (modelsJsonError) { this.showError(`models.json error: ${modelsJsonError}`); } + if (reloadResult.modelFallbackMessage) { + this.showWarning(reloadResult.modelFallbackMessage); + } this.showStatus( savedImplicitProjectTrust ? "Reloaded keybindings, extensions, skills, prompts, themes, and context files; saved project trust" @@ -5516,7 +5529,7 @@ export class InteractiveMode { reloadBoxDismissed = true; } catch (error) { if (!reloadBoxDismissed) { - dismissReloadBox(previousEditor as Component); + dismissReloadBox(this.editor as Component); } this.showError(`Reload failed: ${error instanceof Error ? error.message : String(error)}`); } diff --git a/packages/coding-agent/test/agent-session-dynamic-provider.test.ts b/packages/coding-agent/test/agent-session-dynamic-provider.test.ts index c806efbc3..e54ba627f 100644 --- a/packages/coding-agent/test/agent-session-dynamic-provider.test.ts +++ b/packages/coding-agent/test/agent-session-dynamic-provider.test.ts @@ -176,4 +176,91 @@ describe("AgentSession dynamic provider registration", () => { session.dispose(); }); + + it("replays duplicate provider registrations from a clean extension layer", async () => { + let generation = 1; + const session = await createSession([ + (pi) => { + pi.registerProvider("reload-provider", { + baseUrl: `https://generation-${generation}.test`, + apiKey: "test-key", + api: "openai-completions", + ...(generation === 1 ? { headers: { "x-stale": "yes" } } : {}), + models: [ + { + id: `model-${generation}`, + name: `Model ${generation}`, + reasoning: false, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 128000, + maxTokens: 4096, + }, + ], + }); + }, + (pi) => { + pi.registerProvider("reload-provider", { name: `Winner ${generation}` }); + }, + ]); + await session.modelRuntime.refresh({ allowNetwork: false }); + + expect(session.modelRuntime.getRegisteredProviderConfig("reload-provider")).toMatchObject({ + name: "Winner 1", + headers: { "x-stale": "yes" }, + }); + expect(session.modelRuntime.getModel("reload-provider", "model-1")).toBeDefined(); + + generation = 2; + await session.reload(); + + expect(session.modelRuntime.getRegisteredProviderConfig("reload-provider")).toMatchObject({ + name: "Winner 2", + baseUrl: "https://generation-2.test", + }); + expect(session.modelRuntime.getRegisteredProviderConfig("reload-provider")?.headers).toBeUndefined(); + expect(session.modelRuntime.getModel("reload-provider", "model-1")).toBeUndefined(); + expect(session.modelRuntime.getModel("reload-provider", "model-2")).toBeDefined(); + session.dispose(); + }); + + it("removes disabled providers and falls back when their model was active", async () => { + let enabled = true; + const session = await createSession([ + (pi) => { + if (!enabled) return; + pi.registerProvider("temporary-provider", { + baseUrl: "https://temporary.test", + apiKey: "test-key", + api: "openai-completions", + models: [ + { + id: "temporary-model", + name: "Temporary Model", + reasoning: false, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 128000, + maxTokens: 4096, + }, + ], + }); + }, + ]); + await session.modelRuntime.refresh({ allowNetwork: false }); + const temporaryModel = session.modelRuntime.getModel("temporary-provider", "temporary-model")!; + await session.setModel(temporaryModel); + + enabled = false; + const result = await session.reload(); + + expect(session.modelRuntime.getProvider("temporary-provider")).toBeUndefined(); + expect(session.modelRuntime.getModel("temporary-provider", "temporary-model")).toBeUndefined(); + expect(session.model).toBeDefined(); + expect(session.model?.provider).not.toBe("temporary-provider"); + expect(result.modelFallbackMessage).toContain( + 'Active model "temporary-provider/temporary-model" is no longer available after reload.', + ); + session.dispose(); + }); }); diff --git a/packages/coding-agent/test/extensions-runner.test.ts b/packages/coding-agent/test/extensions-runner.test.ts index 69f20be90..0c6543b6b 100644 --- a/packages/coding-agent/test/extensions-runner.test.ts +++ b/packages/coding-agent/test/extensions-runner.test.ts @@ -8,7 +8,13 @@ import * as os from "node:os"; import * as path from "node:path"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { AuthStorage } from "../src/core/auth-storage.ts"; -import { createExtensionRuntime, discoverAndLoadExtensions, loadExtensions } from "../src/core/extensions/loader.ts"; +import { createEventBus } from "../src/core/event-bus.ts"; +import { + createExtensionRuntime, + discoverAndLoadExtensions, + loadExtensionFromFactory, + loadExtensions, +} from "../src/core/extensions/loader.ts"; import { ExtensionRunner, emitProjectTrustEvent } from "../src/core/extensions/runner.ts"; import type { ExtensionActions, @@ -817,6 +823,60 @@ describe("ExtensionRunner", () => { }); }); + describe("shared event bus", () => { + it("removes listeners owned by an invalidated runtime without breaking replacement communication", async () => { + const eventBus = createEventBus(); + const oldRuntime = createExtensionRuntime(); + let oldApi: Parameters[0]>[0] | undefined; + let peerApi: Parameters[0]>[0] | undefined; + let oldCalls = 0; + const oldExtension = await loadExtensionFromFactory( + (pi) => { + oldApi = pi; + pi.events.on("reload:test", () => oldCalls++); + }, + tempDir, + eventBus, + oldRuntime, + "", + ); + await loadExtensionFromFactory( + (pi) => { + peerApi = pi; + }, + tempDir, + eventBus, + oldRuntime, + "", + ); + const oldRunner = new ExtensionRunner([oldExtension], oldRuntime, tempDir, sessionManager, modelRegistry); + + peerApi!.events.emit("reload:test", undefined); + expect(oldCalls).toBe(1); + oldRunner.invalidate(); + expect(() => oldApi!.events.emit("reload:test", undefined)).toThrow( + "stale after session replacement or reload", + ); + + const newRuntime = createExtensionRuntime(); + let newApi: typeof oldApi; + let newCalls = 0; + await loadExtensionFromFactory( + (pi) => { + newApi = pi; + pi.events.on("reload:test", () => newCalls++); + }, + tempDir, + eventBus, + newRuntime, + "", + ); + newApi!.events.emit("reload:test", undefined); + expect(oldCalls).toBe(1); + expect(newCalls).toBe(1); + }); + }); + describe("provider registration", () => { it("bindCore ignores invalid queued registrations and reports extension error", async () => { const runtime = createExtensionRuntime(); diff --git a/packages/coding-agent/test/suite/agent-session-model-extension.test.ts b/packages/coding-agent/test/suite/agent-session-model-extension.test.ts index b334b1232..7d84838a4 100644 --- a/packages/coding-agent/test/suite/agent-session-model-extension.test.ts +++ b/packages/coding-agent/test/suite/agent-session-model-extension.test.ts @@ -2,7 +2,8 @@ import type { AgentTool, ThinkingLevel } from "@earendil-works/pi-agent-core"; import { fauxAssistantMessage, fauxToolCall, type Model, type Usage } from "@earendil-works/pi-ai"; import { Type } from "typebox"; import { afterEach, describe, expect, it } from "vitest"; -import type { BuildSystemPromptOptions, ExtensionAPI } from "../../src/index.ts"; +import type { BuildSystemPromptOptions, ExtensionAPI, ExtensionFactory, ResourceLoader } from "../../src/index.ts"; +import { createTestExtensionsResult } from "../utilities.ts"; import { createHarness, getAssistantTexts, type Harness } from "./harness.ts"; describe("AgentSession model and extension characterization", () => { @@ -394,4 +395,70 @@ describe("AgentSession model and extension characterization", () => { expect(lifecycleEvents).toEqual(["start:startup", "shutdown:reload", "start:reload"]); }); + + it("invalidates the old runtime after host cleanup and removes disabled extension resources", async () => { + let enabled = true; + const phases: string[] = []; + const factory: ExtensionFactory = (pi) => { + if (!enabled) return; + pi.on("session_shutdown", () => { + phases.push("shutdown"); + }); + pi.on("input", () => ({ action: "continue" })); + pi.registerCommand("removed-command", { handler: async () => {} }); + pi.registerShortcut("ctrl+shift+y", { handler: () => {} }); + pi.registerMessageRenderer("removed-message", () => undefined); + pi.registerEntryRenderer("removed-entry", () => undefined); + pi.registerTool({ + name: "removed_tool", + label: "Removed Tool", + description: "Removed on reload", + parameters: Type.Object({}), + execute: async () => ({ content: [{ type: "text", text: "ok" }], details: {} }), + }); + }; + let extensionsResult = await createTestExtensionsResult([factory]); + const resourceLoader: ResourceLoader = { + getExtensions: () => extensionsResult, + getSkills: () => ({ skills: [], diagnostics: [] }), + getPrompts: () => ({ prompts: [], diagnostics: [] }), + getThemes: () => ({ themes: [], diagnostics: [] }), + getAgentsFiles: () => ({ agentsFiles: [] }), + getSystemPrompt: () => undefined, + getAppendSystemPrompt: () => [], + extendResources: () => {}, + reload: async () => { + extensionsResult = await createTestExtensionsResult([factory]); + }, + }; + const harness = await createHarness({ resourceLoader }); + harnesses.push(harness); + await harness.session.bindExtensions({ shutdownHandler: () => {} }); + const oldRunner = harness.session.extensionRunner; + const oldContext = oldRunner.createContext(); + + expect(oldRunner.getCommand("removed-command")).toBeDefined(); + expect(harness.session.getAllTools().map((tool) => tool.name)).toContain("removed_tool"); + expect(oldRunner.hasHandlers("input")).toBe(true); + expect(oldRunner.getMessageRenderer("removed-message")).toBeDefined(); + expect(oldRunner.getEntryRenderer("removed-entry")).toBeDefined(); + + enabled = false; + await harness.session.reload({ + beforeExtensionInvalidate: () => { + phases.push("host-ui-cleanup"); + expect(oldContext.cwd).toBe(harness.tempDir); + }, + }); + + expect(phases).toEqual(["shutdown", "host-ui-cleanup"]); + expect(() => oldContext.cwd).toThrow("stale after session replacement or reload"); + expect(harness.session.extensionRunner.getCommand("removed-command")).toBeUndefined(); + expect(harness.session.getAllTools().map((tool) => tool.name)).not.toContain("removed_tool"); + expect(harness.session.getActiveToolNames()).not.toContain("removed_tool"); + expect(harness.session.extensionRunner.hasHandlers("input")).toBe(false); + expect(harness.session.extensionRunner.getShortcuts({}).size).toBe(0); + expect(harness.session.extensionRunner.getMessageRenderer("removed-message")).toBeUndefined(); + expect(harness.session.extensionRunner.getEntryRenderer("removed-entry")).toBeUndefined(); + }); }); diff --git a/packages/coding-agent/test/suite/session-loadout-interactive.test.ts b/packages/coding-agent/test/suite/session-loadout-interactive.test.ts index 377ddbbb3..f65fe7e88 100644 --- a/packages/coding-agent/test/suite/session-loadout-interactive.test.ts +++ b/packages/coding-agent/test/suite/session-loadout-interactive.test.ts @@ -37,10 +37,16 @@ type ApplyContext = { showWarning: (message: string) => void; }; +type ReloadBlockedContext = { + session: { isStreaming: boolean; isCompacting: boolean; isBashRunning: boolean }; + showWarning: (message: string) => void; +}; + type InteractiveModePrivate = { maybeRestoreSavedLoadout(this: RestoreContext): Promise; applySessionLoadout(this: ApplyContext, overrides: LoadoutOverride[], persist: boolean): Promise; loadoutActionBlocked(this: Pick, action: "opening" | "applying"): boolean; + handleReloadCommand(this: ReloadBlockedContext): Promise; }; const interactiveModePrototype = InteractiveMode.prototype as unknown as InteractiveModePrivate; @@ -123,6 +129,17 @@ describe("interactive session loadout restore", () => { }); }); +describe("interactive reload blocking", () => { + it("blocks ordinary reload while direct user bash is running", async () => { + const showWarning = vi.fn(); + await interactiveModePrototype.handleReloadCommand.call({ + session: { isStreaming: false, isCompacting: false, isBashRunning: true }, + showWarning, + }); + expect(showWarning).toHaveBeenCalledWith("Wait for the bash command to finish before reloading."); + }); +}); + describe("interactive /loadout application", () => { it("persists a deliberate change, reloads once, and reports unmatched resources", async () => { const harness = await createTestHarness();