pygienium/src/agent-runner.ts

/**
 * agent-runner.ts — spawn isolated sub-agents for analysis and fix phases.
 *
 * The production runner uses pi's `createAgentSession` SDK to spin up a fresh
 * in-memory agent session scoped to the target `cwd`, with the agent
 * definition's system prompt and tool allowlist applied. Because that path
 * needs live model credentials (unsuitable for CI), the runner is backed by an
 * injectable factory: tests swap it for a deterministic fake that executes a
 * tiny instruction protocol embedded in the task string.
 *
 * Instruction protocol (used by the fake runner, harmless to the real one):
 * task lines may begin with `!write <path> <text>` — the fake writes the file
 * and reports it as a finding. Real sub-agents receive the whole task verbatim.
 *
 * @module pygienium/agent-runner
 */

import { mkdir, writeFile } from "node:fs/promises";
import { dirname, isAbsolute, join } from "node:path";
import type {
	AgentSession,
	AgentSessionEvent,
} from "@earendil-works/pi-coding-agent";
import { loadAgents, extensionRoot, type AgentDef } from "./agents.js";

/**
 * Env var overriding the per-agent settle timeout (ms). Guards against a
 * sub-agent session that never settles (stalled provider stream, hung retry /
 * auto-compaction after the final tool call), which otherwise leaves the run
 * stuck mid-phase with no state save, no error, and no completion message.
 */
export const AGENT_TIMEOUT_ENV = "PYGIENIUM_AGENT_TIMEOUT_MS";

/** Default settle timeout per agent phase: 60 minutes. */
const DEFAULT_AGENT_TIMEOUT_MS = 60 * 60_000;

/** Resolve the per-agent settle timeout, honouring the env override. */
export function agentTimeoutMs(): number {
	const raw = process.env[AGENT_TIMEOUT_ENV];
	if (raw && /^\d+$/.test(raw.trim()) && Number(raw.trim()) > 0) {
		return Number(raw.trim());
	}
	return DEFAULT_AGENT_TIMEOUT_MS;
}

export interface AgentTaskOptions {
	/** Absolute working directory for the sub-agent. */
	cwd: string;
	/** Agent name to look up in `agents/*.md`. */
	agentName: string;
	/** The task prompt handed to the sub-agent. */
	task: string;
	/** Optional tool allowlist override (else uses the agent's `allowedTools`). */
	allowedTools?: string[];
	/** Optional explicit agent definition (skips `loadAgents`). */
	agent?: AgentDef;
	/**
	 * Live callback forwarding raw {@link AgentSessionEvent}s from the
	 * sub-agent session. The `pygienium-stream` forwarder in `index.ts`
	 * turns tool_execution_start/end + assistant turns into chat messages.
	 */
	onEvent?: (event: AgentSessionEvent) => void;
	/**
	 * Maximum time the agent run may take before it is aborted and the phase
	 * fails loudly (default {@link agentTimeoutMs}). A session that never
	 * settles — stalled provider stream, hung retry/compaction after its last
	 * tool call — would otherwise hang the check run silently with no state
	 * update and no completion message.
	 */
	timeoutMs?: number;
}

/**
 * Race `promise` against a settle deadline. Returns `{ value }` on success or
 * `{ error }` when the deadline elapsed first (the caller aborts the work).
 * `promise` is still awaited-then-ignored afterwards so late rejections can
 * never surface as unhandled.
 */
