@xynogen/pix-subagent 0.1.0

package/src/index.ts

@@ -0,0 +1,216 @@
+/**
+ * index.ts — pix-subagent extension entry point.
+ *
+ * Registers 3 LLM-callable tools (agent, agent_result, agent_steer),
+ * a live above-editor widget (model shown inline), and a themed
+ * subagent-notification renderer.
+ *
+ * Best-of-both:
+ *   - tintinweb/pi-subagents spawn engine (MIT) — battle-tested createAgentSession path
+ *   - nicobailon/pi-subagents explicit work-splitting (allowed_tools[], model param)
+ *   - pix twist: model name ALWAYS visible in widget + notification
+ */
+
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import { AgentManager } from "./agent-manager.ts";
+import { registerAgents } from "./agent-types.ts";
+import { loadCustomAgents } from "./custom-agents.ts";
+import { listAvailable } from "./model-resolver.ts";
+import {
+	type AgentActivity,
+	createAgentResultTool,
+	createAgentSteerTool,
+	createAgentTool,
+} from "./tools.ts";
+import type { NotificationDetails } from "./types.ts";
+import { registerNotificationRenderer } from "./ui/notification.ts";
+import { AgentWidget } from "./ui/widget.ts";
+import { getLifetimeTotal } from "./usage.ts";
+
+const EXTENSION_KEY = "pix-subagent";
+
+// Reload guard key — stored on globalThis so a dev-reload cleans up stale state
+const CLEANUP_KEY = `__${EXTENSION_KEY}Cleanup`;
+
+export default function registerPixSubagent(pi: ExtensionAPI): void {
+	// ── Cleanup stale timers from a prior reload ───────────────────────────────
+	const g = globalThis as Record<string, unknown>;
+	const prevCleanup = g[CLEANUP_KEY];
+	if (typeof prevCleanup === "function") {
+		try {
+			(prevCleanup as () => void)();
+		} catch {
+			/* best effort */
+		}
+	}
+
+	// ── Agent registry ─────────────────────────────────────────────────────────
+	const reloadCustomAgents = () => {
+		const userAgents = loadCustomAgents(process.cwd());
+		registerAgents(userAgents);
+	};
+	reloadCustomAgents();
+
+	// ── State ──────────────────────────────────────────────────────────────────
+	const agentActivity = new Map<string, AgentActivity>();
+
+	// Debounce: brief hold so agent_result can cancel a notification it just consumed
+	const pendingNudges = new Map<string, ReturnType<typeof setTimeout>>();
+	const NUDGE_HOLD_MS = 200;
+
+	function scheduleNudge(key: string, send: () => void) {
+		cancelNudge(key);
+		pendingNudges.set(
+			key,
+			setTimeout(() => {
+				pendingNudges.delete(key);
+				try {
+					send();
+				} catch {
+					/* ignore stale context errors */
+				}
+			}, NUDGE_HOLD_MS),
+		);
+	}
+
+	function cancelNudge(key: string) {
+		const t = pendingNudges.get(key);
+		if (t != null) {
+			clearTimeout(t);
+			pendingNudges.delete(key);
+		}
+	}
+
+	// ── AgentManager ──────────────────────────────────────────────────────────
+	const manager = new AgentManager(
+		// onComplete — fire subagent-notification nudge for each finished bg agent
+		(record) => {
+			agentActivity.delete(record.id);
+			widget.markFinished(record.id);
+
+			if (record.resultConsumed) {
+				widget.update();
+				return;
+			}
+
+			const totalTokens = getLifetimeTotal(record.lifetimeUsage);
+			const activity = agentActivity.get(record.id);
+			const resultPreview = record.result
+				? record.result.length > 500
+					? `${record.result.slice(0, 500)}…`
+					: record.result
+				: "No output.";
+
+			const details: NotificationDetails = {
+				id: record.id,
+				description: record.description,
+				status: record.status,
+				modelName: record.invocation?.modelName,
+				toolUses: record.toolUses,
+				turnCount: activity?.turnCount ?? 0,
+				maxTurns: activity?.maxTurns,
+				totalTokens,
+				durationMs: record.completedAt
+					? record.completedAt - record.startedAt
+					: 0,
+				error: record.error,
+				resultPreview,
+			};
+
+			scheduleNudge(record.id, () => {
+				if (record.resultConsumed) {
+					widget.update();
+					return;
+				}
+				pi.sendMessage<NotificationDetails>(
+					{
+						customType: "subagent-notification",
+						content: `Agent "${record.description}" ${record.status}.`,
+						display: true,
+						details,
+					},
+					{ deliverAs: "followUp", triggerTurn: true },
+				);
+			});
+
+			widget.update();
+		},
+		4, // maxConcurrent
+		// onStart
+		(_record) => {
+			widget.update();
+		},
+	);
+
+	// ── Widget ─────────────────────────────────────────────────────────────────
+	const widget = new AgentWidget(manager, agentActivity);
+
+	// ── Register renderers ─────────────────────────────────────────────────────
+	registerNotificationRenderer(pi);
+
+	// ── Build initial tool description with model list ─────────────────────────
+	// Model list is empty until session_start provides a modelRegistry.
+	// Tools are registered once; the description is rebuilt on session_start via
+	// re-registering (pi.registerTool replaces by name — verify this at smoke test).
+	// ponytail: if re-register isn't supported mid-session, the list stays empty
+	// first session; acceptable until we find a hook to refresh.
+	let currentModelList: string[] = [];
+
+	function registerTools() {
+		pi.registerTool(
+			createAgentTool(
+				pi as Parameters<typeof manager.spawnAndWait>[0],
+				manager,
+				agentActivity,
+				reloadCustomAgents,
+				currentModelList,
+			),
+		);
+		pi.registerTool(createAgentResultTool(manager, agentActivity));
+		pi.registerTool(createAgentSteerTool(manager));
+	}
+
+	registerTools();
+
+	// ── Lifecycle ──────────────────────────────────────────────────────────────
+	pi.on("session_start", (_event, ctx) => {
+		manager.clearCompleted();
+		agentActivity.clear();
+
+		// Refresh model list from live registry
+		const newList = listAvailable(ctx.modelRegistry);
+		if (newList.join(",") !== currentModelList.join(",")) {
+			currentModelList = newList;
+			// Re-register tools with fresh description
+			// ponytail: if pi throws on duplicate tool names, skip the re-register
+			try {
+				registerTools();
+			} catch {
+				/* description may be stale; non-fatal */
+			}
+		}
+	});
+
+	pi.on("tool_execution_start", (_event, ctx) => {
+		widget.setUICtx(ctx.ui as Parameters<typeof widget.setUICtx>[0]);
+		widget.onTurnStart();
+		widget.ensureTimer();
+	});
+
+	pi.on("session_shutdown", () => {
+		for (const t of pendingNudges.values()) clearTimeout(t);
+		pendingNudges.clear();
+		manager.abortAll();
+		manager.dispose();
+		widget.dispose();
+		agentActivity.clear();
+		if (g[CLEANUP_KEY] === runtimeCleanup) delete g[CLEANUP_KEY];
+	});
+
+	const runtimeCleanup = () => {
+		for (const t of pendingNudges.values()) clearTimeout(t);
+		pendingNudges.clear();
+		widget.dispose();
+	};
+	g[CLEANUP_KEY] = runtimeCleanup;
+}

