pygienium/src/index.ts

/**
 * pygienium — code hygiene extension for pi.
 *
 * Entry point. Registers `/pygienium-help`, auto-registers one
 * `/pygienium-<check>` command per registered `CheckDefinition`, plus the
 * `all`/`resume`/`status`/`export` commands. Adding a check requires ONLY a new
 * file in `src/checks/` plus one `registerCheck(def)` call — no changes here.
 *
 * Check files in `src/checks/` are auto-discovered (every `.ts` except the
 * registry barrel), so they self-register at load time before commands bind.
 *
 * Pi loads this file via jiti at runtime (see `pi.extensions` in package.json).
 * The default export runs once per session; the factory is async so check
 * modules finish registering before command wiring.
 *
 * @module pygienium/index
 */

import { readdir } from "node:fs/promises";
import { dirname, join } from "node:path";
import { fileURLToPath } from "node:url";
import type {
	ExtensionAPI,
	ExtensionCommandContext,
	ExtensionContext,
	MessageRenderer,
	SessionStartEvent,
} from "@earendil-works/pi-coding-agent";
import { Box, Text } from "@earendil-works/pi-tui";
import { registerPygieniumCommands, type PygieniumCtx } from "./commands.js";
import {
	type SendChatMessage,
	type CheckCompletionDetails,
	type ToolCallEntry,
	PHASE_GLYPH,
} from "./phases.js";

/** Startup hint mirrored after piolium's convention. */
export const PYGIENIUM_STARTUP_HINT =
	"Pygienium loaded. Run /pygienium-help for available checks and flags.";

/**
 * Local structural supertypes for the progress-message renderer params.
 * These avoid relying on contextual typing from `MessageRenderer` (which
 * requires resolving pi's internal `Theme`/`CustomMessage` cross-references
 * via node_modules — not always available in dev environments). Using
 * `(...args: any[])` for theme methods makes the type bidirectionally
 * compatible under strict function types, so the cast to `MessageRenderer`
 * in `registerMessageRenderer` is valid.
 */
interface ProgressMessage {
	content: unknown;
	details?: unknown;
}
interface ProgressRenderOptions {
	expanded: boolean;
}
interface ProgressTheme {
	fg: (...args: any[]) => string;
	bg: (...args: any[]) => string;
}

export { buildPygieniumHelpLines } from "./help.js";

/**
 * Import every `checks/*.ts` module (except the registry barrel) so each check
 * file's top-level `registerCheck(def)` call runs before command binding. This
 * is what makes adding a check require zero index.ts changes — drop a file,
 * it self-registers.
 */
async function loadCheckModules(): Promise<void> {
	const dir = join(dirname(fileURLToPath(import.meta.url)), "checks");
	let entries: string[];
	try {
		entries = await readdir(dir);
	} catch {
		return; // no checks dir (e.g. minimal install)
	}
	for (const entry of entries) {
		if (!entry.endsWith(".ts")) continue;
		if (entry === "registry.ts" || entry === "load.ts") continue;
		await import(`./checks/${entry}`);
	}
}

/** Max tool calls shown collapsed in a chat message (matches ralpi). */
const RENDERER_MAX_COLLAPSED = 3;

/**
 * Create a callback to send messages to the main chat window.
 */
function makeSendChatMessage(pi: ExtensionAPI): SendChatMessage {
	return (content: string, meta?: Record<string, unknown>) => {
		pi.sendMessage({
			customType: "pygienium-progress",
			content,
			display: true,
			details: {
				phase: meta?.phase || "info",
				...meta,
			},
		});
	};
}