export async function withSettleTimeout<T>(
	promise: Promise<T>,
	timeoutMs: number,
	label: string,
): Promise<{ value: T } | { error: string }> {
	if (timeoutMs <= 0) {
		try {
			return { value: await promise };
		} catch (err) {
			return { error: err instanceof Error ? err.message : String(err) };
		}
	}
	const settled = promise.then(
		(value) => ({ value }) as { value: T },
		(err) =>
			({ error: err instanceof Error ? err.message : String(err) }) as {
				error: string;
			},
	);
	let timer: ReturnType<typeof setTimeout> | undefined;
	const deadline = new Promise<{ error: string }>((resolve) => {
		timer = setTimeout(() => {
			// ~`1m 30s` / `45s` for the message (timeoutMs is ms; tests use small values).
			const totalSec = Math.round(timeoutMs / 1000);
			const duration =
				totalSec >= 60
					? `${Math.floor(totalSec / 60)}m${totalSec % 60 ? ` ${totalSec % 60}s` : ""}`
					: `${totalSec}s`;
			resolve({
				error: `${label} did not settle within ${duration}; aborted. Check model/provider connectivity, then resume with /pygienium-resume.`,
			});
		}, timeoutMs);
		// Never hold the process open just because a deadline is pending.
		timer.unref?.();
	});
	try {
		return await Promise.race([settled, deadline]);
	} finally {
		if (timer) clearTimeout(timer);
	}
}

export interface AgentRunResult {
	/** Whether the sub-agent completed without throwing. */
	ok: boolean;
	/** The final assistant text emitted by the sub-agent. */
	text: string;
	/** Error message when `ok` is false. */
	error?: string;
}

export type AgentRunner = (opts: AgentTaskOptions) => Promise<AgentRunResult>;

/** Module-level runner (defaults to the SDK-backed runner; tests override it). */
let currentRunner: AgentRunner = defaultAgentRunner;

/** Entry point used by the check-runner. */
export async function runAgentTask(
	opts: AgentTaskOptions,
): Promise<AgentRunResult> {
	// Backstop: even a third-party/custom runner must not be able to hang the
	// run forever. The default runner additionally aborts its session on
	// timeout (see defaultAgentRunner); this race covers every other runner.
	const settled = await withSettleTimeout(
		currentRunner(opts),
		opts.timeoutMs ?? agentTimeoutMs(),
		`sub-agent "${opts.agentName}"`,
	);
	if ("error" in settled) {
		return { ok: false, text: "", error: settled.error };
	}
	return settled.value;
}

/** Override the active agent runner (primarily for tests). */
export function setAgentRunner(runner: AgentRunner): void {
	currentRunner = runner;
}

/** Restore the default SDK-backed agent runner. */
export function resetAgentRunner(): void {
	currentRunner = defaultAgentRunner;
}

/**
 * Real sub-agent runner: spins up an in-memory `AgentSession` scoped to `cwd`,
 * overrides the system prompt with the agent definition's body, restricts tools
 * to the agent's allowlist, and runs the task to completion.
 */
export async function defaultAgentRunner(
	opts: AgentTaskOptions,
): Promise<AgentRunResult> {
	const agents = await loadAgents({ cwd: opts.cwd });
	const agent = opts.agent ?? agents.get(opts.agentName);
	if (!agent) {
		const names = [...agents.keys()];
		return {
			ok: false,
			text: "",
			error:
				`Unknown agent definition: "${opts.agentName}". ` +
				`Available agents: ${names.length > 0 ? names.join(", ") : "(none loaded — check agents/ directories exist)"}. ` +
				`Searched ${extensionRoot()}/agents/ (extension) and ${opts.cwd}/agents/ (project-local). ` +
				`Add agents/${opts.agentName}.md to either location.`,
		};
	}

	// Lazily import the SDK so the rest of the module graph (and tests using the
	// fake runner) never resolve the heavy pi-coding-agent package.
	const {
		createAgentSession,
		DefaultResourceLoader,
		getAgentDir,
		SessionManager,
	} = await import("@earendil-works/pi-coding-agent");
	const loader = new DefaultResourceLoader({
		cwd: opts.cwd,
		agentDir: getAgentDir(),
		systemPromptOverride: () => agent.systemPrompt,
		// Keep the sub-agent isolated: no nested extensions/skills/themes/etc.
		noExtensions: true,
		noSkills: true,
		noThemes: true,
		noPromptTemplates: true,
	});
	await loader.reload();

	const tools = opts.allowedTools ??
		agent.allowedTools ?? ["read", "bash", "grep", "find"];

	const { session } = await createAgentSession({
		cwd: opts.cwd,
		tools,
		sessionManager: SessionManager.inMemory(opts.cwd),
		resourceLoader: loader,
	});

	// A session that never settles (stalled provider stream, hung retry /
	// auto-compaction after its last tool call) must fail the phase loudly
	// instead of hanging the run mid-transition with zero diagnostics. On
	// timeout the session is disposed, which aborts the in-flight run.
	const settled = await withSettleTimeout(
		runSessionToCompletion(session, opts),
		opts.timeoutMs ?? agentTimeoutMs(),
		`sub-agent "${opts.agentName}"`,
	);
	if ("error" in settled) {
		// Abort the still-running session so it can't keep burning provider
		// calls; the background settle path then finishes and disposes too.
		try {
			session.dispose();
		} catch {
			/* ignore dispose errors */
		}
		return { ok: false, text: "", error: settled.error };
	}
	return settled.value;
}

