fix(coding-agent): cache extension imports for session switches

Closes #5905.

packages/coding-agent/CHANGELOG.md

@@ -8,6 +8,7 @@
 
 ### Fixed
 
+- Fixed same-directory session switches to reuse imported extension modules while preserving fresh extension instances and lifecycle events ([#5905](https://github.com/earendil-works/pi/issues/5905)).
 - Fixed deep session branches taking quadratic time to build context or branch paths ([#5909](https://github.com/earendil-works/pi/issues/5909)).
 - Fixed fuzzy `edit` matches to preserve untouched line blocks instead of rewriting the whole file through normalized content ([#5899](https://github.com/earendil-works/pi/issues/5899)).
 - Fixed bash commands through legacy WSL `bash.exe` to pass scripts over stdin so shell variables expand in the target bash ([#5893](https://github.com/earendil-works/pi/issues/5893)).

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

@@ -117,6 +117,30 @@ function getAliases(): Record<string, string> {
 
 type HandlerFn = (...args: unknown[]) => Promise<unknown>;
 
+let extensionCacheCwd: string | undefined;
+let extensionCacheGeneration = 0;
+const extensionCache = new Map<string, ExtensionFactory>();
+
+interface ExtensionCacheToken {
+	cwd: string;
+	generation: number;
+}
+
+export function clearExtensionCache(): void {
+	extensionCache.clear();
+	extensionCacheCwd = undefined;
+	extensionCacheGeneration++;
+}
+
+function useExtensionCacheCwd(cwd: string): ExtensionCacheToken {
+	const resolvedCwd = resolvePath(cwd);
+	if (extensionCacheCwd !== undefined && extensionCacheCwd !== resolvedCwd) {
+		clearExtensionCache();
+	}
+	extensionCacheCwd = resolvedCwd;
+	return { cwd: resolvedCwd, generation: extensionCacheGeneration };
+}
+
 /**
  * Create a runtime with throwing stubs for action methods.
  * Runner.bindCore() replaces these with real implementations.
@@ -328,7 +352,22 @@ function createExtensionAPI(
 	return api;
 }
 
-async function loadExtensionModule(extensionPath: string) {
+function isCurrentCacheToken(cacheToken: ExtensionCacheToken | undefined): cacheToken is ExtensionCacheToken {
+	return (
+		cacheToken !== undefined &&
+		extensionCacheCwd === cacheToken.cwd &&
+		extensionCacheGeneration === cacheToken.generation
+	);
+}
+
+async function loadExtensionModule(extensionPath: string, cacheToken?: ExtensionCacheToken) {
+	if (isCurrentCacheToken(cacheToken)) {
+		const cachedFactory = extensionCache.get(extensionPath);
+		if (cachedFactory) {
+			return cachedFactory;
+		}
+	}
+
 	const jiti = createJiti(import.meta.url, {
 		moduleCache: false,
 		// In Bun binary: use virtualModules for bundled packages (no filesystem resolution)
@@ -339,7 +378,13 @@ async function loadExtensionModule(extensionPath: string) {
 
 	const module = await jiti.import(extensionPath, { default: true });
 	const factory = module as ExtensionFactory;
-	return typeof factory !== "function" ? undefined : factory;
+	if (typeof factory !== "function") {
+		return undefined;
+	}
+	if (isCurrentCacheToken(cacheToken)) {
+		extensionCache.set(extensionPath, factory);
+	}
+	return factory;
 }
 
 /**
@@ -370,11 +415,12 @@ async function loadExtension(
 	cwd: string,
 	eventBus: EventBus,
 	runtime: ExtensionRuntime,
+	cacheToken?: ExtensionCacheToken,
 ): Promise<{ extension: Extension | null; error: string | null }> {
 	const resolvedPath = resolvePath(extensionPath, cwd, { normalizeUnicodeSpaces: true });
 
 	try {
-		const factory = await loadExtensionModule(resolvedPath);
+		const factory = await loadExtensionModule(resolvedPath, cacheToken);
 		if (!factory) {
 			return { extension: null, error: `Extension does not export a valid factory function: ${extensionPath}` };
 		}
@@ -410,20 +456,28 @@ export async function loadExtensionFromFactory(
 /**
  * Load extensions from paths.
  */
-export async function loadExtensions(
+async function loadExtensionsInternal(
 	paths: string[],
 	cwd: string,
 	eventBus?: EventBus,
 	runtime?: ExtensionRuntime,
+	useCache = false,
 ): Promise<LoadExtensionsResult> {
 	const extensions: Extension[] = [];
 	const errors: Array<{ path: string; error: string }> = [];
-	const resolvedCwd = resolvePath(cwd);
+	const cacheToken = useCache ? useExtensionCacheCwd(cwd) : undefined;
+	const resolvedCwd = cacheToken?.cwd ?? resolvePath(cwd);
 	const resolvedEventBus = eventBus ?? createEventBus();
 	const resolvedRuntime = runtime ?? createExtensionRuntime();
 
 	for (const extPath of paths) {
-		const { extension, error } = await loadExtension(extPath, resolvedCwd, resolvedEventBus, resolvedRuntime);
+		const { extension, error } = await loadExtension(
+			extPath,
+			resolvedCwd,
+			resolvedEventBus,
+			resolvedRuntime,
+			cacheToken,
+		);
 
 		if (error) {
 			errors.push({ path: extPath, error });
@@ -442,6 +496,24 @@ export async function loadExtensions(
 	};
 }
 
+export async function loadExtensions(
+	paths: string[],
+	cwd: string,
+	eventBus?: EventBus,
+	runtime?: ExtensionRuntime,
+): Promise<LoadExtensionsResult> {
+	return loadExtensionsInternal(paths, cwd, eventBus, runtime);
+}
+
+export async function loadExtensionsCached(
+	paths: string[],
+	cwd: string,
+	eventBus?: EventBus,
+	runtime?: ExtensionRuntime,
+): Promise<LoadExtensionsResult> {
+	return loadExtensionsInternal(paths, cwd, eventBus, runtime, true);
+}
+
 interface PiManifest {
 	extensions?: string[];
 	themes?: string[];

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

@@ -9,7 +9,12 @@ export type { ResourceCollision, ResourceDiagnostic } from "./diagnostics.ts";
 
 import { canonicalizePath, isLocalPath, resolvePath } from "../utils/paths.ts";
 import { createEventBus, type EventBus } from "./event-bus.ts";
-import { createExtensionRuntime, loadExtensionFromFactory, loadExtensions } from "./extensions/loader.ts";
+import {
+	clearExtensionCache,
+	createExtensionRuntime,
+	loadExtensionFromFactory,
+	loadExtensionsCached,
+} from "./extensions/loader.ts";
 import type { Extension, ExtensionFactory, ExtensionRuntime, LoadExtensionsResult } from "./extensions/types.ts";
 import { DefaultPackageManager, type PathMetadata, type ResolvedResource } from "./package-manager.ts";
 import type { PromptTemplate } from "./prompt-templates.ts";
@@ -206,6 +211,7 @@ export class DefaultResourceLoader implements ResourceLoader {
 	private extensionThemeSourceInfos: Map<string, SourceInfo>;
 	private lastPromptPaths: string[];
 	private lastThemePaths: string[];
+	private loaded: boolean;
 
 	constructor(options: DefaultResourceLoaderOptions) {
 		this.cwd = resolvePath(options.cwd);
@@ -252,6 +258,7 @@ export class DefaultResourceLoader implements ResourceLoader {
 		this.extensionThemeSourceInfos = new Map();
 		this.lastPromptPaths = [];
 		this.lastThemePaths = [];
+		this.loaded = false;
 	}
 
 	getExtensions(): LoadExtensionsResult {
@@ -331,6 +338,10 @@ export class DefaultResourceLoader implements ResourceLoader {
 	}
 
 	async reload(options?: ResourceLoaderReloadOptions): Promise<void> {
+		if (this.loaded) {
+			clearExtensionCache();
+		}
+
 		let preTrustExtensions: LoadExtensionsResult | undefined;
 		if (options?.resolveProjectTrust) {
 			preTrustExtensions = await this.loadProjectTrustExtensions();
@@ -475,6 +486,7 @@ export class DefaultResourceLoader implements ResourceLoader {
 		this.appendSystemPrompt = this.appendSystemPromptOverride
 			? this.appendSystemPromptOverride(baseAppend)
 			: baseAppend;
+		this.loaded = true;
 	}
 
 	private async loadCurrentExtensionSet(options: { includeInlineFactories: boolean }): Promise<LoadExtensionsResult> {
@@ -487,7 +499,7 @@ export class DefaultResourceLoader implements ResourceLoader {
 		const extensionPaths = this.noExtensions
 			? cliEnabledExtensions
 			: this.mergePaths(cliEnabledExtensions, enabledExtensions);
-		const extensionsResult = await loadExtensions(extensionPaths, this.cwd, this.eventBus);
+		const extensionsResult = await loadExtensionsCached(extensionPaths, this.cwd, this.eventBus);
 		if (!options.includeInlineFactories) {
 			return extensionsResult;
 		}
@@ -507,7 +519,7 @@ export class DefaultResourceLoader implements ResourceLoader {
 		preTrustExtensions: LoadExtensionsResult | undefined,
 	): Promise<LoadExtensionsResult> {
 		if (!preTrustExtensions) {
-			const extensionsResult = await loadExtensions(extensionPaths, this.cwd, this.eventBus);
+			const extensionsResult = await loadExtensionsCached(extensionPaths, this.cwd, this.eventBus);
 			const inlineExtensions = await this.loadExtensionFactories(extensionsResult.runtime);
 			extensionsResult.extensions.push(...inlineExtensions.extensions);
 			extensionsResult.errors.push(...inlineExtensions.errors);
@@ -527,7 +539,7 @@ export class DefaultResourceLoader implements ResourceLoader {
 			const resolvedPath = this.resolveExtensionLoadPath(path);
 			return !preloadedByPath.has(resolvedPath) && !failedPreloadPaths.has(resolvedPath);
 		});
-		const remainingExtensions = await loadExtensions(
+		const remainingExtensions = await loadExtensionsCached(
 			remainingPaths,
 			this.cwd,
 			this.eventBus,

packages/coding-agent/test/suite/regressions/extension-factory-cache.test.ts

@@ -0,0 +1,131 @@
+import { existsSync, mkdirSync, rmSync, writeFileSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { afterEach, beforeEach, describe, expect, it } from "vitest";
+import { clearExtensionCache, loadExtensions, loadExtensionsCached } from "../../../src/core/extensions/loader.ts";
+import { DefaultResourceLoader } from "../../../src/core/resource-loader.ts";
+
+interface TestState {
+	moduleLoads?: number;
+	factoryRuns?: number;
+}
+
+function state(): TestState {
+	const global = globalThis as typeof globalThis & { __extensionFactoryCacheTest?: TestState };
+	if (!global.__extensionFactoryCacheTest) {
+		global.__extensionFactoryCacheTest = {};
+	}
+	return global.__extensionFactoryCacheTest;
+}
+
+function resetState(): void {
+	delete (globalThis as typeof globalThis & { __extensionFactoryCacheTest?: TestState }).__extensionFactoryCacheTest;
+}
+
+function writeCountingExtension(filePath: string): void {
+	writeFileSync(
+		filePath,
+		`
+const state = (globalThis.__extensionFactoryCacheTest ??= {});
+state.moduleLoads = (state.moduleLoads ?? 0) + 1;
+
+export default function () {
+	state.factoryRuns = (state.factoryRuns ?? 0) + 1;
+}
+`,
+		"utf-8",
+	);
+}
+
+describe("extension factory cache", () => {
+	const roots: string[] = [];
+
+	function fixture(name: string) {
+		const root = join(tmpdir(), `pi-extension-cache-${name}-${Date.now()}-${Math.random().toString(36).slice(2)}`);
+		const cwd = join(root, "project");
+		const agentDir = join(root, "agent");
+		mkdirSync(cwd, { recursive: true });
+		mkdirSync(agentDir, { recursive: true });
+		roots.push(root);
+		return { root, cwd, agentDir };
+	}
+
+	beforeEach(() => {
+		resetState();
+		clearExtensionCache();
+	});
+
+	afterEach(() => {
+		while (roots.length > 0) {
+			const root = roots.pop();
+			if (root && existsSync(root)) {
+				rmSync(root, { recursive: true, force: true });
+			}
+		}
+		resetState();
+		clearExtensionCache();
+	});
+
+	it("caches extension modules for cached same-cwd loads but reruns factories", async () => {
+		const { root, cwd } = fixture("same-cwd");
+		const extensionPath = join(root, "counting.ts");
+		writeCountingExtension(extensionPath);
+
+		const first = await loadExtensionsCached([extensionPath], cwd);
+		const second = await loadExtensionsCached([extensionPath], cwd);
+
+		expect(state().moduleLoads).toBe(1);
+		expect(state().factoryRuns).toBe(2);
+		expect(first.extensions[0]).not.toBe(second.extensions[0]);
+		expect(first.runtime).not.toBe(second.runtime);
+	});
+
+	it("does not cache direct loadExtensions calls", async () => {
+		const { root, cwd } = fixture("direct");
+		const extensionPath = join(root, "counting.ts");
+		writeCountingExtension(extensionPath);
+
+		await loadExtensions([extensionPath], cwd);
+		await loadExtensions([extensionPath], cwd);
+
+		expect(state().moduleLoads).toBe(2);
+		expect(state().factoryRuns).toBe(2);
+	});
+
+	it("clears the cache on resource loader reload", async () => {
+		const { cwd, agentDir } = fixture("reload");
+		const extensionDir = join(agentDir, "extensions");
+		mkdirSync(extensionDir, { recursive: true });
+		writeCountingExtension(join(extensionDir, "counting.ts"));
+		const loader = new DefaultResourceLoader({
+			cwd,
+			agentDir,
+			noSkills: true,
+			noPromptTemplates: true,
+			noThemes: true,
+		});
+
+		await loader.reload();
+		await loader.reload();
+
+		expect(state().moduleLoads).toBe(2);
+		expect(state().factoryRuns).toBe(2);
+	});
+
+	it("keeps the cache scoped to one cwd", async () => {
+		const { root } = fixture("cross-cwd");
+		const firstCwd = join(root, "first");
+		const secondCwd = join(root, "second");
+		mkdirSync(firstCwd, { recursive: true });
+		mkdirSync(secondCwd, { recursive: true });
+		const extensionPath = join(root, "counting.ts");
+		writeCountingExtension(extensionPath);
+
+		await loadExtensionsCached([extensionPath], firstCwd);
+		await loadExtensionsCached([extensionPath], secondCwd);
+		await loadExtensionsCached([extensionPath], secondCwd);
+
+		expect(state().moduleLoads).toBe(2);
+		expect(state().factoryRuns).toBe(3);
+	});
+});