package/src/invocation-config.ts

@@ -0,0 +1,43 @@
+/**
+ * invocation-config.ts — Resolve per-call agent invocation options.
+ *
+ * Ported from tintinweb/pi-subagents (MIT). Trimmed: dropped isolation/joinMode.
+ */
+
+import type { AgentConfig, ThinkingLevel } from "./types.ts";
+
+interface AgentInvocationParams {
+	model?: string;
+	thinking?: string;
+	max_turns?: number;
+	run_in_background?: boolean;
+	inherit_context?: boolean;
+	isolated?: boolean;
+}
+
+export function resolveAgentInvocationConfig(
+	agentConfig: AgentConfig | undefined,
+	params: AgentInvocationParams,
+): {
+	modelInput?: string;
+	modelFromParams: boolean;
+	thinking?: ThinkingLevel;
+	maxTurns?: number;
+	inheritContext: boolean;
+	runInBackground: boolean;
+	isolated: boolean;
+} {
+	return {
+		modelInput: agentConfig?.model ?? params.model,
+		modelFromParams: agentConfig?.model == null && params.model != null,
+		thinking: (agentConfig?.thinking ?? params.thinking) as
+			| ThinkingLevel
+			| undefined,
+		maxTurns: agentConfig?.maxTurns ?? params.max_turns,
+		inheritContext:
+			agentConfig?.inheritContext ?? params.inherit_context ?? false,
+		runInBackground:
+			agentConfig?.runInBackground ?? params.run_in_background ?? false,
+		isolated: agentConfig?.isolated ?? params.isolated ?? false,
+	};
+}

package/src/model-resolver.ts