/**
 * Run an in-memory agent session to completion, streaming tool events to the
 * chat forwarder, and return the agent's final text. Owns session cleanup.
 */
async function runSessionToCompletion(
	session: AgentSession,
	opts: AgentTaskOptions,
): Promise<AgentRunResult> {
	const acc: SessionEventAccumulator = { text: "" };
	try {
		const unsubscribe = session.subscribe((event: AgentSessionEvent) => {
			applySessionEvent(acc, event, opts.onEvent);
		});
		await session.prompt(opts.task, { expandPromptTemplates: false });
		// Ensure the agent has fully settled (tool calls may still be in-flight
		// after prompt() resolves; piolium's runner calls this for the same
		// reason).
		await session.agent.waitForIdle();
		unsubscribe();
		// Surface session errors that didn't throw but left no useful output.
		// A session ending with stopReason "error" and no text means the model
		// call failed silently — treat that as a failed run, not ok:true. A
		// session that settled with NO text and NO observed events means the
		// model never produced anything at all (dead provider stream, failed
		// start) — reporting that as a successful analysis would let an empty
		// scan masquerade as a clean one, and the check's verify hook would
		// fail only later with a confusing "artifact missing" error.
		const sessionError = emptySessionError(acc);
		if (sessionError) {
			return { ok: false, text: acc.text, error: sessionError };
		}
		return { ok: true, text: acc.text };
	} catch (err) {
		return {
			ok: false,
			text: "",
			error: err instanceof Error ? err.message : String(err),
		};
	} finally {
		try {
			session.dispose();
		} catch {
			/* ignore dispose errors */
		}
	}
}

/** Running capture state while a session streams assistant output. */
export interface SessionEventAccumulator {
	/** Joined assistant text seen so far (text_delta stream). */
	text: string;
	/** Whether any assistant-message or tool event was observed at all. */
	sawMessage: boolean;
	/** stopReason of the final assistant message, when reported. */
	stopReason?: string;
	/** errorMessage of the final assistant message, when reported. */
	errorMessage?: string;
}

/**
 * Classify a settled session's capture: `undefined` when the result is a
 * legitimate (possibly empty-text) outcome, else the error that should fail
 * the run. Kept pure so the decision is unit-testable without a session.
 */
export function emptySessionError(
	acc: SessionEventAccumulator,
): string | undefined {
	if (acc.errorMessage) return acc.errorMessage;
	if (!acc.text.trim()) {
		if (acc.stopReason === "error") {
			return "sub-agent session ended in error with no output.";
		}
		if (!acc.sawMessage) {
			return "sub-agent session settled with no output — no assistant message or tool activity was observed. Check model/provider connectivity, then resume.";
		}
	}
	return undefined;
}

