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, readFile } from "node:fs/promises";
import { dirname, join } from "node:path";
import { homedir } from "node:os";
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,
	PHASE_GLYPH,
} from "./phases.js";
import type { AgentSessionEvent } from "@earendil-works/pi-coding-agent";

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

/** Custom message type for the live tool-event stream (mirrors piolium-stream). */
export const PYGIENIUM_STREAM = "pygienium-stream";

/** Chat rendering style ("verbose" = per-event stream, "compact" = completion-only). */
export type ChatStyle = "verbose" | "compact";

/**
 * Read the pygienium chat style from pi's settings.json.
 * Looks for `pygienium.chatStyle` under `~/.pi/agent/settings.json`.
 * Defaults to "verbose" (piolium-style per-event stream) when absent or unreadable.
 */
async function readChatStyle(): Promise<ChatStyle> {
	try {
		const raw = await readFile(
			join(homedir(), ".pi", "agent", "settings.json"),
			"utf8",
		);
		const settings = JSON.parse(raw) as {
			pygienium?: { chatStyle?: string };
		};
		const style = settings.pygienium?.chatStyle;
		return style === "compact" ? "compact" : "verbose";
	} catch {
		return "verbose";
	}
}

/**
 * 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";

type StreamLineKind = "tool-start" | "tool-end" | "tool-error" | "assistant";

interface StreamLineDetails {
	kind: StreamLineKind;
	phase: string;
	toolName?: string;
	body?: string;
}

/** Pick the one useful argument from a tool-call's args (path/command/…). */
function summarizeArgs(args: unknown): string {
	if (!args || typeof args !== "object") return "";
	const obj = args as Record<string, unknown>;
	const pickKey = [
		"file_path",
		"path",
		"command",
		"pattern",
		"query",
		"url",
	].find((k) => typeof obj[k] === "string");
	if (pickKey) {
		const value = String(obj[pickKey]);
		return value.length > 120 ? `${value.slice(0, 117)}…` : value;
	}
	const json = JSON.stringify(obj);
	return json.length > 120 ? `${json.slice(0, 117)}…` : json;
}

/** Extract joined text from an assistant message's content blocks. */
function extractAssistantText(content: unknown): string {
	if (typeof content === "string") return content;
	if (!Array.isArray(content)) return "";
	return content
		.filter(
			(c) =>
				c && typeof c === "object" && (c as { type?: string }).type === "text",
		)
		.map((c) => (c as { text?: string }).text ?? "")
		.join("");
}

/** Collapse a tool result down to a single short line. */
function summarizeToolResult(result: unknown): string {
	if (result == null) return "";
	if (typeof result === "string") return result;
	if (typeof result === "number" || typeof result === "boolean")
		return String(result);
	if (Array.isArray(result)) {
		return result
			.map((item) => {
				if (typeof item === "string") return item;
				if (
					item &&
					typeof item === "object" &&
					"text" in (item as Record<string, unknown>)
				)
					return String((item as { text?: unknown }).text ?? "");
				return JSON.stringify(item);
			})
			.join("\n");
	}
	if (typeof result !== "object") return "";
	const obj = result as Record<string, unknown>;
	// MCP CallToolResult shape: { content: [{ type: "text", text: "..." }, ...] }
	if (Array.isArray(obj.content)) {
		const unwrapped = summarizeToolResult(obj.content);
		if (unwrapped) return unwrapped;
	}
	const preferKey = ["stdout", "output", "text", "content", "result"].find(
		(k) => typeof obj[k] === "string" && (obj[k] as string).length > 0,
	);
	if (preferKey) return obj[preferKey] as string;
	try {
		return JSON.stringify(obj);
	} catch {
		return "";
	}
}

/** Collapse whitespace and cap a line at `max` chars with an ellipsis. */
function compactLine(text: string, max: number): string {
	const collapsed = text.replace(/\s+/g, " ").trim();
	if (collapsed.length <= max) return collapsed;
	return `${collapsed.slice(0, max - 1)}…`;
}

/**
 * Chat stream forwarder: turns raw sub-agent events into `pygienium-stream`
 * messages (one chat line per start/end/assistant turn), the pygienium
 * analogue of piolium's `makeAgentEventForwarder`. Also exposes
 * `sendPhaseLine` for synthetic progress lines during non-agent phases
 * (verify, cleanup, recon) so the chat doesn't go silent.
 */
interface StreamForwarder {
	/** Forward a raw sub-agent event tagged with a phase label. */
	onAgentEvent(phase: string, event: AgentSessionEvent): void;
	/** Emit a synthetic progress line (e.g. "checking artifacts…"). */
	sendPhaseLine(phase: string, text: string): void;
}

