feat: graph context injection — the session's prior knowledge, steered in at session start (vision-context-injection.md)

One shot per chat, fired at the first turn's onFinish: the session's
topic (its only typed message) is embedded as an EmbeddingGemma retrieval
query; the semantically closest graph propositions are selected under the
session's token budget (pure budget, no similarity threshold — the
decided doctrine) and delivered as a marker-only steer through the wake
pipeline, queuing behind pending questions like every sub-agent result.

- contracts: ChatSettings.graphContextTokens (default 1024, max 4096,
  step 128, 0 = off); GraphContextData + graphContextText; the
  graph-context JobKind + buildGraphContextSteer (id steer_ctx_<chatId>)
- core/graph-context.ts: brute-force cosine top-K, greedy chars/4 fill
  under budget, most-recent UnderstandingLevel fold, count cap 64,
  parse-don't-cast at the marker seam
- server: embedQuery (role-asymmetric query format, same warm singleton
  — vector-space parity by construction); the runner (startGraphContext-
  Injection one-shot gates; runGraphContextJob validates ledger parity);
  wake branches (chatIdOf/jobIdOf, finalizeJob done, recoverStaleJobs
  re-deliver or contextRespawn); abortJob kills context jobs SILENTLY
  (nothing was promised in the lesson); turn-service trigger seam
- adapter: data-graph-context flattens to the canonical block (malformed
  parts have no model face)
- UI: the budget slider (home dock + /chat/new draft, the session-data
  strip); the quiet SystemLine "~537 · 10 propositions"; dock kind
  "Knowledge context"; model-picker PUT now spreads settings (a model
  pick can no longer clear the sibling override)