@@ -0,0 +1,98 @@
+/**
+ * Model resolution: exact match ("provider/modelId") with fuzzy fallback.
+ */
+
+export interface ModelEntry {
+	id: string;
+	name: string;
+	provider: string;
+}
+
+export interface ModelRegistry {
+	find(provider: string, modelId: string): any;
+	getAll(): any[];
+	getAvailable?(): any[];
+}
+
+/**
+ * Resolve a model string to a Model instance.
+ * Tries exact match first ("provider/modelId"), then fuzzy match against all available models.
+ * Returns the Model on success, or an error message string on failure.
+ */
+export function resolveModel(
+	input: string,
+	registry: ModelRegistry,
+): any | string {
+	// Available models (those with auth configured)
+	const all = (registry.getAvailable?.() ?? registry.getAll()) as ModelEntry[];
+	const availableSet = new Set(
+		all.map((m) => `${m.provider}/${m.id}`.toLowerCase()),
+	);
+
+	// 1. Exact match: "provider/modelId" — only if available (has auth)
+	const slashIdx = input.indexOf("/");
+	if (slashIdx !== -1) {
+		const provider = input.slice(0, slashIdx);
+		const modelId = input.slice(slashIdx + 1);
+		if (availableSet.has(input.toLowerCase())) {
+			const found = registry.find(provider, modelId);
+			if (found) return found;
+		}
+	}
+
+	// 2. Fuzzy match against available models
+	const query = input.toLowerCase();
+
+	// Score each model: prefer exact id match > id contains > name contains > provider+id contains
+	let bestMatch: ModelEntry | undefined;
+	let bestScore = 0;
+
+	for (const m of all) {
+		const id = m.id.toLowerCase();
+		const name = m.name.toLowerCase();
+		const full = `${m.provider}/${m.id}`.toLowerCase();
+
+		let score = 0;
+		if (id === query || full === query) {
+			score = 100; // exact
+		} else if (id.includes(query) || full.includes(query)) {
+			score = 60 + (query.length / id.length) * 30; // substring, prefer tighter matches
+		} else if (name.includes(query)) {
+			score = 40 + (query.length / name.length) * 20;
+		} else if (
+			query
+				.split(/[\s\-/]+/)
+				.every(
+					(part) =>
+						id.includes(part) ||
+						name.includes(part) ||
+						m.provider.toLowerCase().includes(part),
+				)
+		) {
+			score = 20; // all parts present somewhere
+		}
+
+		if (score > bestScore) {
+			bestScore = score;
+			bestMatch = m;
+		}
+	}
+
+	if (bestMatch && bestScore >= 20) {
+		const found = registry.find(bestMatch.provider, bestMatch.id);
+		if (found) return found;
+	}
+
+	// 3. No match — list available models
+	const modelList = all
+		.map((m) => `  ${m.provider}/${m.id}`)
+		.sort()
+		.join("\n");
+	return `Model not found: "${input}".\n\nAvailable models:\n${modelList}`;
+}
+
+/** List available models as "provider/id" strings (for tool-description injection). */
+export function listAvailable(registry: ModelRegistry): string[] {
+	const all = (registry.getAvailable?.() ?? registry.getAll()) as ModelEntry[];
+	return all.map((m) => `${m.provider}/${m.id}`).sort();
+}

package/src/once.ts

@@ -0,0 +1,26 @@
+/**
+ * Per-instance idempotency guard for extension activation.
+ *
+ * pix-core (the meta-package) invokes this package's factory, and a standalone
+ * install makes Pi invoke it again — sometimes against the SAME `pi`. We must
+ * dedupe that. But Pi rebuilds the extension runtime on /new, /resume, /fork,
+ * and /reload, handing the factory a BRAND-NEW `pi`; that must re-register.
+ *
+ * Keying the registry on the `pi` instance satisfies both: same instance =>
+ * skip, new instance => run. The registry lives on globalThis because jiti
+ * (`moduleCache: false`) re-evaluates this module on every load pass, so a
+ * module-scoped WeakMap would not be shared between the aggregator pass and the
+ * standalone pass within a single session.
+ */
+export function once(pi: object, key: string, fn: () => void): void {
+	const g = globalThis as { __pixOnce?: WeakMap<object, Set<string>> };
+	const registry = (g.__pixOnce ??= new WeakMap<object, Set<string>>());
+	let loaded = registry.get(pi);
+	if (!loaded) {
+		loaded = new Set<string>();
+		registry.set(pi, loaded);
+	}
+	if (loaded.has(key)) return;
+	loaded.add(key);
+	fn();
+}

