UNPKG

@earendil-works/pi-coding-agent

Version:

Coding agent CLI with read, bash, edit, write tools and session management

710 lines • 29.4 kB
/** * Extension loader - loads TypeScript extension modules using jiti. * */ import * as fs from "node:fs"; import { createRequire } from "node:module"; import * as path from "node:path"; import { fileURLToPath } from "node:url"; import { mcpNamespace } from "@earendil-works/pi-mcp"; import { CONFIG_DIR_NAME, getAgentDir, isBunBinary, isBundledNode } from "../../config.js"; import { canonicalizePath, pathsEqual, resolvePath } from "../../utils/paths.js"; import { createEventBus } from "../event-bus.js"; import { execCommand } from "../exec.js"; import { McpServerRegistry, validateMcpServerConfig } from "../mcp-servers.js"; import { readPiManifest } from "../pi-manifest.js"; import { createSyntheticSourceInfo, getSyntheticPathSource, isSyntheticPath } from "../source-info.js"; import { time } from "../timings.js"; 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; function getCreateJiti() { createJitiPromise ??= (usesEmbeddedModules ? import("./jiti-static-loader.js") : import("./jiti-loader.js")).then((module) => module.createJiti); return createJitiPromise; } let virtualModulesPromise; function getVirtualModules() { virtualModulesPromise ??= import("./virtual-modules.js").then((module) => module.VIRTUAL_MODULES); return virtualModulesPromise; } /** * Get aliases for jiti (used in built Node.js mode). * In compiled binary mode, virtualModules is used instead. */ let _aliases = null; function getAliases() { if (_aliases) return _aliases; const __dirname = path.dirname(fileURLToPath(import.meta.url)); const packageIndex = path.resolve(__dirname, "../..", "index.js"); const typeboxEntry = require.resolve("typebox"); const typeboxCompileEntry = require.resolve("typebox/compile"); const typeboxValueEntry = require.resolve("typebox/value"); const packagesRoot = path.resolve(__dirname, "../../../../"); const resolveWorkspaceOrImport = (workspaceRelativePath, specifier) => { const workspacePath = path.join(packagesRoot, workspaceRelativePath); if (fs.existsSync(workspacePath)) { return workspacePath; } return fileURLToPath(import.meta.resolve(specifier)); }; const piCodingAgentEntry = packageIndex; const piAgentCoreEntry = resolveWorkspaceOrImport("agent/dist/index.js", "@earendil-works/pi-agent-core"); const piTuiEntry = resolveWorkspaceOrImport("tui/dist/index.js", "@earendil-works/pi-tui"); // Extensions resolve the pi-ai root to the extension view: the compat // entrypoint (a strict superset of the core entrypoint) plus TypeBox's Type, // Static and TSchema, so extensions using the old global API and pi-ai's // former TypeBox re-exports keep working. const piAiCompatEntry = path.resolve(__dirname, "pi-ai-extension-view.js"); const piAiModelsEntry = resolveWorkspaceOrImport("ai/dist/models.js", "@earendil-works/pi-ai/models"); const piAiOauthEntry = resolveWorkspaceOrImport("ai/dist/oauth.js", "@earendil-works/pi-ai/oauth"); const piAiProvidersEntry = resolveWorkspaceOrImport("ai/dist/providers/all.js", "@earendil-works/pi-ai/providers/all"); _aliases = { "@earendil-works/pi-coding-agent": piCodingAgentEntry, "@earendil-works/pi-agent-core": piAgentCoreEntry, "@earendil-works/pi-tui": piTuiEntry, "@earendil-works/pi-ai/providers/all": piAiProvidersEntry, "@earendil-works/pi-ai/compat": piAiCompatEntry, "@earendil-works/pi-ai/models": piAiModelsEntry, "@earendil-works/pi-ai/oauth": piAiOauthEntry, "@earendil-works/pi-ai": piAiCompatEntry, "@mariozechner/pi-coding-agent": piCodingAgentEntry, "@mariozechner/pi-agent-core": piAgentCoreEntry, "@mariozechner/pi-tui": piTuiEntry, "@mariozechner/pi-ai/providers/all": piAiProvidersEntry, "@mariozechner/pi-ai/compat": piAiCompatEntry, "@mariozechner/pi-ai/models": piAiModelsEntry, "@mariozechner/pi-ai/oauth": piAiOauthEntry, "@mariozechner/pi-ai": piAiCompatEntry, typebox: typeboxEntry, "typebox/compile": typeboxCompileEntry, "typebox/value": typeboxValueEntry, "@sinclair/typebox": typeboxEntry, "@sinclair/typebox/compile": typeboxCompileEntry, "@sinclair/typebox/value": typeboxValueEntry, }; return _aliases; } let extensionCacheCwd; let extensionCacheGeneration = 0; const extensionCache = new Map(); export function clearExtensionCache() { extensionCache.clear(); extensionCacheCwd = undefined; extensionCacheGeneration++; } function useExtensionCacheCwd(cwd) { const resolvedCwd = resolvePath(cwd); if (extensionCacheCwd !== undefined && !pathsEqual(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. */ export function createExtensionRuntime() { const notInitialized = () => { throw new Error("Extension runtime not initialized. Action methods cannot be called during extension loading."); }; const state = {}; const eventBusUnsubscribers = new Set(); const assertActive = () => { if (state.staleMessage) { throw new Error(state.staleMessage); } }; const runtime = { sendMessage: notInitialized, sendUserMessage: notInitialized, appendEntry: notInitialized, setSessionName: notInitialized, getSessionName: notInitialized, setLabel: notInitialized, getActiveTools: notInitialized, getAllTools: notInitialized, getSettings: notInitialized, setActiveTools: notInitialized, // registerTool() is valid during extension load; refresh is only needed post-bind. refreshTools: () => { }, getCommands: notInitialized, setModel: () => Promise.reject(new Error("Extension runtime not initialized")), getThinkingLevel: notInitialized, setThinkingLevel: notInitialized, flagValues: new Map(), pendingProviderRegistrations: [], pendingNativeProviderRegistrations: [], mcpServers: new McpServerRegistry(), pendingVirtualModelRegistrations: [], createContext: notInitialized, assertActive, invalidate: (message) => { 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. registerProvider: (name, config, extensionPath = "<unknown>") => { runtime.pendingProviderRegistrations.push({ name, config, extensionPath }); }, registerNativeProvider: (provider, extensionPath = "<unknown>") => { runtime.pendingNativeProviderRegistrations.push({ provider, extensionPath }); }, unregisterProvider: (name) => { runtime.pendingProviderRegistrations = runtime.pendingProviderRegistrations.filter((r) => r.name !== name); runtime.pendingNativeProviderRegistrations = runtime.pendingNativeProviderRegistrations.filter((r) => r.provider.id !== name); }, registerVirtualModel: (definition, extensionPath = "<unknown>") => { runtime.pendingVirtualModelRegistrations.push({ definition, extensionPath }); }, unregisterVirtualModel: (provider, id) => { runtime.pendingVirtualModelRegistrations = runtime.pendingVirtualModelRegistrations.filter(({ definition }) => definition.provider !== provider || definition.id !== id); }, }; return runtime; } /** * Create the ExtensionAPI for an extension. * Registration methods write to the extension object. * Action methods delegate to the shared runtime. */ const DEFAULT_MARKDOWN_MESSAGE_TYPES = ["user", "assistant", "assistant-thinking"]; function createExtensionAPI(extension, runtime, cwd, eventBus) { const pendingFlagValues = new Map(); const pendingRuntimeChanges = []; const loadingUnsubscribers = []; let state = "loading"; const assertActive = () => { if (state === "failed") { throw new Error(`Extension "${extension.path}" failed to load and its API is no longer active.`); } runtime.assertActive(); }; const applyRuntimeChange = (change) => { if (state === "loading") pendingRuntimeChanges.push(change); else change(); }; const clearPending = () => { pendingFlagValues.clear(); pendingRuntimeChanges.length = 0; loadingUnsubscribers.length = 0; }; const api = { // Registration methods - write to extension on(event, handler) { assertActive(); const registeredHandler = (...args) => handler(...args); const list = extension.handlers.get(event) ?? []; list.push(registeredHandler); extension.handlers.set(event, list); return () => { const handlers = extension.handlers.get(event); if (!handlers) return; const handlerIndex = handlers.indexOf(registeredHandler); if (handlerIndex === -1) return; handlers.splice(handlerIndex, 1); if (handlers.length === 0) extension.handlers.delete(event); }; }, registerTool(tool) { assertActive(); if (typeof tool.parameters !== "object" || tool.parameters === null || Array.isArray(tool.parameters)) { throw new Error(`Tool "${tool.name}" registered by extension "${extension.path}" must define an object parameter schema.`); } extension.tools.set(tool.name, { definition: tool, sourceInfo: extension.sourceInfo, }); runtime.refreshTools(); }, registerCommand(name, options) { assertActive(); if (typeof name !== "string" || name.length === 0) { throw new Error(`Command registered by extension "${extension.path}" must have a non-empty string name. Use pi.registerCommand("name", { description, handler }).`); } if (typeof options?.handler !== "function") { throw new Error(`Command "/${name}" registered by extension "${extension.path}" must define handler().`); } extension.commands.set(name, { name, sourceInfo: extension.sourceInfo, ...options, }); }, registerShortcut(shortcut, options) { assertActive(); extension.shortcuts.set(shortcut, { shortcut, extensionPath: extension.path, ...options }); }, registerFlag(name, options) { assertActive(); if (options.default !== undefined && typeof options.default !== options.type) { throw new Error(`Invalid default for flag "${name}": expected ${options.type}, got ${typeof options.default}`); } extension.flags.set(name, { name, extensionPath: extension.path, ...options }); if (options.default !== undefined && !runtime.flagValues.has(name)) { if (state === "loading") { if (!pendingFlagValues.has(name)) pendingFlagValues.set(name, options.default); } else { runtime.flagValues.set(name, options.default); } } }, registerMessageRenderer(customType, renderer) { assertActive(); extension.messageRenderers.set(customType, renderer); }, registerMarkdownTransformer(transformer, options) { assertActive(); extension.markdownTransformer = { transformer, messageTypes: new Set(options?.messageTypes ?? DEFAULT_MARKDOWN_MESSAGE_TYPES), }; }, registerEntryRenderer(customType, renderer) { assertActive(); extension.entryRenderers ??= new Map(); extension.entryRenderers.set(customType, renderer); }, registerToolRenderer(resolver) { assertActive(); extension.toolRenderers ??= []; extension.toolRenderers.push(resolver); }, // Flag access - checks extension registered it, reads from runtime getFlag(name) { assertActive(); if (!extension.flags.has(name)) return undefined; return runtime.flagValues.has(name) ? runtime.flagValues.get(name) : pendingFlagValues.get(name); }, // Action methods - delegate to shared runtime sendMessage(message, options) { assertActive(); runtime.sendMessage(message, options); }, sendUserMessage(content, options) { assertActive(); runtime.sendUserMessage(content, options); }, appendEntry(customType, data) { assertActive(); runtime.appendEntry(customType, data); }, setSessionName(name) { assertActive(); runtime.setSessionName(name); }, getSessionName() { assertActive(); return runtime.getSessionName(); }, setLabel(entryId, label) { assertActive(); runtime.setLabel(entryId, label); }, exec(command, args, options) { assertActive(); return execCommand(command, args, options?.cwd ?? cwd, options); }, getActiveTools() { assertActive(); return runtime.getActiveTools(); }, getAllTools() { assertActive(); return runtime.getAllTools(); }, getSettings() { assertActive(); return runtime.getSettings(); }, setActiveTools(toolNames) { assertActive(); runtime.setActiveTools(toolNames); }, getCommands() { assertActive(); return runtime.getCommands(); }, setModel(model) { assertActive(); return runtime.setModel(model); }, getThinkingLevel() { assertActive(); return runtime.getThinkingLevel(); }, setThinkingLevel(level) { assertActive(); runtime.setThinkingLevel(level); }, registerProvider(providerOrName, config) { assertActive(); if (typeof providerOrName === "string") { if (!config) throw new Error("Provider config is required when registering by name"); applyRuntimeChange(() => runtime.registerProvider(providerOrName, config, extension.path)); return; } applyRuntimeChange(() => runtime.registerNativeProvider(providerOrName, extension.path)); }, unregisterProvider(name) { assertActive(); applyRuntimeChange(() => runtime.unregisterProvider(name, extension.path)); }, registerMcpServer(name, config) { assertActive(); const validated = validateMcpServerConfig(name, config); if (typeof validated === "string") { throw new Error(`Invalid MCP server registered by extension "${extension.path}": ${validated}`); } const owner = runtime.mcpServers.get(name)?.extensionPath; if (owner !== undefined && owner !== extension.path) { throw new Error(`MCP server "${name}" is already registered by extension "${owner}"`); } // Names that differ only in `-` and `_` would share a namespace. const clash = runtime.mcpServers .list() .find((server) => server.name !== name && mcpNamespace(server.name) === mcpNamespace(name)); if (clash) throw new Error(`MCP server "${name}" conflicts with registered server "${clash.name}"`); const server = { name, config: structuredClone(validated), extensionPath: extension.path }; applyRuntimeChange(() => runtime.mcpServers.register(server)); }, unregisterMcpServer(name) { assertActive(); applyRuntimeChange(() => runtime.mcpServers.unregister(name, extension.path)); }, getMcpServers() { assertActive(); return runtime.mcpServers.list(); }, registerVirtualModel(model) { assertActive(); // Routing runs after the runner binds, so the context is created per request. The state // comes from the session branch that this router wrote. const definition = { ...model, route: (request) => model.route(request, runtime.createContext()), }; applyRuntimeChange(() => runtime.registerVirtualModel(definition, extension.path)); }, unregisterVirtualModel(provider, id) { assertActive(); applyRuntimeChange(() => runtime.unregisterVirtualModel(provider, id)); }, events: { emit(channel, data) { assertActive(); eventBus.emit(channel, data); }, on(channel, handler) { assertActive(); const unsubscribe = runtime.trackEventBusSubscription(eventBus.on(channel, handler)); if (state === "loading") loadingUnsubscribers.push(unsubscribe); return unsubscribe; }, }, }; return { api, commit: () => { if (state !== "loading") return; runtime.assertActive(); for (const [name, value] of pendingFlagValues) { if (!runtime.flagValues.has(name)) runtime.flagValues.set(name, value); } for (const apply of pendingRuntimeChanges) apply(); state = "active"; clearPending(); }, discard: () => { if (state !== "loading") return; state = "failed"; for (const unsubscribe of loadingUnsubscribers) unsubscribe(); clearPending(); }, }; } function isCurrentCacheToken(cacheToken) { return (cacheToken !== undefined && extensionCacheCwd === cacheToken.cwd && extensionCacheGeneration === cacheToken.generation); } async function loadExtensionModule(extensionPath, cacheToken) { if (isCurrentCacheToken(cacheToken)) { const cachedFactory = extensionCache.get(extensionPath); if (cachedFactory) { return cachedFactory; } } 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, ...resolutionOptions, }); const module = await jiti.import(canonicalizePath(extensionPath), { default: true }); const factory = module; if (typeof factory !== "function") { return undefined; } if (isCurrentCacheToken(cacheToken)) { extensionCache.set(extensionPath, factory); } return factory; } /** * Create an Extension object with empty collections. */ function createExtension(extensionPath, resolvedPath) { const source = getSyntheticPathSource(extensionPath) ?? "local"; const baseDir = isSyntheticPath(extensionPath) ? undefined : path.dirname(resolvedPath); return { path: extensionPath, resolvedPath, sourceInfo: createSyntheticSourceInfo(extensionPath, { source, baseDir }), handlers: new Map(), tools: new Map(), messageRenderers: new Map(), entryRenderers: new Map(), commands: new Map(), flags: new Map(), shortcuts: new Map(), }; } async function initializeExtension(factory, extensionPath, resolvedPath, cwd, eventBus, runtime) { const extension = createExtension(extensionPath, resolvedPath); const load = createExtensionAPI(extension, runtime, cwd, eventBus); try { await factory(load.api); load.commit(); } catch (error) { load.discard(); throw error; } time(`${extensionPath} factory`, "extensions"); return extension; } async function loadExtension(extensionPath, cwd, eventBus, runtime, cacheToken) { const resolvedPath = resolvePath(extensionPath, cwd, { normalizeUnicodeSpaces: true }); try { const factory = await loadExtensionModule(resolvedPath, cacheToken); time(`${extensionPath} module import`, "extensions"); if (!factory) { return { extension: null, error: `Extension does not export a valid factory function: ${extensionPath}` }; } const extension = await initializeExtension(factory, extensionPath, resolvedPath, cwd, eventBus, runtime); return { extension, error: null }; } catch (err) { const message = err instanceof Error ? err.message : String(err); return { extension: null, error: `Failed to load extension: ${message}` }; } } /** * Create an Extension from an inline factory function. */ export async function loadExtensionFromFactory(factory, cwd, eventBus, runtime, extensionPath = "<inline>") { const resolvedCwd = resolvePath(cwd); return initializeExtension(factory, extensionPath, extensionPath, resolvedCwd, eventBus, runtime); } /** * Load extensions from paths. */ async function loadExtensionsInternal(paths, cwd, eventBus, runtime, useCache = false) { const extensions = []; const errors = []; const warnings = []; 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, cacheToken); if (error) { errors.push({ path: extPath, error }); continue; } if (extension) { extensions.push(extension); } } return { extensions, errors, warnings, runtime: resolvedRuntime, }; } export async function loadExtensions(paths, cwd, eventBus, runtime) { return loadExtensionsInternal(paths, cwd, eventBus, runtime); } export async function loadExtensionsCached(paths, cwd, eventBus, runtime) { return loadExtensionsInternal(paths, cwd, eventBus, runtime, true); } function isExtensionFile(name) { return name.endsWith(".ts") || name.endsWith(".js"); } /** * Resolve extension entry points from a directory. * * Checks for: * 1. package.json with "pi.extensions" field -> returns declared paths * 2. index.ts or index.js -> returns the index file * * Returns resolved paths or null if no entry points found. */ function resolveExtensionEntries(dir) { // Check for package.json with "pi" field first const packageJsonPath = path.join(dir, "package.json"); if (fs.existsSync(packageJsonPath)) { const manifest = readPiManifest(packageJsonPath); if (manifest?.extensions?.length) { const entries = []; for (const extPath of manifest.extensions) { const resolvedExtPath = path.resolve(dir, extPath); if (fs.existsSync(resolvedExtPath)) { entries.push(resolvedExtPath); } } if (entries.length > 0) { return entries; } } } // Check for index.ts or index.js const indexTs = path.join(dir, "index.ts"); const indexJs = path.join(dir, "index.js"); if (fs.existsSync(indexTs)) { return [indexTs]; } if (fs.existsSync(indexJs)) { return [indexJs]; } return null; } /** * Discover extensions in a directory. * * Discovery rules: * 1. Direct files: `extensions/*.ts` or `*.js` → load * 2. Subdirectory with index: `extensions/* /index.ts` or `index.js` → load * 3. Subdirectory with package.json: `extensions/* /package.json` with "pi" field → load what it declares * * No recursion beyond one level. Complex packages must use package.json manifest. */ function discoverExtensionsInDir(dir) { if (!fs.existsSync(dir)) { return []; } const discovered = []; try { const entries = fs.readdirSync(dir, { withFileTypes: true }); for (const entry of entries) { const entryPath = path.join(dir, entry.name); // 1. Direct files: *.ts or *.js if ((entry.isFile() || entry.isSymbolicLink()) && isExtensionFile(entry.name)) { discovered.push(entryPath); continue; } // 2 & 3. Subdirectories if (entry.isDirectory() || entry.isSymbolicLink()) { const entries = resolveExtensionEntries(entryPath); if (entries) { discovered.push(...entries); } } } } catch { return []; } return discovered; } /** * Discover and load extensions from standard locations. */ export async function discoverAndLoadExtensions(configuredPaths, cwd, agentDir = getAgentDir(), eventBus) { const resolvedCwd = resolvePath(cwd); const resolvedAgentDir = resolvePath(agentDir); const allPaths = []; const seen = new Set(); const addPaths = (paths) => { for (const p of paths) { const resolved = path.resolve(p); if (!seen.has(resolved)) { seen.add(resolved); allPaths.push(p); } } }; // 1. Project-local extensions: cwd/${CONFIG_DIR_NAME}/extensions/ const localExtDir = path.join(resolvedCwd, CONFIG_DIR_NAME, "extensions"); addPaths(discoverExtensionsInDir(localExtDir)); // 2. Global extensions: agentDir/extensions/ const globalExtDir = path.join(resolvedAgentDir, "extensions"); addPaths(discoverExtensionsInDir(globalExtDir)); // 3. Explicitly configured paths for (const p of configuredPaths) { const resolved = resolvePath(p, resolvedCwd, { normalizeUnicodeSpaces: true }); if (fs.existsSync(resolved) && fs.statSync(resolved).isDirectory()) { // Check for package.json with pi manifest or index.ts const entries = resolveExtensionEntries(resolved); if (entries) { addPaths(entries); continue; } // No explicit entries - discover individual files in directory addPaths(discoverExtensionsInDir(resolved)); continue; } addPaths([resolved]); } return loadExtensions(allPaths, resolvedCwd, eventBus); } //# sourceMappingURL=loader.js.map