/**
 * Interpret one session event into the running accumulator and forward the
 * stream-driving events to the chat.
 *
 * MUST never throw: it runs synchronously inside the SDK's event pipeline
 * (`session.subscribe` listeners are invoked from `processEvents`/`_emit`).
 * A throw there is not contained — the SDK's run-failure path re-emits
 * failure events through the same callback, so a callback that always throws
 * recurses until stack overflow and crashes the host, stranding run-state at
 * the phase boundary with no save and no completion. Guard every shape (the
 * runtime events are partial/streaming and weaker than their types) and never
 * let the cosmetic chat forwarder take the run down.
 */
export function applySessionEvent(
	acc: SessionEventAccumulator,
	event: AgentSessionEvent,
	forward?: (event: AgentSessionEvent) => void,
): void {
	if (!event) return;
	try {
		// Any message or tool event means the session actually ran — a settled
		// capture with none of these is a dead session, not an empty scan.
		if (
			event.type === "message_update" ||
			event.type === "message_end" ||
			event.type === "tool_execution_start" ||
			event.type === "tool_execution_end"
		) {
			acc.sawMessage = true;
		}
		if (
			event.type === "message_update" &&
			event.assistantMessageEvent?.type === "text_delta"
		) {
			acc.text += event.assistantMessageEvent.delta ?? "";
		}
		if (event.type === "message_end") {
			// Capture the full assistant text from the finalized message —
			// models that don't stream text_delta (or truncate) still surface
			// their output here. Prefer the streamed text when non-empty.
			const message = event.message as
				| {
						role?: string;
						content?: unknown;
						stopReason?: string;
						errorMessage?: string;
				  }
				| undefined;
			if (message) {
				if (message.stopReason) acc.stopReason = message.stopReason;
				if (message.errorMessage) acc.errorMessage = message.errorMessage;
				if (message.role === "assistant") {
					const full = extractAssistantText(message.content).trim();
					if (full && !acc.text.trim()) acc.text = full;
				}
			}
		}
	} catch (err) {
		// A malformed event must not crash the SDK event pipeline; log and skip.
		console.error("pygienium: error processing sub-agent event", err);
	}
	// Forward the stream-driving events to the chat forwarder; it turns each
	// into its own `pygienium-stream` message (see index.ts). Cosmetic chat
	// UI: a broken forwarder must not fail or hang the agent run.
	if (
		event.type === "tool_execution_start" ||
		event.type === "tool_execution_end" ||
		event.type === "message_end"
	) {
		try {
			forward?.(event);
		} catch (err) {
			console.error("pygienium: error forwarding sub-agent event", err);
		}
	}
}

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

/**
 * Fake agent runner for tests: it understands a tiny instruction protocol
 * embedded in the task so a no-op check can produce deterministic findings
 * and write marker files without a model. Recognised instructions (one per
 * line, leading-whitespace tolerant):
 *
 *   !write <path> <text...>   — write text to path (relative to cwd); recorded
 *   !echo <text...>            — appended to findings
 *
 * The agent's emitted findings text is the collected `!echo`/`!write` lines.
 */
export const fakeAgentRunner: AgentRunner = async (opts) => {
	const lines = opts.task.split(/\r?\n/);
	const findings: string[] = [];
	try {
		for (const line of lines) {
			const trimmed = line.trim();
			const write = /^!write\s+(\S+)\s*(.*)$/.exec(trimmed);
			if (write) {
				const rel = write[1] as string;
				const content = (write[2] ?? "").replace(/^["']|["']$/g, "");
				const full = isAbsolute(rel) ? rel : join(opts.cwd, rel);
				await mkdir(dirname(full), { recursive: true });
				await writeFile(full, content + "\n", "utf8");
				findings.push(`wrote ${rel}`);
				continue;
			}
			const echo = /^!echo\s+(.*)$/.exec(trimmed);
			if (echo) {
				findings.push((echo[1] ?? "").replace(/^["']|["']$/g, ""));
			}
		}
		return { ok: true, text: findings.join("\n") };
	} catch (err) {
		return {
			ok: false,
			text: findings.join("\n"),
			error: err instanceof Error ? err.message : String(err),
		};
	}
};