function makeStreamForwarder(pi: ExtensionAPI): StreamForwarder {
	const send = (details: StreamLineDetails, fallback: string) => {
		pi.sendMessage<StreamLineDetails>({
			customType: PYGIENIUM_STREAM,
			content: fallback,
			display: true,
			details,
		});
	};

	const onAgentEvent = (phase: string, event: AgentSessionEvent): void => {
		switch (event.type) {
			case "tool_execution_start": {
				const body = summarizeArgs(event.args);
				send(
					{ kind: "tool-start", phase, toolName: event.toolName, body },
					`[${phase}] → ${event.toolName}${body ? `  ${body}` : ""}`,
				);
				return;
			}
			case "tool_execution_end": {
				const body = compactLine(summarizeToolResult(event.result), 200);
				const kind: StreamLineKind = event.isError ? "tool-error" : "tool-end";
				const marker = event.isError ? "✗" : "←";
				send(
					{ kind, phase, toolName: event.toolName, body },
					`[${phase}] ${marker} ${event.toolName}${body ? `  ${body}` : ""}`,
				);
				return;
			}
			case "message_end": {
				const message = event.message as {
					role?: string;
					content?: unknown;
				};
				if (message.role !== "assistant") return;
				const text = extractAssistantText(message.content).trim();
				if (!text) return;
				const head = compactLine(text, 240);
				send({ kind: "assistant", phase, body: head }, `[${phase}] ${head}`);
				return;
			}
		}
	};

	const sendPhaseLine = (phase: string, text: string): void => {
		send({ kind: "assistant", phase, body: text }, `[${phase}] ${text}`);
	};

	return { onAgentEvent, sendPhaseLine };
}

/**
 * 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}`);
	}
}

/**
 * Create a callback to send completion 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;
					error?: string;
			  }
			| undefined;

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

		const completion = details?.completion;
		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 (!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,
	);

	// Live tool-event stream renderer: one chat line per tool start/end and
	// assistant turn, indented so ends nest under their start. Mirrors
	// piolium's PIOLIUM_STREAM renderer.
	pi.registerMessageRenderer<StreamLineDetails>(
		PYGIENIUM_STREAM,
		(message, _options, theme) => {
			const details = message.details;
			if (!details || typeof details !== "object") {
				const fallback =
					typeof message.content === "string" ? message.content : "";
				return new Text(theme.fg("muted", fallback), 0, 0);
			}
			const { kind, phase, toolName, body } = details;
			// Indent end/error lines so they visually nest under the matching
			// start line. The pad width matches the "[phase] " prefix.
			const phaseTag = theme.fg("accent", `[${phase}]`);
			const indent = " ".repeat(phase.length + 3);
			let line: string;
			switch (kind) {
				case "tool-start": {
					const arrow = theme.fg("muted", "→");
					const name = theme.fg("toolTitle", theme.bold(toolName ?? ""));
					const args = body ? ` ${theme.fg("muted", body)}` : "";
					line = `${phaseTag} ${arrow} ${name}${args}`;
					break;
				}
				case "tool-end": {
					const arrow = theme.fg("success", "←");
					const result = body
						? ` ${theme.fg("dim", body)}`
						: ` ${theme.fg("dim", "(ok)")}`;
					line = `${indent}${arrow}${result}`;
					break;
				}
				case "tool-error": {
					const marker = theme.fg("error", "✗");
					const result = body
						? ` ${theme.fg("error", body)}`
						: ` ${theme.fg("error", "failed")}`;
					line = `${indent}${marker}${result}`;
					break;
				}
				case "assistant":
					line = `${phaseTag} ${theme.fg("muted", body ?? "")}`;
					break;
				default:
					line =
						typeof message.content === "string"
							? theme.fg("muted", message.content)
							: "";
			}
			return new Text(line, 0, 0);
		},
	);

	const forwarder = makeStreamForwarder(pi);

	// Read chat style from pi's settings.json. When "compact", suppress the
	// per-event stream + synthetic phase lines so only the completion message
	// (with its expandable phase tree) shows — the ralpi-style rendering.
	const chatStyle = await readChatStyle();
	const verbose = chatStyle === "verbose";
	const onAgentEvent = verbose ? forwarder.onAgentEvent : undefined;
	const sendPhaseLine = verbose ? forwarder.sendPhaseLine : undefined;

	registerPygieniumCommands((name, options) => {
		pi.registerCommand(name, {
			description: options.description,
			handler: (args: string, ctx: ExtensionCommandContext) => {
				// Create PygieniumCtx with chat + stream callbacks.
				const pygieniumCtx: PygieniumCtx = {
					cwd: ctx.cwd,
					mode: ctx.mode,
					hasUI: ctx.hasUI,
					ui: ctx.ui,
					sendChatMessage,
					onAgentEvent,
					sendPhaseLine,
				};
				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");
		},
	);
}