Live-smoked end to end on the real embedder + real GLM: seeded ledger
(10 propositions), session on "modal logic: possible worlds and the
necessity box" — job in the dock, queued behind the prober's question,
delivered on the answer, SystemLine in the transcript, wake turn reads
the block and continues the probe. Gates: tsc, lint, 459/459, build.
This commit is contained in:
amosblomqvist
2026-09-15 07:54:50 +03:00
parent ecb293a77d
commit 1d45d77303
25 changed files with 1443 additions and 28 deletions
+20 -3
View File
@@ -4,14 +4,15 @@ import { createProberAgent, createTeacherAgent, type TeachingAgent } from "../ag
import { piEngineFor } from "../agent/models.ts";
import { createJobRunner } from "../agent/subagents/runner.ts";
import { createExtractionPipeline } from "../agent/subagents/extraction.ts";
import { embedNewNodes } from "./embeddings.ts";
import { runGraphContextJob, type GraphContextDeps } from "../agent/subagents/graph-context.ts";
import { EMBEDDING_DIM, EMBEDDING_MODEL, embedNewNodes, embedQuery } from "./embeddings.ts";
import type { ModelRole } from "../../config.ts";
import { claimSteers, deliverJobResult, setJobRespawn, setTurnRunner } from "./jobs.ts";
import { claimSteers, deliverJobResult, setContextRespawn, setJobRespawn, setTurnRunner } from "./jobs.ts";
import { recordToPiMessages } from "./pi/record-adapter.ts";
import { createPiTurnResponse } from "./pi/ui-stream.ts";
import { stores } from "./stores.ts";
import { beginTurn, endTurn } from "./turn-liveness.ts";
import { persistTurn, triggerExtraction, withMessageIds } from "./turn-service.ts";
import { persistTurn, triggerExtraction, triggerGraphContext, withMessageIds } from "./turn-service.ts";
import type { CallUsage } from "../core/contracts.ts";
/**
@@ -48,6 +49,18 @@ export const extraction = createExtractionPipeline({
embedder: embedNewNodes,
});
/** Graph context injection (vision-context-injection.md): the runner's
* deps — same storage ports, the REAL embedder as the query seam (same
* warm singleton as extraction, so vector-space parity holds by
* construction), and the wake pipeline as the delivery seam. */
export const graphContext: GraphContextDeps = {
chatStore: stores.chatStore,
jobStore: stores.jobStore,
graphStore: stores.graphStore,
embedder: { model: EMBEDDING_MODEL, dim: EMBEDDING_DIM, embedQuery },
deliver: deliverJobResult,
};
// Wake-pipeline wiring (contracts.ts "THE ASYNC PROTOCOL" §4): the pipeline
// fires its server-side teacher turns through THIS composition root, and
// crash recovery re-spawns orphaned sub-agent loops through THIS runner.
@@ -65,6 +78,7 @@ setTurnRunner(async (chatId, messages) => {
return res.status;
});
setJobRespawn((jobId) => void jobRunner.runJob(jobId));
setContextRespawn((jobId) => void runGraphContextJob(graphContext, jobId));
/** The message-shape gate the SDK's UIMessageValidator owned: a malformed
* body must fail BEFORE the stream exists — the path where no onFinish will
@@ -223,6 +237,9 @@ export async function handleChatRequest(
});
}
triggerExtraction(stores, extraction, chatId);
// The graph READ channel: one-shot prior-knowledge injection at
// session start — a no-op for every session past its start.
triggerGraphContext(graphContext, chatId);
// LAST: a poller that sees not-in-flight must read an already-persisted
// record — endTurn only after persistTurn has done its job.
endTurn(chatId);
+18
View File
@@ -42,6 +42,24 @@ function embedInput(statement: string): string {
return `title: none | text: ${statement}`;
}
/** EmbeddingGemma's retrieval-QUERY format (model card) — the topic side of
* graph context injection (vision-context-injection.md). Role-asymmetric
* BY DESIGN: queries and documents use different trained templates; the
* requirement is same model + same truncation (one vector space), each
* side in the format the model was trained for. */
function queryInput(text: string): string {
return `task: search result | query: ${text}`;
}
/** Embed one retrieval query (the session's topic) — same warm singleton,
* same 512-dim MRL truncation as the ledger. The one embedding call the
* injection job makes; tens–~200ms warm. */
export async function embedQuery(text: string): Promise<number[]> {
const extractor = await getPipeline();
const output = await extractor([queryInput(text)], { pooling: "mean", normalize: true });
return truncateNormalize((output.tolist() as number[][])[0]!);
}
/** Truncate an L2-normalized vector to EMBEDDING_DIM (MRL: leading dims
* carry most of the signal) and re-normalize so cosine = dot product. */
function truncateNormalize(vec: readonly number[]): number[] {
+56 -7
View File
@@ -34,6 +34,7 @@
import type { UIMessage } from "ai";
import {
buildAbortSteer,
buildGraphContextSteer,
buildSubagentSteer,
isSubagentFailure,
steerDelivered,
@@ -60,6 +61,9 @@ interface QueuedSteer {
interface WakeState {
turnRunner?: TurnRunner;
respawn?: JobRespawn;
/** Crash recovery for context-injection jobs (they are not sub-agent
* loops — their runner takes the graph deps; chat-handler wires it). */
contextRespawn?: (jobId: string) => void;
/** chatId → steers awaiting delivery, in arrival order. */
queues: Map<string, QueuedSteer[]>;
/** Chats whose pump loop is currently active (serialization per chat). */
@@ -87,6 +91,13 @@ export function setJobRespawn(respawn: JobRespawn): void {
state().respawn = respawn;
}
/** chat-handler wires this at module init: the graph-context runner
* (subagents/graph-context.ts runGraphContextJob — re-runnable from the
* job record alone, which is all crash recovery has). */
export function setContextRespawn(respawn: (jobId: string) => void): void {
state().contextRespawn = respawn;
}
const POLL_MS = 400;
/** After the gates open, wait a beat before firing — a client that just
* answered a question is about to send its continuation; don't 409 it.
@@ -110,11 +121,14 @@ function delay(ms: number): Promise<void> {
return new Promise((resolve) => setTimeout(resolve, ms));
}
/** The job a steer belongs to, if it is a job steer (result/abort);
/** The job a steer belongs to, if it is a job steer (result/abort/context);
* probe-handoff steers have no job — null. */
function jobIdOf(steer: SteerMessage): string | null {
const meta = steer.metadata as { subagentResult?: { jobId: string } };
return meta.subagentResult?.jobId ?? null;
const meta = steer.metadata as {
subagentResult?: { jobId: string };
graphContext?: { jobId: string };
};
return meta.subagentResult?.jobId ?? meta.graphContext?.jobId ?? null;
}
/** Finalize a delivered steer's job bookkeeping — job steers only. */
@@ -220,6 +234,12 @@ export function abortJob(jobId: string): AbortOutcome {
if (!job) return { ok: false, reason: "not_found" };
if (job.status !== "running") return { ok: false, reason: "not_running" };
// Context-injection jobs abort SILENTLY: nothing was ever promised in the
// lesson (server-initiated scaffolding, unlike a teacher-dispatched
// sub-agent), so no abort steer — flipping the record and killing the
// run is the whole story. Steps 1–3 below are shared; step 4 is skipped.
const isContext = job.kind === "graph-context";
fireJobAbort(jobId); // no-op when the loop already returned (delivering)
// Drop any queued steers for this job (their deliver() resolves unfired).
@@ -234,16 +254,22 @@ export function abortJob(jobId: string): AbortOutcome {
}
stores.jobStore.updateJob(jobId, { status: "error", error: "Aborted by learner", activity: "aborted" });
void deliverJobResult(buildAbortSteer(job)).catch(() => {
/* deliverJobResult marks the job error on unrecoverable failure */
});
if (!isContext) {
void deliverJobResult(buildAbortSteer(job)).catch(() => {
/* deliverJobResult marks the job error on unrecoverable failure */
});
}
return { ok: true };
}
function chatIdOf(steer: SteerMessage): string {
// The probe handoff carries its chat directly (no job record exists).
// The probe handoff and the context injection carry their chat directly
// (the context steer's job record would say it too, but the metadata is
// the cheap path — same rule as the handoff).
const handoff = (steer.metadata as { probeHandoff?: { chatId: string } }).probeHandoff;
if (handoff) return handoff.chatId;
const injected = (steer.metadata as { graphContext?: { chatId: string } }).graphContext;
if (injected) return injected.chatId;
// The chat is resolvable from the job record; the steer's metadata
// carries jobId, which indexes the store.
const jobId = jobIdOf(steer);
@@ -352,6 +378,12 @@ async function pump(chatId: string): Promise<void> {
function finalizeJob(jobId: string): void {
const job = stores.jobStore.getJob(jobId);
if (!job || job.status !== "running") return;
// Context-injection jobs have no sub-agent result shape — a delivered
// block is a success, full stop (isSubagentFailure would misread it).
if (job.kind === "graph-context") {
stores.jobStore.updateJob(jobId, { status: "done", activity: "delivered" });
return;
}
const result = job.result;
const failed = result ? isSubagentFailure(job.kind, result as never) : true;
stores.jobStore.updateJob(jobId, {
@@ -418,6 +450,23 @@ function recoverStaleJobs(): void {
if (now - Date.parse(job.updatedAt) < STALE_JOB_MS) continue;
if (!stores.chatStore.getChat(job.chatId)) continue;
if (isTurnRunning(job.chatId)) continue; // a wake turn may be delivering it right now
if (job.kind === "graph-context") {
// Context-injection orphans: a stored block re-delivers in its own
// steer shape; a job that died before selecting re-runs from the
// record (the graph runner — not the sub-agent respawn seam).
if (job.result) {
stores.jobStore.updateJob(job.id, { activity: "recovering: re-delivering" });
void deliverJobResult(
buildGraphContextSteer({ chatId: job.chatId, jobId: job.id }, job.result as never),
).catch(() => {
/* deliverJobResult marks the job error on failure */
});
} else if (s.contextRespawn) {
stores.jobStore.updateJob(job.id, { activity: "recovering: re-running" });
s.contextRespawn(job.id);
}
continue;
}
if (job.result) {
// The wake never fired (or crashed mid-flight): re-deliver.
const steer = buildSubagentSteer(job, job.result as never);
+12 -4
View File
@@ -1,7 +1,8 @@
import type { UIMessage } from "ai";
import type { AgentMessage } from "@earendil-works/pi-agent-core";
import type { AssistantMessage, Message, ToolResultMessage } from "@earendil-works/pi-ai";
import { HANDOFF_TEXT, verdictText } from "../../core/contracts/chat.ts";
import { graphContextOfMessage } from "../../core/graph-context.ts";
import { graphContextText, HANDOFF_TEXT, verdictText } from "../../core/contracts/chat.ts";
/**
* The record → pi-ai transcript adapter (docs/pi-engine-swap.md §4). The
@@ -55,12 +56,19 @@ function zeroUsage() {
/** The model-facing text of a marker-only user message, or null when the
* message has no model face at all (UI-only data parts). */
function systemTextOf(parts: ReadonlyArray<{ type: string; data?: unknown }>): string | null {
for (const part of parts) {
function systemTextOf(message: { parts: ReadonlyArray<{ type: string; data?: unknown }> }): string | null {
for (const part of message.parts) {
if (part.type === "data-verdict") {
return verdictText((part.data as { accepted?: unknown } | undefined)?.accepted === true);
}
if (part.type === "data-probe-handoff") return HANDOFF_TEXT;
// Graph context injection: the block text derives from the part's
// validated data (parse-don't-cast — a malformed part has no model
// face, and the message then never reaches the provider).
if (part.type === "data-graph-context") {
const data = graphContextOfMessage({ parts: message.parts });
return data ? graphContextText(data) : null;
}
}
return null;
}
@@ -79,7 +87,7 @@ export function recordToPiMessages(messages: UIMessage[]): Message[] {
// wording lives in the adapter's reach (core contracts), never
// authored as fake human speech in the transcript.
if (texts.length === 0) {
const flat = systemTextOf(message.parts);
const flat = systemTextOf(message);
if (flat === null) continue; // pure UI data part, no model face
texts.push(flat);
}
+20
View File
@@ -21,6 +21,7 @@ import { pendingQuizAnswers } from "../core/graph.ts";
import type { ChatStore, GraphStore } from "../core/ports.ts";
import { EXTRACTION } from "../../config.ts";
import type { ExtractionPipeline } from "../agent/subagents/extraction.ts";
import { startGraphContextInjection, type GraphContextDeps } from "../agent/subagents/graph-context.ts";
/**
* The canonical state write: the client owns the message list between turns
@@ -98,3 +99,22 @@ export function triggerExtraction(
console.error(`[chat ${chatId}] extraction trigger failed:`, err);
}
}
/**
* The graph READ channel (vision-context-injection.md): one shot per chat,
* at session start — the first typed message is the session's topic; the
* semantically closest propositions are selected under the session's token
* budget and steered in through the wake pipeline. Fire-and-forget like
* triggerExtraction; the one-shot gates (topic, budget, prior injection)
* live in the runner module, not here — this seam only decides WHEN to ask.
* Skipped under tests, same doctrine: the runner is tested directly with a
* scripted embedder.
*/
export function triggerGraphContext(deps: GraphContextDeps, chatId: string): void {
if (process.env.NODE_ENV === "test") return;
try {
startGraphContextInjection(deps, chatId);
} catch (err) {
console.error(`[chat ${chatId}] graph-context trigger failed:`, err);
}
}