package/src/prompts.ts

@@ -0,0 +1,111 @@
+/**
+ * prompts.ts — System prompt builder for agents.
+ */
+
+import type { AgentConfig, EnvInfo } from "./types.ts";
+
+/** Extra sections to inject into the system prompt (memory, skills, etc.). */
+export interface PromptExtras {
+	/** Persistent memory content to inject (first 200 lines of MEMORY.md + instructions). */
+	memoryBlock?: string;
+	/** Preloaded skill contents to inject. */
+	skillBlocks?: { name: string; content: string }[];
+}
+
+/**
+ * Build the system prompt for an agent from its config.
+ *
+ * - "replace" mode: env header + config.systemPrompt (full control, no parent identity)
+ * - "append" mode: parent system prompt + sub-agent context + env header + config.systemPrompt
+ * - "append" with empty systemPrompt: pure parent clone
+ *
+ * Both modes include an `<active_agent name="${config.name}"/>` tag so downstream
+ * extensions (e.g. permission/policy systems) can resolve per-agent policy
+ * inside the child session by parsing the system prompt. In replace mode the tag
+ * is prepended; in append mode it follows the shared inherited content so the
+ * parent prompt forms an identical, cacheable byte prefix with the parent
+ * session (the LLM's KV cache can then reuse those tokens across every spawn).
+ *
+ * @param parentSystemPrompt  The parent agent's effective system prompt (for append mode).
+ * @param extras  Optional extra sections to inject (memory, preloaded skills).
+ */
+export function buildAgentPrompt(
+	config: AgentConfig,
+	cwd: string,
+	env: EnvInfo,
+	parentSystemPrompt?: string,
+	extras?: PromptExtras,
+): string {
+	const activeAgentTag = `<active_agent name="${config.name}"/>\n\n`;
+
+	const envBlock = `# Environment
+Working directory: ${cwd}
+${env.isGitRepo ? `Git repository: yes\nBranch: ${env.branch}` : "Not a git repository"}
+Platform: ${env.platform}`;
+
+	// Build optional extras suffix
+	const extraSections: string[] = [];
+	if (extras?.memoryBlock) {
+		extraSections.push(extras.memoryBlock);
+	}
+	if (extras?.skillBlocks?.length) {
+		for (const skill of extras.skillBlocks) {
+			extraSections.push(
+				`\n# Preloaded Skill: ${skill.name}\n${skill.content}`,
+			);
+		}
+	}
+	const extrasSuffix =
+		extraSections.length > 0 ? `\n\n${extraSections.join("\n")}` : "";
+
+	if (config.promptMode === "append") {
+		const identity = parentSystemPrompt || genericBase;
+
+		const bridge = `<sub_agent_context>
+You are operating as a sub-agent invoked to handle a specific task.
+- Use the read tool instead of cat/head/tail
+- Use the edit tool instead of sed/awk
+- Use the write tool instead of echo/heredoc
+- Use the find tool instead of bash find/ls for file search
+- Use the grep tool instead of bash grep/rg for content search
+- Make independent tool calls in parallel
+- Use absolute file paths
+- Do not use emojis
+- Be concise but complete
+</sub_agent_context>`;
+
+		const customSection = config.systemPrompt?.trim()
+			? `\n\n<agent_instructions>\n${config.systemPrompt}\n</agent_instructions>`
+			: "";
+
+		// Place shared/stable content first so the LLM's KV cache can reuse the
+		// inherited prefix across all subagent invocations. The parent prompt is
+		// placed verbatim (no wrapper tag) so it forms an identical byte prefix
+		// with the parent session, maximising KV cache hits. The <active_agent>
+		// tag and env block vary per call and are placed after the cached prefix.
+		return (
+			identity +
+			"\n\n" +
+			bridge +
+			"\n\n" +
+			activeAgentTag +
+			envBlock +
+			customSection +
+			extrasSuffix
+		);
+	}
+
+	// "replace" mode — env header + the config's full system prompt
+	const replaceHeader = `You are a pi coding agent sub-agent.
+You have been invoked to handle a specific task autonomously.
+
+${envBlock}`;
+
+	return `${activeAgentTag + replaceHeader}\n\n${config.systemPrompt}${extrasSuffix}`;
+}
+
+/** Fallback base prompt when parent system prompt is unavailable in append mode. */
+const genericBase = `# Role
+You are a general-purpose coding agent for complex, multi-step tasks.
+You have full access to read, write, edit files, and execute commands.
+Do what has been asked; nothing more, nothing less.`;