feat(coding-agent): defer extension loader dependencies

Load ordinary jiti lazily on Node while retaining jiti/static for Bun and SEA binaries. This closes #9540.

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
 

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";

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";

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<string, unknown> = {
-	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<typeof createJiti> | undefined;
+
+function getCreateJiti(): Promise<typeof createJiti> {
+	createJitiPromise ??= (usesEmbeddedModules ? import("./jiti-static-loader.ts") : import("./jiti-loader.ts")).then(
+		(module) => module.createJiti,
+	);
+	return createJitiPromise;
+}
+
+let virtualModulesPromise: Promise<Record<string, unknown>> | undefined;
+
+function getVirtualModules(): Promise<Record<string, unknown>> {
+	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 });

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<string, unknown> = {
+	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,
+};

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();
 

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<string, unknown>;
+}
+
+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();
+	});
+});