learnharness/src/agent/subagents/runner.ts

import { isToolUIPart, type LanguageModel, type UIMessage } from "ai";
import {
	buildSubagentSteer,
	type DemoResult,
	type DemoToolInput,
	type JobKind,
	type JobRecord,
	type ResearchResult,
	type ResearchToolInput,
	type SubagentSteerMessage,
	type VisualResult,
	type VisualToolInput,
} from "../../core/contracts.ts";
import type { ChatStore, JobStore } from "../../core/ports.ts";
import { runDemoBuilder } from "./demo-builder.ts";
import { createResearcherAgent } from "./researcher.ts";
import { runVisualMaker } from "./visual-maker.ts";

/**
 * The background job runner (contracts.ts "THE ASYNC PROTOCOL" §2–3), created
 * via createJobRunner({ chatStore, jobStore, deliver }) — the composition root
 * (chat-handler) injects the storage ports and the WAKE SEAM (the server-side
 * steer delivery, src/server/jobs.ts); this module never imports an adapter.
 *
 *   runJob(jobId) — runs the kind's sub-agent loop in the background
 *     (the job tools' execute fires this via the spawn seam, fire-and-forget),
 *     updating the job's activity line on every iteration. On completion the
 *     runner builds the STEER MESSAGE (buildSubagentSteer) and hands it to
 *     `deliver`; the wake pipeline fires the server-side teacher turn and
 *     flips the job to done/error when the lesson has actually resumed.
 *
 *   On failure the result is the failure shape (empty html/svg + "FAILED: …"
 *   notes; research: failure brief) — and it STILL steers, because the
 *   teacher must learn the outcome. If delivery itself fails, the job is
 *     marked error (the learner's activity card shows it).
 */

/**
 * Test seam: the spawn wiring has no injection point at the tool layer, so
 * tests swap in a MockLanguageModelV4 here to keep job tests off the network.
 * Undefined in production. (The jobs spawn through the module-level KIND
 * runners, which read these at call time.)
 */
export const researchModelOverride: { current?: LanguageModel } = {};
export const demoModelOverride: { current?: LanguageModel } = {};
export const visualModelOverride: { current?: LanguageModel } = {};

/** The dispatch-time tool input (from src/agent/tools/jobs.ts schemas). */
type JobToolInput = ResearchToolInput | DemoToolInput | VisualToolInput;

export interface JobRunnerDeps {
	chatStore: ChatStore;
	jobStore: JobStore;
	/**
	 * Hand a finished job's steer to the wake pipeline — resolves when the
	 * steer is delivered (the wake turn completed and the job status
	 * flipped); rejects when delivery fails unrecoverably.
	 */
	deliver: (steer: SubagentSteerMessage) => Promise<void>;
}

export interface JobRunner {
	/** Run one job's sub-agent loop to completion (background; never throws).
	 *  The tool input rides along from dispatch (the spawn fires mid-stream,
	 *  before the tool part is persisted — only crash-recovery re-spawns
	 *  resolve the input from the transcript). */
	runJob(jobId: string, input?: ResearchToolInput | DemoToolInput | VisualToolInput): Promise<void>;
}

