learnharness/src/agent/subagents/graph-context.ts

/**
 * Graph context injection — the background runner
 * (docs/vision-context-injection.md, settled 2026-09-14).
 *
 * One shot per chat, at session start: the first learner message is the
 * session's TOPIC; it is embedded as a retrieval QUERY (the embedder's
 * model-card format), the semantically closest graph propositions are
 * selected under the session's token budget (pure core: selectGraphContext),
 * and delivered as a marker-only STEER through the wake pipeline — the same
 * gates as sub-agent results: it queues behind a pending quiz and rides the
 * turn after the learner's answer. The prober typically opens with a run of
 * quiz/ask calls, so delivery lands mid-probing, before plan_init — where
 * prior knowledge matters most.
 *
 * No teacher-model call: the job's only model spend is ONE embedding
 * query (the local embedder was free, provider ones are not — wave 6
 * records its prompt tokens on the job record). It still carries a
 * JobRecord (kind "graph-context") so the dock shows it, and it
 * registers its AbortController so the dock's END kills it — but an abort
 * steers NOTHING: nothing was ever promised in the lesson (the injection is
 * server-initiated scaffolding, unlike a teacher-dispatched sub-agent).
 *
 * The topic gate is structural: injection fires only while the session has
 * EXACTLY ONE text-bearing user message — the topic. Old in-flight sessions
 * (several typed messages) never inject (no retroactive injection), and a
 * session whose first turn happened with injection off stays off even if
 * the budget is raised later (set per-session, before the session starts).
 */

import {
	accumulateJobUsage,
	buildGraphContextSteer,
	steerDelivered,
	type JobRecord,
	type SteerMessage,
} from "../../core/contracts.ts";
import {
	GRAPH_CONTEXT_TOKENS_DEFAULT,
	type ChatSettings,
} from "../../core/contracts/chat.ts";
import { hasContext, selectGraphContext } from "../../core/graph-context.ts";
import type { UIMessage } from "ai";
import { registerJobAbort, unregisterJobAbort } from "./aborts.ts";
import type { ChatStore, GraphStore, JobStore } from "../../core/ports.ts";

/** The embedder as the runner needs it (BYOK wave 6): the model slug the
 *  ledger must match (parity: one vector space) + the query seam.
 *  Resolved FRESH per run by the composition root (deps.embedder) so a
 *  settings PUT or a key save between dispatch and spawn lands; null
 *  model = no embedding configured — the run skips honestly. The query
 *  carries its prompt-token count so the job records its spend (the
 *  local embedder was free; provider ones are not). */
export interface GraphContextEmbedder {
	model: string | null;
	embedQuery: (text: string) => Promise<{ vector: number[]; promptTokens?: number }>;
}

export interface GraphContextDeps {
	chatStore: Pick<ChatStore, "getChat">;
	jobStore: Pick<JobStore, "createJob" | "getJob" | "updateJob" | "findJobByToolCallId">;
	graphStore: Pick<GraphStore, "readNodes" | "readEvents" | "readEmbeddings">;
	/** The embedder RESOLVER (wave 6): called once at run start — fresh
	 *  setting + key reads each run. null = no embedding model configured. */
	embedder: () => GraphContextEmbedder | null;
	/** The wake seam — jobs.ts deliverJobResult (resolves on delivery). */
	deliver: (steer: SteerMessage) => Promise<void>;
}

/** The injection job's synthetic toolCallId — no tool part exists; the id
 *  is deterministic per chat so findJobByToolCallId doubles as the one-shot
 *  guard, mirroring the spawn-dedup the tool jobs get for free. */
export function contextToolCallId(chatId: string): string {
	return `ctx_${chatId}`;
}

/** The session's effective budget: settings override, else the default.
 *  Absent ≠ 0 — the decided default is ON (1024 tokens). */
export function effectiveContextBudget(settings: ChatSettings | undefined): number {
	return settings?.graphContextTokens ?? GRAPH_CONTEXT_TOKENS_DEFAULT;
}

/** The session's topic: its ONLY text-bearing user message, if exactly one
 *  exists. Steers (verdicts, handoffs, answers) carry no text part, so
 *  quiz rounds during probing never break the gate. */
export function topicOf(messages: UIMessage[]): string | null {
	let topic: string | null = null;
	let count = 0;
	for (const m of messages) {
		if (m.role !== "user") continue;
		const text = m.parts
			.filter((p) => p.type === "text")
			.map((p) => (p as { text?: unknown }).text)
			.filter((t): t is string => typeof t === "string" && t.trim().length > 0)
			.join("\n");
		if (text.length === 0) continue;
		count++;
		topic = text;
	}
	return count === 1 ? topic : null;
}

/**
 * The one-shot trigger (fire-and-forget, called from turn-service at every
 * turn's onFinish): starts the injection job if and only if the session is
 * at its start (exactly one typed message — the topic), the budget is on,
 * and nothing was injected before (delivered steer OR any prior context
 * job — success or failure: one shot means one shot). Returns the created
 * job for tests; null when nothing was started.
 */
