diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 68c223758..3bb3d5cef 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -35,6 +35,7 @@ - Moved compaction, branch summarization, and retry spinners into the editor border alongside the working indicator. Custom editors use the same embedding opt-in for all status spinners. - Enabled strict-prefer JSON-schema sampling by default for built-in `read`, `bash`, `powershell`, `edit`, and `write` tools, without requiring `PI_EXPERIMENTAL`. Extensions can re-register tool definitions with `constrainedSampling: false`. - Formatted Bash and PowerShell tool durations of at least one minute as minutes and seconds, with hours when needed ([#9628](https://github.com/earendil-works/pi/issues/9628)). +- Deferred the extension compiler and bundled virtual modules until a filesystem extension is loaded, reducing the baseline SDK import cost ([#9540](https://github.com/earendil-works/pi/issues/9540)). ### Fixed diff --git a/packages/coding-agent/src/core/extensions/jiti-loader.ts b/packages/coding-agent/src/core/extensions/jiti-loader.ts new file mode 100644 index 000000000..ae0c45972 --- /dev/null +++ b/packages/coding-agent/src/core/extensions/jiti-loader.ts @@ -0,0 +1,3 @@ +// Normal Node runtimes use jiti's lazy transform so Babel is loaded only when +// native loading fails and an extension needs transformation. +export { createJiti } from "jiti"; diff --git a/packages/coding-agent/src/core/extensions/jiti-static-loader.ts b/packages/coding-agent/src/core/extensions/jiti-static-loader.ts new file mode 100644 index 000000000..2d8987ba4 --- /dev/null +++ b/packages/coding-agent/src/core/extensions/jiti-static-loader.ts @@ -0,0 +1,3 @@ +// Compiled binaries need jiti's static entry so Bun and SEA bundlers embed the +// Babel transform. The module itself remains lazy until an extension is loaded. +export { createJiti } from "jiti/static"; diff --git a/packages/coding-agent/src/core/extensions/loader.ts b/packages/coding-agent/src/core/extensions/loader.ts index 56a773ab0..ab7ec239a 100644 --- a/packages/coding-agent/src/core/extensions/loader.ts +++ b/packages/coding-agent/src/core/extensions/loader.ts @@ -7,24 +7,10 @@ import * as fs from "node:fs"; import { createRequire } from "node:module"; import * as path from "node:path"; import { fileURLToPath } from "node:url"; -import * as _bundledPiAgentCore from "@earendil-works/pi-agent-core"; import type { Provider } from "@earendil-works/pi-ai"; -import * as _bundledPiAiCompat from "@earendil-works/pi-ai/compat"; -import * as _bundledPiAiOauth from "@earendil-works/pi-ai/oauth"; -import * as _bundledPiAiProviders from "@earendil-works/pi-ai/providers/all"; import type { KeyId } from "@earendil-works/pi-tui"; -import * as _bundledPiTui from "@earendil-works/pi-tui"; -import { createJiti } from "jiti/static"; -// Static imports of packages that extensions may use. -// These MUST be static so Bun bundles them into the compiled binary. -// The virtualModules option then makes them available to extensions. -import * as _bundledTypebox from "typebox"; -import * as _bundledTypeboxCompile from "typebox/compile"; -import * as _bundledTypeboxValue from "typebox/value"; +import type { createJiti } from "jiti"; import { CONFIG_DIR_NAME, getAgentDir, isBunBinary, isBundledNode } from "../../config.ts"; -// NOTE: This import works because loader.ts exports are NOT re-exported from index.ts, -// avoiding a circular dependency. Extensions can import from @earendil-works/pi-coding-agent. -import * as _bundledPiCodingAgent from "../../index.ts"; import { resolvePath } from "../../utils/paths.ts"; import { createEventBus, type EventBus } from "../event-bus.ts"; import type { ExecOptions } from "../exec.ts"; @@ -46,39 +32,29 @@ import type { ToolDefinition, } from "./types.ts"; -/** Modules available to extensions via virtualModules (for compiled binaries) */ -const VIRTUAL_MODULES: Record = { - typebox: _bundledTypebox, - "typebox/compile": _bundledTypeboxCompile, - "typebox/value": _bundledTypeboxValue, - "@sinclair/typebox": _bundledTypebox, - "@sinclair/typebox/compile": _bundledTypeboxCompile, - "@sinclair/typebox/value": _bundledTypeboxValue, - "@earendil-works/pi-agent-core": _bundledPiAgentCore, - "@earendil-works/pi-tui": _bundledPiTui, - // Extensions resolve the pi-ai root to the compat entrypoint (a strict - // superset of the core entrypoint): existing extensions using the old - // global API keep working at runtime until compat is removed. - "@earendil-works/pi-ai": _bundledPiAiCompat, - "@earendil-works/pi-ai/compat": _bundledPiAiCompat, - "@earendil-works/pi-ai/oauth": _bundledPiAiOauth, - "@earendil-works/pi-ai/providers/all": _bundledPiAiProviders, - "@earendil-works/pi-coding-agent": _bundledPiCodingAgent, - "@mariozechner/pi-agent-core": _bundledPiAgentCore, - "@mariozechner/pi-tui": _bundledPiTui, - "@mariozechner/pi-ai": _bundledPiAiCompat, - "@mariozechner/pi-ai/compat": _bundledPiAiCompat, - "@mariozechner/pi-ai/oauth": _bundledPiAiOauth, - "@mariozechner/pi-ai/providers/all": _bundledPiAiProviders, - "@mariozechner/pi-coding-agent": _bundledPiCodingAgent, -}; - const require = createRequire(import.meta.url); const isNodeSeaBinary = ("sea" in process.features && process.features.sea === true) || process.getBuiltinModule("node:sea")?.isSea() === true; const isTypeScriptSourceRuntime = !isBunBinary && path.extname(fileURLToPath(import.meta.url)) === ".ts"; +const usesEmbeddedModules = isBunBinary || isNodeSeaBinary || isBundledNode; + +let createJitiPromise: Promise | undefined; + +function getCreateJiti(): Promise { + createJitiPromise ??= (usesEmbeddedModules ? import("./jiti-static-loader.ts") : import("./jiti-loader.ts")).then( + (module) => module.createJiti, + ); + return createJitiPromise; +} + +let virtualModulesPromise: Promise> | undefined; + +function getVirtualModules(): Promise> { + virtualModulesPromise ??= import("./virtual-modules.ts").then((module) => module.VIRTUAL_MODULES); + return virtualModulesPromise; +} /** * Get aliases for jiti (used in built Node.js mode). @@ -508,16 +484,18 @@ async function loadExtensionModule(extensionPath: string, cacheToken?: Extension } } - const jiti = createJiti(import.meta.url, { + const createJitiImpl = await getCreateJiti(); + // Compiled binaries and the bundled Node distribution use embedded modules. + // Source TypeScript reuses host modules and root tsconfig paths. Unbundled + // Node builds use dist aliases and do not need the bundled virtual modules. + const resolutionOptions = usesEmbeddedModules + ? { virtualModules: await getVirtualModules(), tryNative: false } + : isTypeScriptSourceRuntime + ? { virtualModules: await getVirtualModules(), tsconfigPaths: true } + : { alias: getAliases() }; + const jiti = createJitiImpl(import.meta.url, { moduleCache: false, - // Compiled binaries and the bundled Node distribution use embedded modules. - // Source TypeScript reuses host modules and root tsconfig paths. Unbundled - // Node builds use dist aliases. - ...(isBunBinary || isNodeSeaBinary || isBundledNode - ? { virtualModules: VIRTUAL_MODULES, tryNative: false } - : isTypeScriptSourceRuntime - ? { virtualModules: VIRTUAL_MODULES, tsconfigPaths: true } - : { alias: getAliases() }), + ...resolutionOptions, }); const module = await jiti.import(extensionPath, { default: true }); diff --git a/packages/coding-agent/src/core/extensions/virtual-modules.ts b/packages/coding-agent/src/core/extensions/virtual-modules.ts new file mode 100644 index 000000000..66a71f695 --- /dev/null +++ b/packages/coding-agent/src/core/extensions/virtual-modules.ts @@ -0,0 +1,38 @@ +import * as bundledPiAgentCore from "@earendil-works/pi-agent-core"; +import * as bundledPiAiCompat from "@earendil-works/pi-ai/compat"; +import * as bundledPiAiOauth from "@earendil-works/pi-ai/oauth"; +import * as bundledPiAiProviders from "@earendil-works/pi-ai/providers/all"; +import * as bundledPiTui from "@earendil-works/pi-tui"; +import * as bundledTypebox from "typebox"; +import * as bundledTypeboxCompile from "typebox/compile"; +import * as bundledTypeboxValue from "typebox/value"; +// This import is safe because loader.ts exports are not re-exported from index.ts. +// Extensions can therefore import from @earendil-works/pi-coding-agent. +import * as bundledPiCodingAgent from "../../index.ts"; + +/** Modules available to extensions in source and compiled binary runtimes. */ +export const VIRTUAL_MODULES: Record = { + typebox: bundledTypebox, + "typebox/compile": bundledTypeboxCompile, + "typebox/value": bundledTypeboxValue, + "@sinclair/typebox": bundledTypebox, + "@sinclair/typebox/compile": bundledTypeboxCompile, + "@sinclair/typebox/value": bundledTypeboxValue, + "@earendil-works/pi-agent-core": bundledPiAgentCore, + "@earendil-works/pi-tui": bundledPiTui, + // Extensions resolve the pi-ai root to the compat entrypoint (a strict + // superset of the core entrypoint): existing extensions using the old + // global API keep working at runtime until compat is removed. + "@earendil-works/pi-ai": bundledPiAiCompat, + "@earendil-works/pi-ai/compat": bundledPiAiCompat, + "@earendil-works/pi-ai/oauth": bundledPiAiOauth, + "@earendil-works/pi-ai/providers/all": bundledPiAiProviders, + "@earendil-works/pi-coding-agent": bundledPiCodingAgent, + "@mariozechner/pi-agent-core": bundledPiAgentCore, + "@mariozechner/pi-tui": bundledPiTui, + "@mariozechner/pi-ai": bundledPiAiCompat, + "@mariozechner/pi-ai/compat": bundledPiAiCompat, + "@mariozechner/pi-ai/oauth": bundledPiAiOauth, + "@mariozechner/pi-ai/providers/all": bundledPiAiProviders, + "@mariozechner/pi-coding-agent": bundledPiCodingAgent, +}; diff --git a/packages/coding-agent/test/suite/regressions/8237-node-sea-extension-loading.test.ts b/packages/coding-agent/test/suite/regressions/8237-node-sea-extension-loading.test.ts index 160bf7dc3..5ba41fdf0 100644 --- a/packages/coding-agent/test/suite/regressions/8237-node-sea-extension-loading.test.ts +++ b/packages/coding-agent/test/suite/regressions/8237-node-sea-extension-loading.test.ts @@ -9,13 +9,28 @@ const state = vi.hoisted(() => { }); return { originalGetBuiltinModule, + jitiModuleLoads: 0, + virtualModulesLoads: 0, createJiti: vi.fn((_id: unknown, _options: unknown) => ({ import: vi.fn(async () => () => {}), })), }; }); -vi.mock("jiti/static", () => ({ createJiti: state.createJiti })); +vi.mock("jiti/static", () => { + state.jitiModuleLoads++; + return { createJiti: state.createJiti }; +}); + +vi.mock("../../../src/core/extensions/virtual-modules.ts", () => { + state.virtualModulesLoads++; + return { + VIRTUAL_MODULES: { + typebox: {}, + "@earendil-works/pi-coding-agent": {}, + }, + }; +}); import { loadExtensions } from "../../../src/core/extensions/loader.ts"; @@ -32,10 +47,16 @@ afterAll(() => { }); describe("Node SEA extension loading", () => { - it("uses bundled virtual modules instead of filesystem aliases", async () => { + // Regression test for #8237 and #9540. + it("loads jiti and bundled virtual modules only when importing an extension", async () => { + expect(state.jitiModuleLoads).toBe(0); + expect(state.virtualModulesLoads).toBe(0); + const result = await loadExtensions(["/extension.ts"], "/"); expect(result.errors).toEqual([]); + expect(state.jitiModuleLoads).toBe(1); + expect(state.virtualModulesLoads).toBe(1); expect(result.extensions).toHaveLength(1); expect(state.createJiti).toHaveBeenCalledOnce(); diff --git a/packages/coding-agent/test/suite/regressions/9540-extension-loader-lazy.test.ts b/packages/coding-agent/test/suite/regressions/9540-extension-loader-lazy.test.ts new file mode 100644 index 000000000..97c006830 --- /dev/null +++ b/packages/coding-agent/test/suite/regressions/9540-extension-loader-lazy.test.ts @@ -0,0 +1,58 @@ +import { describe, expect, it, vi } from "vitest"; + +const state = vi.hoisted(() => ({ + jitiModuleLoads: 0, + jitiStaticModuleLoads: 0, + virtualModulesLoads: 0, + createJiti: vi.fn((_id: unknown, _options: unknown) => ({ + import: vi.fn(async () => () => {}), + })), +})); + +vi.mock("jiti", () => { + state.jitiModuleLoads++; + return { createJiti: state.createJiti }; +}); + +vi.mock("jiti/static", () => { + state.jitiStaticModuleLoads++; + return { createJiti: state.createJiti }; +}); + +vi.mock("../../../src/core/extensions/virtual-modules.ts", () => { + state.virtualModulesLoads++; + return { VIRTUAL_MODULES: {} }; +}); + +import { loadExtensions } from "../../../src/core/extensions/loader.ts"; + +interface JitiOptionsProbe { + alias?: unknown; + tryNative?: boolean; + tsconfigPaths?: boolean; + virtualModules?: Record; +} + +describe("extension loader lazy imports", () => { + // Regression test for #9540. + it("defers ordinary jiti and its virtual modules until importing an extension", async () => { + expect(state.jitiModuleLoads).toBe(0); + expect(state.jitiStaticModuleLoads).toBe(0); + expect(state.virtualModulesLoads).toBe(0); + + const result = await loadExtensions(["/extension.ts"], "/"); + + expect(result.errors).toEqual([]); + expect(result.extensions).toHaveLength(1); + expect(state.jitiModuleLoads).toBe(1); + expect(state.jitiStaticModuleLoads).toBe(0); + expect(state.virtualModulesLoads).toBe(1); + expect(state.createJiti).toHaveBeenCalledOnce(); + + const options = state.createJiti.mock.calls[0][1] as JitiOptionsProbe; + expect(options.tryNative).toBeUndefined(); + expect(options.tsconfigPaths).toBe(true); + expect(options.alias).toBeUndefined(); + expect(options.virtualModules).toBeDefined(); + }); +});