export function createJobRunner(deps: JobRunnerDeps): JobRunner {
	const { chatStore, jobStore, deliver } = deps;

	/** The kind→runner map: one sub-agent loop per job kind. */
	// Each kind's runner narrows the union input itself (dispatch-time inputs
	// are kind-correct by construction; transcript fallbacks are cast after).
	const KIND_RUNNERS: Record<JobKind, (jobId: string, input?: JobToolInput) => Promise<void>> = {
		research: (jobId, input) => runResearchJob(jobId, researchModelOverride.current, input),
		demo: (jobId, input) => runDemoJob(jobId, demoModelOverride.current, input),
		visual: (jobId, input) => runVisualJob(jobId, visualModelOverride.current, input),
	};

	async function runJob(jobId: string, input?: JobToolInput): Promise<void> {
		const job = jobStore.getJob(jobId);
		if (!job) {
			console.error(`[job ${jobId}] not found — nothing to run`);
			return;
		}
		if (job.status !== "running") return; // already finished / re-run after crash recovery
		const run = KIND_RUNNERS[job.kind];
		if (!run) {
			console.error(`[job ${jobId}] unknown kind: ${job.kind}`);
			jobStore.updateJob(jobId, { status: "error", error: `unknown kind ${job.kind}`, activity: "failed" });
			return;
		}
		await run(jobId, input);
	}

	/** Resolve the job's tool input: dispatch passes it directly; only a
	 *  crash-recovery re-spawn (no input) falls back to scanning the
	 *  transcript (by then onFinish has long persisted the part). */
	function resolveInput(
		job: JobRecord,
		input: JobToolInput | undefined,
	): ResearchToolInput | DemoToolInput | VisualToolInput | undefined {
		if (input) return input;
		const chat = chatStore.getChat(job.chatId);
		if (!chat) return undefined;
		return (findToolPart(chat.messages, job.toolCallId)?.input ?? undefined) as
			| ResearchToolInput
			| DemoToolInput
			| VisualToolInput
			| undefined;
	}

	// ─── shared plumbing ────────────────────────────────────────────────────

	function findToolPart(messages: UIMessage[], toolCallId: string) {
		for (const m of messages) {
			for (const part of m.parts) {
				if (isToolUIPart(part) && part.toolCallId === toolCallId) return part;
			}
		}
		return null;
	}

	/**
	 * The tail of every job: store the result, build the steer, hand it to
	 * the wake pipeline. The pipeline flips the status when the lesson has
	 * actually resumed — until then the job reads "running" so the client's
	 * activity card keeps polling and re-syncs once, at the end.
	 */
	async function finishJob(jobId: string, result: ResearchResult | DemoResult | VisualResult): Promise<void> {
		const job = jobStore.getJob(jobId);
		if (!job) return;
		jobStore.updateJob(jobId, { result, activity: "delivering result" });
		const steer = buildSubagentSteer(job, result);
		try {
			await deliver(steer);
		} catch (err) {
			const reason = err instanceof Error ? err.message : String(err);
			console.error(`[job ${jobId}] steer delivery failed:`, err);
			jobStore.updateJob(jobId, { status: "error", error: reason, activity: "delivery failed" });
		}
	}

	/** Chat-record + input resolution shared by all three runners. */
	function jobContext(jobId: string): { job: JobRecord; chatId: string } | null {
		const job = jobStore.getJob(jobId);
		if (!job || job.status !== "running") return null;
		const chat = chatStore.getChat(job.chatId);
		if (!chat) {
			jobStore.updateJob(jobId, { status: "error", error: `Chat not found: ${job.chatId}`, activity: "failed" });
			return null;
		}
		return { job, chatId: job.chatId };
	}

	// ─── research ───────────────────────────────────────────────────────────

	async function runResearchJob(jobId: string, model?: LanguageModel, passedInput?: JobToolInput): Promise<void> {
		const ctx = jobContext(jobId);
		if (!ctx) return;
		const { job } = ctx;

		const input = (passedInput ?? resolveInput(job, passedInput)) as ResearchToolInput | undefined;
		const question = typeof input?.question === "string" && input.question.trim() ? input.question : "(no question)";
		const context = typeof input?.context === "string" && input.context.trim() ? input.context : undefined;

		const prompt = [
			`Research question: ${question}`,
			context ? `Context (what this is needed for): ${context}` : null,
			"Produce your final brief in the required format.",
		]
			.filter(Boolean)
			.join("\n\n");

		const fetchedUrls: string[] = [];
		const agent = createResearcherAgent(model);

		try {
			jobStore.updateJob(jobId, { activity: "researching" });
			const result = await agent.generate({
				prompt,
				onToolExecutionStart: ({ toolCall }) => {
					const input = toolCall.input as { query?: string; url?: string } | undefined;
					if (toolCall.toolName === "web_fetch" && typeof input?.url === "string") {
						fetchedUrls.push(input.url);
					}
					jobStore.updateJob(jobId, { activity: activityLine(toolCall.toolName, input) });
				},
			});

			const brief = result.text.trim();
			if (!brief) throw new Error("the researcher produced no text");
			const research: ResearchResult = { brief, sources: parseSources(brief, fetchedUrls) };
			await finishJob(jobId, research);
		} catch (err) {
			const reason = err instanceof Error ? err.message : String(err);
			jobStore.updateJob(jobId, { error: reason });
			await finishJob(jobId, { brief: `Research failed: ${reason}`, sources: [] });
		}
	}

	// ─── demo + visual ──────────────────────────────────────────────────────

	async function runDemoJob(jobId: string, model?: LanguageModel, passedInput?: JobToolInput): Promise<void> {
		const ctx = jobContext(jobId);
		if (!ctx) return;
		const { job } = ctx;

		const input = (passedInput ?? resolveInput(job, passedInput)) as DemoToolInput | undefined;
		const title = typeof input?.title === "string" && input.title.trim() ? input.title : "Interactive demo";
		const brief = typeof input?.brief === "string" ? input.brief : "";

		try {
			jobStore.updateJob(jobId, { activity: "building demo" });
			const result = await runDemoBuilder(
				{ title, brief },
				{
					model,
					onActivity: (line) => jobStore.updateJob(jobId, { activity: line }),
				},
			);
			if (isDemoFailure(result)) jobStore.updateJob(jobId, { error: result.notes ?? "demo build failed" });
			await finishJob(jobId, result);
		} catch (err) {
			const reason = err instanceof Error ? err.message : String(err);
			jobStore.updateJob(jobId, { error: reason });
			await finishJob(jobId, { title, html: "", tested: true, notes: `FAILED: ${reason}` });
		}
	}

	async function runVisualJob(jobId: string, model?: LanguageModel, passedInput?: JobToolInput): Promise<void> {
		const ctx = jobContext(jobId);
		if (!ctx) return;
		const { job } = ctx;

		const input = (passedInput ?? resolveInput(job, passedInput)) as VisualToolInput | undefined;
		const brief = typeof input?.brief === "string" ? input.brief : "";

		try {
			jobStore.updateJob(jobId, { activity: "drawing figure" });
			const result = await runVisualMaker(
				{ brief },
				{
					model,
					onActivity: (line) => jobStore.updateJob(jobId, { activity: line }),
				},
			);
			if (isVisualFailure(result)) jobStore.updateJob(jobId, { error: result.notes ?? "visual build failed" });
			await finishJob(jobId, result);
		} catch (err) {
			const reason = err instanceof Error ? err.message : String(err);
			jobStore.updateJob(jobId, { error: reason });
			await finishJob(jobId, { svg: "", notes: `FAILED: ${reason}` });
		}
	}

	return { runJob };
}

