/** * Slim main-agent selector extension (spec MS-001). * * Purpose: discover prompt-persona agent definitions, switch the active persona * via the `/mainagent` command, append the persona body to Pi's system prompt * each turn, apply the persona's model/thinking/tools to the current session * (session-scoped: new sessions keep Pi's configured defaults), persist the * selection, and reflect it in the status widget. An interactive TUI menu * (bare `/mainagent`) also edits the model/effort/description of a * definition file and activates the selected agent. * * Tool filters layer across agent switches: an agent without its own tools * configuration keeps the restrictions applied by the previous one (restored * on deactivation; faithful to piagents.ts_). * * Supported frontmatter keys (anything else in the definition file is * silently ignored by this extension): * name — unique persona name; falls back to the file basename. * description — short display description (picker, list rendering). * model — provider/model applied at activation (session-scoped). * thinking — thinking level applied at activation (session-scoped). * tools — optional tool whitelist, wildcards allowed. * excludeTools — optional tool blacklist, wildcards allowed. * prompt — path to an external Markdown prompt file (absolute, or * relative to the definition file); its content becomes the * persona body, prepended to any inline body below the * frontmatter. * * Prompt semantics: the persona body is APPENDED to Pi's system prompt each * turn (append-only). The base prompt, APPEND_SYSTEM.md, project context * files and skills are never stripped. Frontmatter keys belonging to other * extensions (e.g. pi-subagents' `systemPromptMode`, `inheritProjectContext`, * `inheritSkills`, `acceptanceRole`) are not part of this contract and are * ignored here. * * Out of scope (vs piagents.ts_): the subagent orchestration tool and command. */ import fs from "node:fs"; import path from "node:path"; import { DynamicBorder, getAgentDir, parseFrontmatter, } from "@earendil-works/pi-coding-agent"; import type { ExtensionAPI, ExtensionContext, Theme, } from "@earendil-works/pi-coding-agent"; import type { ThinkingLevel } from "@earendil-works/pi-agent-core"; import { clampThinkingLevel, getSupportedThinkingLevels } from "@earendil-works/pi-ai"; import type { Api, Model } from "@earendil-works/pi-ai"; import { matchesKey, SelectList, Text } from "@earendil-works/pi-tui"; import type { SelectItem, SelectListTheme } from "@earendil-works/pi-tui"; // --------------------------------------------------------------------------- // Types // --------------------------------------------------------------------------- /** A discovered prompt-persona agent definition. */ interface AgentDef { /** Absolute path of the source definition file (display/logging only). */ filePath: string; /** Unique agent name used for activation and persistence. */ name: string; /** Short display description. */ description?: string; /** Provider/model identifier applied to the session at activation. */ model?: string; /** Thinking level applied to the session at activation. */ thinking?: ThinkingLevel; /** Tool whitelist (wildcards allowed) applied to the session at activation. */ tools?: string[]; /** Tool blacklist (wildcards allowed) applied to the session at activation. */ excludeTools?: string[]; /** Persona markdown appended to Pi's system prompt. */ body: string; } /** Field updates applied to an agent definition file from the interactive menu. */ interface AgentDefUpdates { /** New provider/model identifier. */ model?: string; /** New thinking level. */ thinking?: ThinkingLevel; /** New description text (empty clears the field). */ description?: string; } /** A model and effort pair chosen in the combined picker dialog. */ interface ModelEffortChoice { /** Chosen provider/model identifier. */ model: string; /** Chosen thinking level. */ effort: ThinkingLevel; } // --------------------------------------------------------------------------- // Constants // --------------------------------------------------------------------------- const LOG_PREFIX = "[main-agent]"; const COMMAND_NAME = "mainagent"; const WIDGET_ID = "main-agent-status"; const AGENTS_SUBDIR = "agents"; const STATE_FILENAME = "agent-state.json"; const PROMPT_METADATA_KEY = "prompt"; const FRONTMATTER_SEPARATOR = "\n\n"; const UNSUPPORTED_YAML_SCALAR_PREFIX = /^[!&*[{|>-]/; const MODEL_ID_SEPARATOR = "/"; const MAX_PROMPT_FILE_BYTES = 64 * 1024; const DEACTIVATION_KEYWORD = "off"; const THINKING_LEVELS: readonly ThinkingLevel[] = [ "off", "minimal", "low", "medium", "high", "xhigh", "max", ]; const MODEL_METADATA_KEY = "model"; const THINKING_METADATA_KEY = "thinking"; const DESCRIPTION_METADATA_KEY = "description"; const FRONTMATTER_BOUNDARY = "---"; const DESCRIPTION_BLOCK_INDENT = " "; const UTF8_BOM = "\uFEFF"; const DEFAULT_EFFORT: ThinkingLevel = "medium"; const MAX_VISIBLE_MENU_ITEMS = 10; const CURRENT_MODEL_LABEL = "current"; const ACTION_CHANGE_MODEL = "Change model & effort"; const ACTION_EDIT_DESCRIPTION = "Edit description"; const ACTION_SWITCH = "Switch to this agent"; const ACTION_BACK = "Back"; // --------------------------------------------------------------------------- // Tool-name normalization (frontmatter metadata) // --------------------------------------------------------------------------- /** * Normalizes a supported scalar tool name for runtime tool filtering. * * Input: raw YAML scalar text. Output: normalized tool name, if supported. * Side effects: none. */ function normalizeToolName(value: string): string | undefined { const trimmed = value.trim(); if (!trimmed || UNSUPPORTED_YAML_SCALAR_PREFIX.test(trimmed)) return undefined; const isQuoted = (trimmed.startsWith('"') && trimmed.endsWith('"')) || (trimmed.startsWith("'") && trimmed.endsWith("'")); const toolName = isQuoted ? trimmed.slice(1, -1).trim() : trimmed; return toolName || undefined; } /** * Parses a tools-list frontmatter field (YAML list or CSV) for runtime tool * filtering. * * Input: raw frontmatter value. Output: normalized tool names, when any. * Side effects: none. */ function parseToolList(value: unknown): string[] | undefined { // Accept either a YAML sequence or a comma-separated string form. const values = Array.isArray(value) ? value : typeof value === "string" ? value.split(",") : []; const tools = values .filter((tool): tool is string => typeof tool === "string") .map(normalizeToolName) .filter((toolName): toolName is string => Boolean(toolName)); return tools.length > 0 ? tools : undefined; } /** * Validates a raw thinking-level frontmatter value against the supported levels. * * Input: raw frontmatter value. Output: validated thinking level, if supported. * Side effects: none. */ function parseThinkingLevel(value: unknown): ThinkingLevel | undefined { if (typeof value !== "string") return undefined; const level = value.trim().toLowerCase() as ThinkingLevel; return (THINKING_LEVELS as readonly string[]).includes(level) ? level : undefined; } /** * Escapes regex metacharacters in one wildcard-pattern segment. * * Input: literal pattern segment. Output: escaped segment. * Side effects: none. */ function escapeToolPatternSegment(segment: string): string { return segment.replace(/[.*+?^${}()|[\]\\]/g, "\\$&"); } /** * Checks whether a tool name matches a pattern containing '*' wildcards. * * Input: candidate tool name and pattern (e.g. "ssh-manager_*"). Output: match result. * Side effects: none. */ function matchesToolPattern(toolName: string, pattern: string): boolean { const regex = new RegExp(`^${pattern.split("*").map(escapeToolPatternSegment).join(".*")}$`); return regex.test(toolName); } /** * Expands wildcard entries ("ssh-manager_*") against the available tool names. * * Inputs: whitelist entries and registered tool names. Output: effective tool list. * Exact entries stay verbatim; patterns matching nothing are dropped. An empty * available list returns the patterns unchanged (safe fallback). * Side effects: none. */ function expandToolPatterns(patterns: string[], availableToolNames: readonly string[]): string[] { if (availableToolNames.length === 0 || !patterns.some((entry) => entry.includes("*"))) { return patterns; } const resolved: string[] = []; for (const entry of patterns) { if (!entry.includes("*")) { resolved.push(entry); continue; } for (const name of availableToolNames) { if (matchesToolPattern(name, entry) && !resolved.includes(name)) resolved.push(name); } } return resolved; } /** * Lists registered tool names for wildcard expansion. * * Input: Pi extension API. Output: registered tool names; empty when unavailable. * Side effects: none. */ function listRegisteredToolNames(pi: ExtensionAPI): string[] { try { return pi.getAllTools().map((tool) => tool.name); } catch { // Tool metadata may be unavailable before Pi finishes binding its actions. return []; } } // --------------------------------------------------------------------------- // Paths // --------------------------------------------------------------------------- /** * Resolves the directory containing agent definition files. * * Input: Pi's configured agent directory. Output: absolute agents directory. * Side effects: none. */ function agentsDir(): string { return path.join(getAgentDir(), AGENTS_SUBDIR); } /** * Resolves the persisted active-agent state file path. * * Input: Pi's configured agent directory. Output: absolute state-file path. * Side effects: none. */ function stateFile(): string { return path.join(getAgentDir(), STATE_FILENAME); } // --------------------------------------------------------------------------- // Persistence // --------------------------------------------------------------------------- /** * Reads the persisted active-agent name from agent-state.json. * * Input: state file contents. Output: active agent name, or undefined when no * selection is persisted or the file is missing, corrupt, or wrong-shaped. * Side effects: reads the state file. */ function readState(): string | undefined { try { // Parse the persisted JSON state and validate the `{ active }` shape. const parsed: unknown = JSON.parse(fs.readFileSync(stateFile(), "utf-8")); if (typeof parsed !== "object" || parsed === null || Array.isArray(parsed)) { return undefined; } const active = (parsed as { active?: unknown }).active; if (typeof active === "string" && active.length > 0) return active; return undefined; } catch { // Missing or malformed state simply means "no persisted selection". return undefined; } } /** * Persists the selected agent name or the deactivated state. * * Input: agent name or `null`. Output: whether persistence succeeded. * Side effects: writes agent-state.json; logs write failures. */ function writeState(active: string | null): boolean { try { fs.writeFileSync(stateFile(), JSON.stringify({ active }, null, 2), "utf-8"); return true; } catch (err) { // Report failures so the caller can show a warning notification. console.error(`${LOG_PREFIX} failed to write state:`, err); return false; } } // --------------------------------------------------------------------------- // Agent discovery // --------------------------------------------------------------------------- /** * Reads the optional Markdown prompt referenced by an agent definition. * * Inputs: parsed prompt metadata and the definition file path. Output: trimmed * Markdown content, or an empty string when no prompt file is configured. * Side effects: reads the configured prompt file. */ function readPromptFile(prompt: unknown, definitionPath: string): string { // Keep inline-only agent definitions valid when they omit the prompt field. if (typeof prompt !== "string" || !prompt.trim()) return ""; // Absolute prompt paths are supported by design; relative ones resolve // against the definition that declares the prompt. const configuredPath = prompt.trim(); const promptPath = path.isAbsolute(configuredPath) ? configuredPath : path.resolve(path.dirname(definitionPath), configuredPath); // Refuse non-regular files (a FIFO would block the read) and oversized ones. const stat = fs.statSync(promptPath); if (!stat.isFile() || stat.size > MAX_PROMPT_FILE_BYTES) { throw new Error( `prompt file must be a regular file of at most ${MAX_PROMPT_FILE_BYTES} bytes: ${promptPath}`, ); } return fs.readFileSync(promptPath, "utf-8").trim(); } /** * Discovers and normalizes agent definitions from the agents directory. * * Input: none (uses the configured agent directory). Output: name-sorted agent * definitions; empty catalog when the directory is missing. * Side effects: reads agent files from the filesystem; logs per-file failures. */ function loadAgents(): AgentDef[] { // Discover the optional agent-definition directory before reading its files. const dir = agentsDir(); let entries: fs.Dirent[]; try { if (!fs.existsSync(dir)) { console.warn(`${LOG_PREFIX} agent definitions directory is unavailable: ${dir}`); return []; } entries = fs.readdirSync(dir, { withFileTypes: true }); } catch (err) { console.error(`${LOG_PREFIX} failed to discover agent definitions in "${dir}":`, err); return []; } // Read each Markdown definition and convert it into an agent record. const agents: AgentDef[] = []; const seenNames = new Set(); for (const entry of entries) { if (!entry.isFile() || !entry.name.endsWith(".md")) continue; const filePath = path.join(dir, entry.name); try { if (fs.statSync(filePath).size > MAX_PROMPT_FILE_BYTES) { throw new Error(`definition file exceeds ${MAX_PROMPT_FILE_BYTES} bytes`); } const raw = fs.readFileSync(filePath, "utf-8"); const { frontmatter: meta, body: inlineBody } = parseFrontmatter>(raw); const promptBody = readPromptFile(meta[PROMPT_METADATA_KEY], filePath); // External prompt file first, then the inline body. const body = [promptBody, inlineBody.trim()] .filter((value) => value.length > 0) .join(FRONTMATTER_SEPARATOR); const declaredName = typeof meta.name === "string" ? meta.name.trim() : ""; const name = declaredName || path.basename(entry.name, ".md") || entry.name; if (seenNames.has(name)) { console.warn(`${LOG_PREFIX} duplicate agent name "${name}" in "${filePath}", skipped`); continue; } seenNames.add(name); agents.push({ filePath, name, description: typeof meta.description === "string" && meta.description ? meta.description : undefined, model: typeof meta.model === "string" ? meta.model : undefined, thinking: parseThinkingLevel(meta.thinking), tools: parseToolList(meta.tools), excludeTools: parseToolList(meta.excludeTools ?? meta.exclude_tools), body, }); } catch (err) { // Isolate per-file failures so one bad definition cannot break loading. console.error(`${LOG_PREFIX} failed to load agent definition "${filePath}":`, err); } } // Return definitions in a stable order for display and completion. agents.sort((first, second) => first.name.localeCompare(second.name)); return agents; } // --------------------------------------------------------------------------- // Command rendering // --------------------------------------------------------------------------- /** * Builds the display-only metadata line for one agent in the list output. * * Input: agent definition. Output: bracketed metadata line, or empty string * when the agent defines no metadata. * Side effects: none. */ function buildMetaLine(agent: AgentDef): string { const parts = [ agent.model, agent.thinking ? `thinking:${agent.thinking}` : undefined, agent.tools ? `tools:${agent.tools.join(",")}` : undefined, agent.excludeTools ? `exclude:${agent.excludeTools.join(",")}` : undefined, ].filter((part): part is string => Boolean(part)); return parts.length > 0 ? ` [${parts.join(" | ")}]` : ""; } /** * Renders the `/mainagent` catalog listing. * * Inputs: discovered agents and the currently active one. Output: multi-line * listing text with active-state markers. * Side effects: none. */ function renderAgentList(agents: AgentDef[], active: AgentDef | null): string { const lines = agents.map((agent) => { const marker = active && active.name === agent.name ? "●" : "○"; return ` ${marker} ${agent.name}${buildMetaLine(agent)}\n ${agent.description ?? ""}`; }); return `Agents (${agents.length}):\n${lines.join("\n")}`; } // --------------------------------------------------------------------------- // Definition editing (frontmatter surgery) // --------------------------------------------------------------------------- /** * Checks whether a frontmatter line continues a block-scalar value. * * Input: one frontmatter line. Output: whether it is indented content. * Side effects: none. */ function isBlockValueLine(line: string): boolean { return line.length > 0 && /^\s/.test(line); } /** * Finds the extent of a frontmatter key's value: the key line plus its * continuation lines (indented content, and blank lines only when further * indented content follows before the next column-0 line). Indented `#` * comment lines are consumed too: they are indistinguishable from folded * scalar content, so replacements drop them (accepted tradeoff). * * Inputs: frontmatter lines and the key line index. Output: exclusive end * index of the value extent. * Side effects: none. */ function findKeyExtent(lines: string[], keyIndex: number): number { let end = keyIndex + 1; while (end < lines.length) { if (isBlockValueLine(lines[end])) { end++; continue; } if (lines[end] === "") { // Blank lines belong to the block only when indented content follows. let lookahead = end + 1; while (lookahead < lines.length && lines[lookahead] === "") lookahead++; if (lookahead < lines.length && isBlockValueLine(lines[lookahead])) { end = lookahead + 1; continue; } } break; } return end; } /** * Replaces or inserts a single-line frontmatter key. * * Inputs: frontmatter lines, key, and scalar value. Output: none (mutates * lines). Continuation lines of a multi-line hand-authored value are consumed * so they cannot orphan and corrupt the YAML. * Side effects: none. */ function setFrontmatterScalar(lines: string[], key: string, value: string): void { const index = lines.findIndex((line) => line.startsWith(`${key}:`)); if (index === -1) { lines.push(`${key}: ${value}`); return; } const end = findKeyExtent(lines, index); lines.splice(index, end - index, `${key}: ${value}`); } /** * Replaces or inserts the description frontmatter key as a YAML block scalar. * * Inputs: frontmatter lines and the new description. Output: none (mutates * lines). An empty description becomes an explicit empty quoted scalar so the * loader drops the field. * Side effects: none. */ function setFrontmatterDescription(lines: string[], value: string): void { const keyIndex = lines.findIndex((line) => line.startsWith(`${DESCRIPTION_METADATA_KEY}:`)); const replacement = value.trim() ? [ `${DESCRIPTION_METADATA_KEY}: |`, ...value.trimEnd().split("\n").map((line) => (line ? DESCRIPTION_BLOCK_INDENT + line : "")), ] : [`${DESCRIPTION_METADATA_KEY}: ""`]; if (keyIndex === -1) lines.push(...replacement); else lines.splice(keyIndex, findKeyExtent(lines, keyIndex) - keyIndex, ...replacement); } /** * Builds the confirmation summary lines for pending definition updates. * * Inputs: agent definition and updates. Output: changed-field summary lines. * Side effects: none. */ function buildUpdateSummary(agent: AgentDef, updates: AgentDefUpdates): string[] { const lines = [path.basename(agent.filePath)]; if (updates.model !== undefined && updates.model !== agent.model) { lines.push(`model: ${agent.model ?? "—"} → ${updates.model}`); } if (updates.thinking !== undefined && updates.thinking !== agent.thinking) { lines.push(`effort: ${agent.thinking ?? "—"} → ${updates.thinking}`); } if (updates.description !== undefined) lines.push("description: updated"); return lines; } /** * Applies one set of menu updates to parsed frontmatter lines. * * Inputs: frontmatter lines and updates. Output: none (mutates lines). * Side effects: none. */ function applyFrontmatterUpdates(lines: string[], updates: AgentDefUpdates): void { if (updates.model !== undefined) { setFrontmatterScalar(lines, MODEL_METADATA_KEY, updates.model); } if (updates.thinking !== undefined) { setFrontmatterScalar(lines, THINKING_METADATA_KEY, updates.thinking); } if (updates.description !== undefined) { setFrontmatterDescription(lines, updates.description); } } /** * Checks whether the loader would parse frontmatter that this writer's * strict opening delimiter (`---` + line ending) cannot rewrite. * * Input: file contents without a BOM. Output: true when the loader sees a * non-empty frontmatter block this writer cannot match (e.g. a `----` or * `--- ` opener); an empty payload is body-only for the loader too. * Side effects: none. */ function hasUnsupportedFrontmatter(content: string): boolean { // Mirror the loader's own detection: any "---"-prefixed first line plus a // "\n---" closer, with a non-empty YAML payload between them. const normalized = content.replace(/\r\n/g, "\n").replace(/\r/g, "\n"); if (!normalized.startsWith(FRONTMATTER_BOUNDARY)) return false; const closerIndex = normalized.indexOf(`\n${FRONTMATTER_BOUNDARY}`, FRONTMATTER_BOUNDARY.length); return ( closerIndex !== -1 && normalized.slice(FRONTMATTER_BOUNDARY.length + 1, closerIndex) !== "" ); } /** * Applies definition updates to one agent `.md` file without touching its body. * * Inputs: raw file contents and the updates to apply. Output: rewritten file * contents with an updated frontmatter block. * Side effects: none (pure string transform). */ function updateFrontmatterFields(raw: string, updates: AgentDefUpdates): string { // Keep a UTF-8 BOM out of the rewritten frontmatter and re-emit it first. const bom = raw.startsWith(UTF8_BOM) ? UTF8_BOM : ""; const content = bom ? raw.slice(UTF8_BOM.length) : raw; const match = /^---(\r?\n)([\s\S]*?)\r?\n---/.exec(content); // Refuse openers this writer cannot rewrite instead of demoting the // loader-visible frontmatter into the persona body; body-only files get a // fresh block prepended instead of guessed-at repairs. if (!match) { if (hasUnsupportedFrontmatter(content)) { throw new Error("unsupported frontmatter opening delimiter"); } const freshEol = content.includes("\r\n") ? "\r\n" : "\n"; const fresh: string[] = []; applyFrontmatterUpdates(fresh, updates); return bom + [FRONTMATTER_BOUNDARY, ...fresh, FRONTMATTER_BOUNDARY, "", content].join(freshEol); } // Rebuild only the frontmatter using its own line ending; the body below // the closing delimiter is preserved byte-for-byte. const lines = match[2].split(/\r?\n/); applyFrontmatterUpdates(lines, updates); return ( bom + [FRONTMATTER_BOUNDARY, ...lines, FRONTMATTER_BOUNDARY].join(match[1]) + content.slice(match[0].length) ); } // --------------------------------------------------------------------------- // Interactive menu (TUI dialogs) // --------------------------------------------------------------------------- /** Configuration for one framed SelectList dialog. */ interface SelectDialogSpec { /** Dialog title rendered above the list. */ title: string; /** Selectable items. */ items: SelectItem[]; /** Keymap hint rendered below the list. */ help: string; /** Maps the selected item value to the dialog result. */ resolve: (value: string) => T; /** Extra lines re-evaluated on every render (e.g. the live effort value). */ dynamicLines?: (theme: Theme) => string[]; /** Consumes raw input before the list sees it; returns true when handled. */ onRawInput?: (data: string) => boolean; /** Item value preselected when the dialog opens. */ initialSelection?: string; /** Notified whenever the highlighted item changes (arrow keys, clicks). */ onSelectionChange?: (value: string) => void; } /** * Builds the `provider/id` identifier used as list label and persisted value. * * Input: model. Output: its canonical identifier. Side effects: none. */ function modelIdentifier(model: Model): string { return `${model.provider}${MODEL_ID_SEPARATOR}${model.id}`; } /** * Lists the models selectable in the menu, with their metadata. * * Input: extension context. Output: deduplicated (by identifier), sorted * model objects; scoped models take precedence over the full catalogue. * Side effects: none. */ function buildModelChoices(ctx: ExtensionContext): Model[] { const models: readonly Model[] = ctx.scopedModels.length > 0 ? ctx.scopedModels.map((scoped) => scoped.model) : ctx.modelRegistry.getAvailable(); const byIdentifier = new Map>( models.map((model): [string, Model] => [modelIdentifier(model), model]), ); return [...byIdentifier.keys()].sort().map((identifier) => byIdentifier.get(identifier)!); } /** * Builds the shared visual theme for menu SelectLists. * * Input: extension theme. Output: SelectList theme callbacks. * Side effects: none. */ function buildSelectListTheme(theme: Theme): SelectListTheme { return { selectedPrefix: (text) => theme.fg("accent", text), selectedText: (text) => theme.fg("accent", text), description: (text) => theme.fg("muted", text), scrollInfo: (text) => theme.fg("dim", text), noMatch: (text) => theme.fg("warning", text), }; } /** * Shows one framed SelectList dialog and resolves the selection. * * Inputs: extension context and dialog configuration. Output: resolved * selection, or null when cancelled. * Side effects: renders a modal dialog with keyboard focus. */ function showSelectDialog( ctx: ExtensionContext, spec: SelectDialogSpec, ): Promise { return ctx.ui.custom((tui, theme, _keybindings, done) => { const topBorder = new DynamicBorder((line) => theme.fg("accent", line)); const bottomBorder = new DynamicBorder((line) => theme.fg("accent", line)); const titleText = new Text(theme.fg("accent", theme.bold(spec.title)), 1, 0); const helpText = new Text(theme.fg("dim", spec.help), 1, 0); const selectList = new SelectList(spec.items, MAX_VISIBLE_MENU_ITEMS, buildSelectListTheme(theme)); selectList.onSelect = (item) => done(spec.resolve(item.value)); selectList.onCancel = () => done(null); selectList.onSelectionChange = (item) => spec.onSelectionChange?.(item.value); if (spec.initialSelection !== undefined) { const initialIndex = spec.items.findIndex( (item) => item.value === spec.initialSelection, ); if (initialIndex >= 0) selectList.setSelectedIndex(initialIndex); } return { render: (width) => [ ...topBorder.render(width), ...titleText.render(width), ...(spec.dynamicLines ? spec.dynamicLines(theme) : []), ...selectList.render(width), ...helpText.render(width), ...bottomBorder.render(width), ], invalidate: () => { selectList.invalidate(); titleText.invalidate(); helpText.invalidate(); }, handleInput: (data: string) => { if (!spec.onRawInput?.(data)) selectList.handleInput(data); tui.requestRender(); }, }; }); } /** * Shows the agent selection dialog. * * Inputs: extension context, discovered agents, and the active one. Output: * chosen agent name, or null when cancelled. * Side effects: renders a modal dialog. */ function showAgentPicker( ctx: ExtensionContext, agents: readonly AgentDef[], active: AgentDef | null, ): Promise { const items: SelectItem[] = agents.map((agent) => ({ value: agent.name, label: agent.name, description: [ active === agent ? "● active" : undefined, `model: ${agent.model ?? "—"}`, `effort: ${agent.thinking ?? "—"}`, ] .filter((part): part is string => Boolean(part)) .join(" · "), })); return showSelectDialog(ctx, { title: "Main agent", items, help: "↑↓ select · enter open · esc cancel", resolve: (value) => value, }); } /** * Shows the per-agent action menu. * * Inputs: extension context and the agent being managed. Output: chosen action * label, or null when cancelled. * Side effects: renders a modal dialog. */ function showAgentActions(ctx: ExtensionContext, agent: AgentDef): Promise { const title = `${agent.name} — model: ${agent.model ?? "—"} · effort: ${agent.thinking ?? "—"}`; return showSelectDialog(ctx, { title, items: [ { value: ACTION_CHANGE_MODEL, label: ACTION_CHANGE_MODEL, description: "←→ cycles the effort while picking the model", }, { value: ACTION_EDIT_DESCRIPTION, label: ACTION_EDIT_DESCRIPTION }, { value: ACTION_SWITCH, label: ACTION_SWITCH }, { value: ACTION_BACK, label: ACTION_BACK }, ], help: "↑↓ select · enter confirm · esc cancel", resolve: (value) => value, }); } /** * Shows the combined model and effort picker: ↑↓ picks the model, ←→ cycles * the effort level shown live between the title and the list. Only the * levels supported by the highlighted model are offered, and the current * effort is clamped to them via the same SDK helpers the session runtime * uses for setThinkingLevel, so a saved pair never silently diverges at * runtime. * * Inputs: extension context and the agent being edited. Output: the chosen * model/effort pair, or null when cancelled. * Side effects: renders a modal dialog; notifies when no model is available. */ function showModelEffortPicker( ctx: ExtensionContext, agent: AgentDef, ): Promise { const models = buildModelChoices(ctx); if (models.length === 0) { ctx.ui.notify("No models available to pick from.", "warning"); return Promise.resolve(null); } const initialModel = models.find((model) => modelIdentifier(model) === agent.model) ?? models[0]; let currentModel: Model = initialModel; let effort: ThinkingLevel = clampThinkingLevel( initialModel, agent.thinking ?? DEFAULT_EFFORT, ); return showSelectDialog(ctx, { title: `Model & effort — ${agent.name}`, items: models.map((model) => { const identifier = modelIdentifier(model); return { value: identifier, label: identifier, description: identifier === agent.model ? CURRENT_MODEL_LABEL : undefined, }; }), initialSelection: modelIdentifier(initialModel), help: "↑↓ model · ←→ effort · enter save · esc cancel", onSelectionChange: (identifier) => { const model = models.find((entry) => modelIdentifier(entry) === identifier); if (!model) return; currentModel = model; effort = clampThinkingLevel(model, effort); }, resolve: (identifier) => { const model = models.find((entry) => modelIdentifier(entry) === identifier) ?? currentModel; return { model: identifier, effort: clampThinkingLevel(model, effort) }; }, dynamicLines: (theme) => [ "", theme.fg("accent", ` effort: ‹ ${effort} ›`), ], onRawInput: (data) => { const isLeft = matchesKey(data, "left"); const isRight = matchesKey(data, "right"); if (!isLeft && !isRight) return false; const levels = getSupportedThinkingLevels(currentModel); const index = Math.max(0, levels.indexOf(effort)); const direction = isRight ? 1 : -1; effort = levels[(index + direction + levels.length) % levels.length]; return true; }, }); } // --------------------------------------------------------------------------- // Extension // --------------------------------------------------------------------------- /** * Registers the main-agent selector features with Pi. * * Input: Pi extension API. Output: none; command and lifecycle handlers register. * Side effects: reads agent files, persists selections, updates the status * widget, and appends the persona body to the system prompt each turn. */ function registerMainAgentFeatures(pi: ExtensionAPI): void { // Initialize discovered definitions, the mutable active selection, and the // session runtime baseline captured before the first agent setting. const agents = loadAgents(); let active: AgentDef | null = null; let runtimeBaseline: { activeTools: string[]; thinkingLevel: ThinkingLevel; model: ExtensionContext["model"]; } | null = null; /** * Finds a discovered agent by its configured name (exact match first, then * case-insensitive fallback for typed input). * * Input: agent name. Output: matching definition, if any. * Side effects: none. */ const findAgent = (name: string): AgentDef | undefined => agents.find((agent) => agent.name === name) ?? agents.find((agent) => agent.name.toLowerCase() === name.toLowerCase()); /** * Updates the TUI widget with the active agent name. * * Input: extension context. Output: none. * Side effects: replaces the status widget content. */ const updateAgentStatus = (ctx: ExtensionContext): void => { ctx.ui.setWidget( WIDGET_ID, [`Main Agent: ${active?.name ?? "none"}`], { placement: "belowEditor" }, ); }; /** * Captures runtime settings before the first agent setting is applied. * * Input: extension context. Output: whether a baseline is available. * Side effects: reads Pi runtime settings and stores them in session state. */ const captureBaseline = (ctx: ExtensionContext): boolean => { // Preserve an existing baseline across agent changes in the session. if (runtimeBaseline) return true; try { // Snapshot the tools, thinking level, and model from the active runtime. runtimeBaseline = { activeTools: [...pi.getActiveTools()], thinkingLevel: pi.getThinkingLevel(), model: ctx.model, }; return true; } catch (err) { console.error(`${LOG_PREFIX} failed to capture runtime baseline:`, err); return false; } }; /** * Restores the runtime settings captured before agent activation. * * Input: stored runtime baseline. Output: whether all settings were restored. * Side effects: changes Pi tools, thinking level, and model. */ const restoreBaseline = async (): Promise => { // Return immediately when no agent settings altered this runtime. const baseline = runtimeBaseline; if (!baseline) return true; // Track partial restore failures while attempting every setting. let restored = true; try { pi.setActiveTools([...baseline.activeTools]); } catch (err) { restored = false; console.error(`${LOG_PREFIX} failed to restore active tools:`, err); } // Restore the original model only when a baseline model was available. if (baseline.model) { try { const ok = await pi.setModel(baseline.model); if (!ok) { restored = false; console.error(`${LOG_PREFIX} failed to restore model: setter returned false`); } } catch (err) { restored = false; console.error(`${LOG_PREFIX} failed to restore model:`, err); } } try { pi.setThinkingLevel(baseline.thinkingLevel); } catch (err) { restored = false; console.error(`${LOG_PREFIX} failed to restore thinking level:`, err); } // Clear the snapshot only after all original settings were restored. if (restored) runtimeBaseline = null; return restored; }; /** * Applies an agent's tool whitelist/blacklist to the active runtime. * * Input: agent definition. Output: summary of the applied filters, if any. * Side effects: changes Pi's active tools. */ const applyAgentTools = (agent: AgentDef): string | undefined => { const registeredToolNames = listRegisteredToolNames(pi); const whitelist = agent.tools && agent.tools.length > 0 ? expandToolPatterns(agent.tools, registeredToolNames) : undefined; const excludedTools = agent.excludeTools ? expandToolPatterns(agent.excludeTools, registeredToolNames) : []; if (!whitelist && excludedTools.length === 0) return undefined; // With only a blacklist, start from the active tools so existing runtime // restrictions survive. const baseTools = whitelist ?? pi.getActiveTools(); pi.setActiveTools(baseTools.filter((toolName) => !excludedTools.includes(toolName))); const toolsLabel = whitelist && agent.tools ? whitelist.length === agent.tools.length ? agent.tools.join(",") : `${agent.tools.join(",")} → ${whitelist.length} tools` : undefined; const excludeLabel = excludedTools.length > 0 ? `exclude: ${excludedTools.join(",")}` : undefined; return [toolsLabel, excludeLabel] .filter((label): label is string => Boolean(label)) .join(" | "); }; /** * Resolves and applies an agent's provider/model identifier to this session. * * Inputs: agent definition and extension context. Output: summary label when * the model was applied, otherwise undefined. * Side effects: changes Pi's model; notifies the TUI on failures. */ const applyAgentModel = async ( agent: AgentDef, ctx: ExtensionContext, ): Promise => { const configured = agent.model; if (!configured) return undefined; const separatorIndex = configured.indexOf(MODEL_ID_SEPARATOR); if (separatorIndex === -1) { ctx.ui.notify( `Agent "${agent.name}": model "${configured}" must use the provider${MODEL_ID_SEPARATOR}model form`, "warning", ); return undefined; } const provider = configured.slice(0, separatorIndex); const modelId = configured.slice(separatorIndex + MODEL_ID_SEPARATOR.length); const model = ctx.modelRegistry.find(provider, modelId); if (!model) { ctx.ui.notify( `Agent "${agent.name}": model "${configured}" not found in the registry`, "warning", ); return undefined; } try { if (!(await pi.setModel(model))) { console.error(`${LOG_PREFIX} setModel returned false`); ctx.ui.notify(`Agent "${agent.name}": unable to set model ${configured}`, "warning"); return undefined; } return `model: ${configured}`; } catch (err) { console.error(`${LOG_PREFIX} setModel failed:`, err); ctx.ui.notify(`Agent "${agent.name}": unable to set model ${configured}`, "warning"); return undefined; } }; /** * Applies an agent's model, thinking level, and tool filters to this session. * * Inputs: agent definition and extension context. Output: descriptions of the * settings that were applied successfully. * Side effects: changes Pi tools, thinking level, and model; notifies the TUI * on partial failures. */ const applyAgentRuntime = async ( agent: AgentDef, ctx: ExtensionContext, ): Promise => { const applied: string[] = []; // Apply the optional tool whitelist and blacklist first. try { const toolsSummary = applyAgentTools(agent); if (toolsSummary) applied.push(toolsSummary); } catch (err) { console.error(`${LOG_PREFIX} setActiveTools failed:`, err); } // Resolve and apply the optional provider/model identifier first: // setModel overwrites the thinking level with the per-model or global // default, so the agent's effort must be applied after the model switch. if (agent.model) { const modelLabel = await applyAgentModel(agent, ctx); if (modelLabel) applied.push(modelLabel); } // Apply the optional thinking level (validated at load time); the runtime // clamps it against the model set above. if (agent.thinking) { try { pi.setThinkingLevel(agent.thinking); applied.push(`thinking: ${agent.thinking}`); } catch (err) { console.error(`${LOG_PREFIX} setThinkingLevel failed:`, err); } } return applied; }; /** * Applies an agent selection or deactivation end-to-end. * * Inputs: agent definition or null, extension context. Output: none. * Side effects: captures/restores the runtime baseline, updates the selection, * the widget, the state file, and notifies the TUI. */ const applySelection = async (agent: AgentDef | null, ctx: ExtensionContext): Promise => { // Deactivation: restore the session's pre-agent runtime settings. if (!agent) { if (!captureBaseline(ctx)) { ctx.ui.notify("Unable to capture the pre-agent runtime state", "error"); return; } active = null; updateAgentStatus(ctx); const persisted = writeState(null); const restored = await restoreBaseline(); if (persisted && restored) { ctx.ui.notify("Main agent deactivated.", "info"); } else if (!persisted && !restored) { ctx.ui.notify( "Main agent deactivated, but state persistence and runtime restore are incomplete.", "warning", ); } else if (!persisted) { ctx.ui.notify("Main agent deactivated, but the state could not be saved.", "warning"); } else { ctx.ui.notify("Main agent deactivated, but the runtime restore is incomplete.", "warning"); } return; } // Activation: stop before changing settings if the original runtime cannot // be preserved for a later restore. if (!captureBaseline(ctx)) { ctx.ui.notify( `Agent "${agent.name}": unable to capture the pre-agent runtime state`, "error", ); return; } active = agent; updateAgentStatus(ctx); const persisted = writeState(agent.name); const applied = await applyAgentRuntime(agent, ctx); const runtimeLabel = applied.length > 0 ? ` [${applied.join(" | ")}]` : ""; if (persisted) { ctx.ui.notify(`Active agent: ${agent.name}${runtimeLabel}`, "info"); } else { ctx.ui.notify( `Active agent: ${agent.name}${runtimeLabel} (state could not be saved)`, "warning", ); } }; /** * Persists definition updates end-to-end: confirms with the user, rewrites * the `.md` frontmatter, mutates the in-memory definition, and live-applies * runtime changes when the edited agent is the active one. * * Inputs: agent definition, updates, extension context. Output: none. * Side effects: writes the definition file, updates the in-memory catalog, * may change Pi's runtime model/thinking; notifies the TUI. */ const saveAgentUpdates = async ( agent: AgentDef, updates: AgentDefUpdates, ctx: ExtensionContext, ): Promise => { const summary = buildUpdateSummary(agent, updates).join("\n"); const confirmed = await ctx.ui.confirm( "Write changes to the agent definition?", summary, ); if (!confirmed) return; // Write the file first; only a successful write updates the session state. try { const raw = fs.readFileSync(agent.filePath, "utf-8"); fs.writeFileSync(agent.filePath, updateFrontmatterFields(raw, updates), "utf-8"); } catch (err) { console.error(`${LOG_PREFIX} failed to update agent definition "${agent.filePath}":`, err); ctx.ui.notify( `Unable to update ${path.basename(agent.filePath)}: definition unchanged`, "error", ); return; } // Keep the in-memory catalog coherent with the file for this session. if (updates.model !== undefined) agent.model = updates.model; if (updates.thinking !== undefined) agent.thinking = updates.thinking; if (updates.description !== undefined) { agent.description = updates.description.trim() || undefined; } // Live-apply model/effort edits when the edited agent is the active one; // description edits need no runtime change. if (active === agent && (updates.model !== undefined || updates.thinking !== undefined)) { const applied = await applyAgentRuntime(agent, ctx); const runtimeLabel = applied.length > 0 ? ` [${applied.join(" | ")}]` : ""; ctx.ui.notify(`Updated "${agent.name}"${runtimeLabel}`, "info"); return; } ctx.ui.notify(`Updated "${agent.name}" (${path.basename(agent.filePath)})`, "info"); }; /** * Runs the combined model/effort picker and saves the resulting changes. * * Inputs: agent definition and extension context. Output: none. * Side effects: opens the picker dialog; may write the definition file and * change the runtime; notifies the TUI. */ const handleModelEffortChange = async ( agent: AgentDef, ctx: ExtensionContext, ): Promise => { const choice = await showModelEffortPicker(ctx, agent); if (!choice) return; const updates: AgentDefUpdates = {}; if (choice.model !== agent.model) updates.model = choice.model; if (choice.effort !== agent.thinking) updates.thinking = choice.effort; if (updates.model === undefined && updates.thinking === undefined) { ctx.ui.notify("No changes to save.", "info"); return; } await saveAgentUpdates(agent, updates, ctx); }; /** * Runs the description editor and saves the resulting changes. * * Inputs: agent definition and extension context. Output: none. * Side effects: opens the editor dialog; may write the definition file and * notify the TUI. */ const handleDescriptionChange = async ( agent: AgentDef, ctx: ExtensionContext, ): Promise => { const value = await ctx.ui.editor( `Edit description — ${agent.name}`, (agent.description ?? "").trim(), ); if (value === undefined) return; if (value.trim() === (agent.description ?? "").trim()) { ctx.ui.notify("No changes to save.", "info"); return; } await saveAgentUpdates(agent, { description: value }, ctx); }; /** * Runs the interactive agent menu loop shown by a bare `/mainagent` * (TUI only). * * Input: extension context. Output: none. * Side effects: opens dialogs; may save definition edits and switch agents. */ const runAgentMenu = async (ctx: ExtensionContext): Promise => { let name = await showAgentPicker(ctx, agents, active); while (name !== null) { const agent = findAgent(name); if (!agent) return; const action = await showAgentActions(ctx, agent); if (action === null) return; if (action === ACTION_SWITCH) { await applySelection(agent, ctx); return; } if (action === ACTION_BACK) { name = await showAgentPicker(ctx, agents, active); continue; } if (action === ACTION_CHANGE_MODEL) await handleModelEffortChange(agent, ctx); else if (action === ACTION_EDIT_DESCRIPTION) await handleDescriptionChange(agent, ctx); } }; // Restore a persisted selection at registration when the agent still exists. const persisted = readState(); if (typeof persisted === "string") { active = findAgent(persisted) ?? null; if (!active) { console.warn( `${LOG_PREFIX} state references unknown agent "${persisted}", ignoring`, ); } } pi.on("session_start", async (_event, ctx) => { // Re-apply a persisted agent to this session's runtime without // re-persisting the selection: another Pi instance may have deactivated // the agent since this state was written. if (active) { if (captureBaseline(ctx)) { await applyAgentRuntime(active, ctx); } else { ctx.ui.notify("Unable to capture the pre-agent runtime state", "error"); } } updateAgentStatus(ctx); }); pi.on("session_shutdown", async () => { // Give the session back its pre-agent runtime settings. await restoreBaseline(); }); pi.on("before_agent_start", (event, _ctx) => { // Reapply the tool filters; other code may have changed tools mid-session. if ( active && ((active.tools && active.tools.length > 0) || (active.excludeTools && active.excludeTools.length > 0)) ) { try { applyAgentTools(active); } catch (err) { console.error(`${LOG_PREFIX} per-turn tool reapply failed:`, err); } } // Keep the unmodified Pi prompt when no persona body is active. if (!active || !active.body) return; return { systemPrompt: event.systemPrompt + FRONTMATTER_SEPARATOR + active.body }; }); pi.registerCommand(COMMAND_NAME, { description: "Show or change the active agent. No arguments: interactive menu (TUI). Usage: /mainagent [name|off]. E.g. /mainagent gv-dev", /** * Provides agent names that begin with the command argument prefix. * * Input: completion prefix. Output: matching completion items. * Side effects: none. */ getArgumentCompletions: (prefix: string) => { const lowerPrefix = prefix.toLowerCase(); const agentItems = agents .filter((agent) => agent.name.toLowerCase().startsWith(lowerPrefix)) .map((agent) => ({ value: agent.name, label: agent.name, description: agent.description, })); // Offer the deactivation keyword unless a real agent claims that name. const offItems = DEACTIVATION_KEYWORD.startsWith(lowerPrefix) && !findAgent(DEACTIVATION_KEYWORD) ? [ { value: DEACTIVATION_KEYWORD, label: DEACTIVATION_KEYWORD, description: "Deactivate persona", }, ] : []; return [...agentItems, ...offItems]; }, /** * Lists, activates, or deactivates the agent selected by `/mainagent`. * * Inputs: command arguments and extension context. Output: none. * Side effects: changes the active persona, persists the selection, updates * the widget, and notifies the TUI. */ handler: async (args: string, ctx: ExtensionContext) => { // Normalize the optional command argument before branching on intent. const arg = args.trim(); // No argument: interactive menu in the TUI, static listing elsewhere. if (!arg) { if (agents.length === 0) { ctx.ui.notify( `No agents found. Create files in ${agentsDir()}${path.sep}*.md`, "info", ); return; } if (ctx.mode === "tui") { await runAgentMenu(ctx); return; } ctx.ui.notify(renderAgentList(agents, active), "info"); return; } // Resolve the requested agent definition by name; real agents win over // the deactivation aliases. const agent = findAgent(arg); if (!agent && arg.toLowerCase() === DEACTIVATION_KEYWORD) { await applySelection(null, ctx); return; } if (!agent) { ctx.ui.notify( `Agent "${arg}" not found. Available: ${agents.map((a) => a.name).join(", ")}`, "error", ); return; } // Activate the persona and apply its runtime settings. await applySelection(agent, ctx); }, }); } /** Registers every pi-mainagent capability from a single extension entry point. */ export default function piMainagentExtension(pi: ExtensionAPI): void { registerMainAgentFeatures(pi); }