Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
18 kB · 603 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604/** * Extension loader - loads TypeScript extension modules using jiti. * */
import * as fs from "node:fs";import { createRequire } from "node:module";import * as os from "node:os";import * as path from "node:path";import { fileURLToPath } from "node:url";import type { KeyId } from "@earendil-works/pi-tui";import type { Jiti } from "jiti";import { CONFIG_DIR_NAME, getAgentDir, isBunBinary } from "../../config.js";import { createEventBus, type EventBus } from "../event-bus.js";import type { ExecOptions } from "../exec.js";import { execCommand } from "../exec.js";import { createSyntheticSourceInfo } from "../source-info.js";import { time } from "../timings.js";import type { Extension, ExtensionAPI, ExtensionFactory, ExtensionRuntime, LoadExtensionsResult, MessageRenderer, ProviderConfig, RegisteredCommand, ToolDefinition,} from "./types.js";
const require = createRequire(import.meta.url);
/** * Get aliases for jiti (used in Node.js/development mode). * In Bun binary mode, virtualModules is used instead. */let _aliases: Record<string, string> | null = null;
function getAliases(): Record<string, string> { 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: string, specifier: string): string => { 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"); const piAiEntry = resolveWorkspaceOrImport("ai/dist/index.js", "@earendil-works/pi-ai"); const piAiOauthEntry = resolveWorkspaceOrImport("ai/dist/oauth.js", "@earendil-works/pi-ai/oauth");
_aliases = { "@earendil-works/pi-coding-agent": piCodingAgentEntry, "@earendil-works/pi-agent-core": piAgentCoreEntry, "@earendil-works/pi-tui": piTuiEntry, "@earendil-works/pi-ai": piAiEntry, "@earendil-works/pi-ai/oauth": piAiOauthEntry, "@mariozechner/pi-coding-agent": piCodingAgentEntry, "@mariozechner/pi-agent-core": piAgentCoreEntry, "@mariozechner/pi-tui": piTuiEntry, "@mariozechner/pi-ai": piAiEntry, "@mariozechner/pi-ai/oauth": piAiOauthEntry, typebox: typeboxEntry, "typebox/compile": typeboxCompileEntry, "typebox/value": typeboxValueEntry, "@sinclair/typebox": typeboxEntry, "@sinclair/typebox/compile": typeboxCompileEntry, "@sinclair/typebox/value": typeboxValueEntry, };
return _aliases;}
const UNICODE_SPACES = /[\u00A0\u2000-\u200A\u202F\u205F\u3000]/g;
function normalizeUnicodeSpaces(str: string): string { return str.replace(UNICODE_SPACES, " ");}
function expandPath(p: string): string { const normalized = normalizeUnicodeSpaces(p); if (normalized.startsWith("~/")) { return path.join(os.homedir(), normalized.slice(2)); } if (normalized.startsWith("~")) { return path.join(os.homedir(), normalized.slice(1)); } return normalized;}
function resolvePath(extPath: string, cwd: string): string { const expanded = expandPath(extPath); if (path.isAbsolute(expanded)) { return expanded; } return path.resolve(cwd, expanded);}
type HandlerFn = (...args: unknown[]) => Promise<unknown>;
/** * Create a runtime with throwing stubs for action methods. * Runner.bindCore() replaces these with real implementations. */export function createExtensionRuntime(): ExtensionRuntime { const notInitialized = () => { throw new Error("Extension runtime not initialized. Action methods cannot be called during extension loading."); }; const state: { staleMessage?: string } = {}; const assertActive = () => { if (state.staleMessage) { throw new Error(state.staleMessage); } };
const runtime: ExtensionRuntime = { sendMessage: notInitialized, sendUserMessage: notInitialized, appendEntry: notInitialized, setSessionName: notInitialized, getSessionName: notInitialized, setLabel: notInitialized, getActiveTools: notInitialized, getAllTools: 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: [], assertActive, invalidate: (message) => { 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()."; }, // 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 }); }, unregisterProvider: (name) => { runtime.pendingProviderRegistrations = runtime.pendingProviderRegistrations.filter((r) => r.name !== name); }, };
return runtime;}
/** * Create the ExtensionAPI for an extension. * Registration methods write to the extension object. * Action methods delegate to the shared runtime. */function createExtensionAPI( extension: Extension, runtime: ExtensionRuntime, cwd: string, eventBus: EventBus,): ExtensionAPI { const api = { on(event: string, handler: HandlerFn): void { runtime.assertActive(); const list = extension.handlers.get(event) ?? []; list.push(handler); extension.handlers.set(event, list); },
registerTool(tool: ToolDefinition): void { runtime.assertActive(); extension.tools.set(tool.name, { definition: tool, sourceInfo: extension.sourceInfo, }); runtime.refreshTools(); },
registerCommand(name: string, options: Omit<RegisteredCommand, "name" | "sourceInfo">): void { runtime.assertActive(); extension.commands.set(name, { name, sourceInfo: extension.sourceInfo, ...options, }); },
registerShortcut( shortcut: KeyId, options: { description?: string; handler: (ctx: import("./types.js").ExtensionContext) => Promise<void> | void; }, ): void { runtime.assertActive(); extension.shortcuts.set(shortcut, { shortcut, extensionPath: extension.path, ...options }); },
registerFlag( name: string, options: { description?: string; type: "boolean" | "string"; default?: boolean | string }, ): void { runtime.assertActive(); extension.flags.set(name, { name, extensionPath: extension.path, ...options }); if (options.default !== undefined && !runtime.flagValues.has(name)) { runtime.flagValues.set(name, options.default); } },
registerMessageRenderer<T>(customType: string, renderer: MessageRenderer<T>): void { runtime.assertActive(); extension.messageRenderers.set(customType, renderer as MessageRenderer); },
getFlag(name: string): boolean | string | undefined { runtime.assertActive(); if (!extension.flags.has(name)) return undefined; return runtime.flagValues.get(name); },
sendMessage(message, options): void { runtime.assertActive(); runtime.sendMessage(message, options); },
sendUserMessage(content, options): void { runtime.assertActive(); runtime.sendUserMessage(content, options); },
appendEntry(customType: string, data?: unknown): void { runtime.assertActive(); runtime.appendEntry(customType, data); },
setSessionName(name: string): void | Promise<void> { runtime.assertActive(); return runtime.setSessionName(name); },
getSessionName(): string | undefined { runtime.assertActive(); return runtime.getSessionName(); },
setLabel(entryId: string, label: string | undefined): void { runtime.assertActive(); runtime.setLabel(entryId, label); },
exec(command: string, args: string[], options?: ExecOptions) { runtime.assertActive(); // Read the host-supplied env at call time so per-session vars (e.g. // herdr pane identity) are current, then let an explicit options.env win. const sessionEnv = runtime.getExecEnv?.(); const env = sessionEnv || options?.env ? { ...sessionEnv, ...options?.env } : undefined; return execCommand(command, args, options?.cwd ?? cwd, { ...options, env }); },
getActiveTools(): string[] { runtime.assertActive(); return runtime.getActiveTools(); },
getAllTools() { runtime.assertActive(); return runtime.getAllTools(); },
setActiveTools(toolNames: string[]): void { runtime.assertActive(); runtime.setActiveTools(toolNames); },
getCommands() { runtime.assertActive(); return runtime.getCommands(); },
setModel(model) { runtime.assertActive(); return runtime.setModel(model); },
getThinkingLevel() { runtime.assertActive(); return runtime.getThinkingLevel(); },
setThinkingLevel(level) { runtime.assertActive(); runtime.setThinkingLevel(level); },
registerProvider(name: string, config: ProviderConfig) { runtime.assertActive(); runtime.registerProvider(name, config, extension.path); },
unregisterProvider(name: string) { runtime.assertActive(); runtime.unregisterProvider(name, extension.path); },
events: eventBus, } as ExtensionAPI;
return api;}
// Replaced with `true` by the esbuild CLI bundle (scripts/bundle.mjs); stays// undefined in unbundled dist/ and under tsx.declare const __PI_BUNDLED__: boolean | undefined;const isBundledCli = typeof __PI_BUNDLED__ !== "undefined" && __PI_BUNDLED__ === true;
let extensionJiti: Jiti | undefined;
async function createExtensionJiti(): Promise<Jiti> { // jiti and the bundled virtual modules are loaded lazily so that importing // the loader (which nearly every startup path does transitively) doesn't pay // for the full package graph; both specifiers are literals, so Bun still // bundles them into the compiled binary. const { createJiti } = await import("jiti/static"); // In the Bun binary and the esbuild CLI bundle: serve pi packages from // virtualModules so extensions share the bundle's module instances // (file-path aliases would load a second, divergent copy of each package). // Also disable tryNative so jiti handles ALL imports (not just the entry point) // In Node.js/dev: use aliases to resolve to node_modules paths return createJiti(import.meta.url, { // Module bodies evaluate once per process; factories still re-run per // load, so top-level extension edits only take effect after a restart. moduleCache: true, ...(isBunBinary || isBundledCli ? { virtualModules: (await import("./bundled-modules.js")).VIRTUAL_MODULES, tryNative: false } : { alias: getAliases() }), });}
async function loadExtensionModule(extensionPath: string) { extensionJiti ??= await createExtensionJiti(); const module = await extensionJiti.import(extensionPath, { default: true }); const factory = module as ExtensionFactory; return typeof factory !== "function" ? undefined : factory;}
/** * Create an Extension object with empty collections. */function createExtension(extensionPath: string, resolvedPath: string): Extension { const source = extensionPath.startsWith("<") && extensionPath.endsWith(">") ? extensionPath.slice(1, -1).split(":")[0] || "temporary" : "local"; const baseDir = extensionPath.startsWith("<") ? undefined : path.dirname(resolvedPath);
return { path: extensionPath, resolvedPath, sourceInfo: createSyntheticSourceInfo(extensionPath, { source, baseDir }), handlers: new Map(), tools: new Map(), messageRenderers: new Map(), commands: new Map(), flags: new Map(), shortcuts: new Map(), };}
async function loadExtension( extensionPath: string, cwd: string, eventBus: EventBus, runtime: ExtensionRuntime,): Promise<{ extension: Extension | null; error: string | null }> { const resolvedPath = resolvePath(extensionPath, cwd);
try { const factory = await loadExtensionModule(resolvedPath); if (!factory) { return { extension: null, error: `Extension does not export a valid factory function: ${extensionPath}` }; }
const extension = createExtension(extensionPath, resolvedPath); const api = createExtensionAPI(extension, runtime, cwd, eventBus); await factory(api);
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: ExtensionFactory, cwd: string, eventBus: EventBus, runtime: ExtensionRuntime, extensionPath = "<inline>",): Promise<Extension> { const extension = createExtension(extensionPath, extensionPath); const api = createExtensionAPI(extension, runtime, cwd, eventBus); await factory(api); return extension;}
/** * Load extensions from paths. */export async function loadExtensions(paths: string[], cwd: string, eventBus?: EventBus): Promise<LoadExtensionsResult> { const extensions: Extension[] = []; const errors: Array<{ path: string; error: string }> = []; const resolvedEventBus = eventBus ?? createEventBus(); const runtime = createExtensionRuntime();
for (const extPath of paths) { const { extension, error } = await loadExtension(extPath, cwd, resolvedEventBus, runtime); time(`loader.ext:${path.basename(extPath)}`);
if (error) { errors.push({ path: extPath, error }); continue; }
if (extension) { extensions.push(extension); } }
return { extensions, errors, runtime, };}
interface PiManifest { extensions?: string[]; themes?: string[]; skills?: string[]; prompts?: string[];}
function readPiManifest(packageJsonPath: string): PiManifest | null { try { const content = fs.readFileSync(packageJsonPath, "utf-8"); const pkg = JSON.parse(content); if (pkg.pi && typeof pkg.pi === "object") { return pkg.pi as PiManifest; } return null; } catch { return null; }}
function isExtensionFile(name: string): boolean { 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: string): string[] | null { const packageJsonPath = path.join(dir, "package.json"); if (fs.existsSync(packageJsonPath)) { const manifest = readPiManifest(packageJsonPath); if (manifest?.extensions?.length) { const entries: string[] = []; 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; } } }
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: string): string[] { if (!fs.existsSync(dir)) { return []; }
const discovered: string[] = [];
try { const entries = fs.readdirSync(dir, { withFileTypes: true });
for (const entry of entries) { const entryPath = path.join(dir, entry.name);
if ((entry.isFile() || entry.isSymbolicLink()) && isExtensionFile(entry.name)) { discovered.push(entryPath); continue; }
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: string[], cwd: string, agentDir: string = getAgentDir(), eventBus?: EventBus,): Promise<LoadExtensionsResult> { const allPaths: string[] = []; const seen = new Set<string>();
const addPaths = (paths: string[]) => { for (const p of paths) { const resolved = path.resolve(p); if (!seen.has(resolved)) { seen.add(resolved); allPaths.push(p); } } };
const localExtDir = path.join(cwd, CONFIG_DIR_NAME, "extensions"); addPaths(discoverExtensionsInDir(localExtDir));
const globalExtDir = path.join(agentDir, "extensions"); addPaths(discoverExtensionsInDir(globalExtDir));
for (const p of configuredPaths) { const resolved = resolvePath(p, cwd); if (fs.existsSync(resolved) && fs.statSync(resolved).isDirectory()) { const entries = resolveExtensionEntries(resolved); if (entries) { addPaths(entries); continue; } addPaths(discoverExtensionsInDir(resolved)); continue; }
addPaths([resolved]); }
return loadExtensions(allPaths, cwd, eventBus);}