export default async function pygieniumExtension(
	pi: ExtensionAPI,
): Promise<void> {
	// Self-register every shipped check before wiring commands.
	await loadCheckModules();

	const sendChatMessage = makeSendChatMessage(pi);

	// Register custom message renderer for pygienium progress messages.
	// Renders an expandable phase tree: collapsed shows the header + a hint,
	// expanded (Ctrl+O) shows every phase with its status and notes.
	const progressRenderer = (
		message: ProgressMessage,
		{ expanded }: ProgressRenderOptions,
		theme: ProgressTheme,
	) => {
		const details = message.details as
			| {
					phase?: string;
					completion?: CheckCompletionDetails;
					toolCalls?: ToolCallEntry[];
					error?: string;
			  }
			| undefined;

		const lines: string[] = [];
		lines.push(String(message.content));

		const completion = details?.completion;
		const toolCalls = details?.toolCalls;
		if (completion) {
			if (expanded) {
				// Expanded: show every phase with status glyph + branch.
				const phases = completion.phases;
				for (let i = 0; i < phases.length; i++) {
					const entry = phases[i];
					if (!entry) continue;
					const isLast = i === phases.length - 1;
					const branch = isLast ? "  └── " : "  ├── ";
					const glyph = PHASE_GLYPH[entry.status] ?? "?";
					const tag = theme.fg("accent", entry.label);
					const note = entry.note ? ` · ${entry.note}` : "";
					lines.push(`${branch}${glyph} ${tag}${note}`);
				}
				if (completion.error) {
					lines.push(theme.fg("error", `  error: ${completion.error}`));
				}
			} else {
				// Collapsed: summary line + hint.
				const done = completion.phases.filter(
					(p) => p.status === "complete",
				).length;
				const total = completion.phases.length;
				const hint = completion.error
					? theme.fg("error", `  ├── ${completion.error}`)
					: theme.fg(
							"dim",
							`  ├── ${done}/${total} phases · press Ctrl+O for detail`,
						);
				lines.push(hint);
			}
		} else if (toolCalls && toolCalls.length > 0) {
			const all = toolCalls;
			if (expanded) {
				for (let i = 0; i < all.length; i++) {
					const entry = all[i]!;
					const isLast = i === all.length - 1;
					const branch = isLast ? "  └── " : "  ├── ";
					const tag = theme.fg("accent", `[${entry.name}]`);
					lines.push(`${branch}${tag} ${entry.label}`);
				}
			} else {
				const shown = all.slice(-RENDERER_MAX_COLLAPSED);
				const remaining = all.length - shown.length;
				if (remaining > 0) {
					lines.push(theme.fg("dim", `  ├── ${remaining} more`));
				}
				for (let i = 0; i < shown.length; i++) {
					const entry = shown[i]!;
					const isLast = i === shown.length - 1;
					const branch = isLast ? "  └── " : "  ├── ";
					const tag = theme.fg("accent", `[${entry.name}]`);
					lines.push(`${branch}${tag} ${entry.label}`);
				}
			}
			if (details?.error) {
				lines.push(theme.fg("error", `  error: ${details.error}`));
			}
		} else if (!expanded) {
			lines.push(theme.fg("dim", "  ├── press Ctrl+O for detail"));
		}

		const text = lines.join("\n");
		const box = new Box(1, 1, (t: string) => theme.bg("customMessageBg", t));
		box.addChild(new Text(text, 0, 0));
		return box;
	};
	pi.registerMessageRenderer(
		"pygienium-progress",
		progressRenderer as MessageRenderer,
	);

	registerPygieniumCommands((name, options) => {
		pi.registerCommand(name, {
			description: options.description,
			handler: (args: string, ctx: ExtensionCommandContext) => {
				// Create PygieniumCtx with sendChatMessage callback
				const pygieniumCtx: PygieniumCtx = {
					cwd: ctx.cwd,
					mode: ctx.mode,
					hasUI: ctx.hasUI,
					ui: ctx.ui,
					sendChatMessage,
				};
				return options.handler(args, pygieniumCtx);
			},
		});
	});

	pi.on(
		"session_start",
		async (_event: SessionStartEvent, ctx: ExtensionContext) => {
			if (!ctx.hasUI) return;
			ctx.ui.notify(PYGIENIUM_STARTUP_HINT, "info");
		},
	);
}