// ─── module-level pure helpers ──────────────────────────────────────────────

/** Failure shape (W6 brief): html === "" with "FAILED: …" notes. */
function isDemoFailure(result: DemoResult): boolean {
	return typeof result.html !== "string" || result.html.length === 0;
}

function isVisualFailure(result: VisualResult): boolean {
	return typeof result.svg !== "string" || result.svg.length === 0;
}

/** The activity line shown in the UI: `web_search: "query"` / `web_fetch: host`. */
function activityLine(toolName: string, input: { query?: string; url?: string } | undefined): string {
	if (toolName === "web_search") return `web_search: "${input?.query ?? ""}"`;
	if (toolName === "web_fetch") return `web_fetch: ${hostOf(input?.url)}`;
	return toolName;
}

function hostOf(url: string | undefined): string {
	if (!url) return "(no url)";
	try {
		return new URL(url).hostname;
	} catch {
		return url.slice(0, 40);
	}
}

/** Sources for the result: markdown links cited in the brief, else fetched URLs. */
function parseSources(brief: string, fetchedUrls: string[]): ResearchResult["sources"] {
	const sources: ResearchResult["sources"] = [];
	const seen = new Set<string>();
	for (const match of brief.matchAll(/\[([^\]]+)\]\((https?:\/\/[^)\s]+)\)/g)) {
		const title = (match[1] ?? "").trim();
		const url = match[2]!;
		if (!title || seen.has(url)) continue;
		seen.add(url);
		sources.push({ title, url });
	}
	if (sources.length === 0) {
		for (const url of fetchedUrls) {
			if (seen.has(url)) continue;
			seen.add(url);
			sources.push({ title: hostOf(url), url });
		}
	}
	return sources.slice(0, 10);
}