export function startGraphContextInjection(deps: GraphContextDeps, chatId: string): JobRecord | null {
	const { chatStore, jobStore } = deps;
	const chat = chatStore.getChat(chatId);
	if (!chat) return null;
	if (effectiveContextBudget(chat.settings) <= 0) return null;
	if (topicOf(chat.messages) === null) return null;
	// One per chat, ever: a delivered steer OR a prior job (any status).
	if (steerDelivered(chat.messages, `steer_ctx_${chatId}`)) return null;
	if (jobStore.findJobByToolCallId(chatId, contextToolCallId(chatId))) return null;

	const job = jobStore.createJob({
		kind: "graph-context",
		chatId,
		toolCallId: contextToolCallId(chatId),
		activity: "searching knowledge graph",
	});
	void runGraphContextJob(deps, job.id).catch((err: unknown) => {
		console.error(`[chat ${chatId}] graph-context injection crashed:`, err);
	});
	return job;
}

/** Run one injection job to completion (background; never throws). Also the
 *  crash-recovery respawn seam: a job record orphaned "running" by a
 *  process crash re-runs from here (the record is the only input needed). */
export async function runGraphContextJob(deps: GraphContextDeps, jobId: string): Promise<void> {
	const { chatStore, jobStore, graphStore, deliver } = deps;
	const job = jobStore.getJob(jobId);
	if (!job || job.status !== "running") return;

	const fail = (reason: string) => {
		jobStore.updateJob(jobId, { status: "error", error: reason, activity: "failed" });
	};
	const chat = chatStore.getChat(job.chatId);
	if (!chat) return fail(`chat not found: ${job.chatId}`);
	// BYOK (wave 6): fresh resolution per run; null = honestly off (the
	// dispatch gate already skips unconfigured sessions — this only fires
	// on a config change racing a respawn, and the error says exactly why).
	// The resolver returns null for BOTH an unset model and a set model
	// whose key vanished, so the message names both, never a false "null".
	const embedder = deps.embedder();
	if (!embedder || embedder.model === null) {
		return fail("no embedding model available — injection skipped (graph.embeddingModel unset, or its provider key is missing)");
	}
	const topic = topicOf(chat.messages);
	if (topic === null) return fail("no injectable topic (session past its start)");
	const budget = effectiveContextBudget(chat.settings);

	// Real-kill wiring: the dock's END fires this controller; the runner
	// checks it before every phase transition and unwinds silently.
	const controller = new AbortController();
	registerJobAbort(jobId, controller);
	try {
		jobStore.updateJob(jobId, { activity: "embedding topic" });
		const ledger = graphStore.readEmbeddings();
		// Staleness is the MODEL identity (wave 6): a given model serves one
		// dim by default, so model parity is space parity; the dim is checked
		// defensively against the served vector below.
		if (ledger.model !== embedder.model) {
			return fail("embeddings ledger is stale (model mismatch) — injection skipped");
		}
		if (Object.keys(ledger.vectors).length === 0) {
			return fail("the knowledge graph has no embeddings yet — nothing to inject");
		}
		const query = await embedder.embedQuery(topic);
		if (controller.signal.aborted) return;
		// The defensive dim half of the old parity check: a provider-side dim
		// change (rare) must fail the query, never silently mismatch.
		if (query.vector.length !== ledger.dim) {
			return fail(
				`embedding dim mismatch: query is ${query.vector.length}, ledger holds ${ledger.dim}`,
			);
		}
		// Money channel (wave 6): the one embedding call's prompt tokens fold
		// into the job record — priced via the EMBEDDING_PRICES overlay.
		if (query.promptTokens !== undefined) {
			jobStore.updateJob(jobId, {
				usage: accumulateJobUsage(job.usage, embedder.model, {
					inputTokens: query.promptTokens,
					outputTokens: 0,
				}),
			});
		}

		jobStore.updateJob(jobId, { activity: "selecting propositions" });
		const data = selectGraphContext({
			chatId: job.chatId,
			query: query.vector,
			vectors: ledger.vectors,
			nodes: graphStore.readNodes(),
			events: graphStore.readEvents(),
			budget,
		});
		if (!hasContext(data)) return fail("nothing relevant fit the budget");
		if (controller.signal.aborted) return;

		// The result rides the record (crash recovery re-delivers it); the
		// wake pipeline flips the job done when the lesson has resumed.
		jobStore.updateJob(jobId, { result: data, activity: "delivering context" });
		try {
			await deliver(buildGraphContextSteer({ chatId: job.chatId, jobId }, data));
		} catch (err) {
			console.error(`[chat ${job.chatId}] graph-context steer delivery failed:`, err);
			fail("context delivery failed");
		}
	} catch (err) {
		if (controller.signal.aborted) return; // the abort path owns the record
		console.error(`[chat ${job.chatId}] graph-context injection failed:`, err);
		fail(err instanceof Error ? err.message : String(err));
	} finally {
		unregisterJobAbort(jobId);
	}
}