diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 1e2a54d6..31f4b5f0 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/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)). diff --git a/packages/coding-agent/src/core/extensions/loader.ts b/packages/coding-agent/src/core/extensions/loader.ts index 081d2d11..dcaf389f 100644 --- a/packages/coding-agent/src/core/extensions/loader.ts +++ b/packages/coding-agent/src/core/extensions/loader.ts @@ -117,6 +117,30 @@ function getAliases(): Record { type HandlerFn = (...args: unknown[]) => Promise; +let extensionCacheCwd: string | undefined; +let extensionCacheGeneration = 0; +const extensionCache = new Map(); + +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 { 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 { + return loadExtensionsInternal(paths, cwd, eventBus, runtime); +} + +export async function loadExtensionsCached( + paths: string[], + cwd: string, + eventBus?: EventBus, + runtime?: ExtensionRuntime, +): Promise { + return loadExtensionsInternal(paths, cwd, eventBus, runtime, true); +} + interface PiManifest { extensions?: string[]; themes?: string[]; diff --git a/packages/coding-agent/src/core/resource-loader.ts b/packages/coding-agent/src/core/resource-loader.ts index b35787af..18486ead 100644 --- a/packages/coding-agent/src/core/resource-loader.ts +++ b/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; 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 { + 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 { @@ -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 { 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, diff --git a/packages/coding-agent/test/suite/regressions/extension-factory-cache.test.ts b/packages/coding-agent/test/suite/regressions/extension-factory-cache.test.ts new file mode 100644 index 00000000..ced01224 --- /dev/null +++ b/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); + }); +});