Master seam: plan-block emission + turn usage capture
- src/core/plan-blocks.ts: pure emission builder — full block on successful plan_init, neighborhood block when the derived plan state moved since the newest existing block; the as-of state reads from the block's frozen snapshot (transcript quiz parts mutate in place on answer, so re-derivation is post-mutation blind). Idempotent by construction. - turn-service persistTurn enriches with withPlanBlocks at the canonical write; chat.tsx adopts the record once the server agrees the turn is done (endTurn flips only after persist). - chat-handler: per-request usage collector via agent-level onStepFinish (StepResult.usage), appended at onFinish to ChatRecord.turns via appendTurnUsage, idempotent per response message id.
This commit is contained in:
@@ -1,6 +1,6 @@
|
||||
import * as fs from "node:fs";
|
||||
import * as path from "node:path";
|
||||
import { ToolLoopAgent, stepCountIs, type LanguageModel, type ToolSet } from "ai";
|
||||
import { ToolLoopAgent, stepCountIs, type LanguageModel, type LanguageModelUsage, type ToolSet } from "ai";
|
||||
import { PATHS } from "../../config.ts";
|
||||
import { modelFor } from "./models.ts";
|
||||
import { questionTools } from "./tools.ts";
|
||||
@@ -58,7 +58,16 @@ function withGraphContext(graphContext: string | null | undefined): string {
|
||||
return `${instructions}\n\n---\n\n${graphContext}`;
|
||||
}
|
||||
|
||||
export function createTeacherAgent(model?: LanguageModel, tools?: ToolSet, graphContext?: string | null) {
|
||||
export function createTeacherAgent(
|
||||
model?: LanguageModel,
|
||||
tools?: ToolSet,
|
||||
graphContext?: string | null,
|
||||
// Money channel (wave 4): the chat-handler passes a per-turn collector;
|
||||
// every model call's usage lands in the chat record's turns log
|
||||
// (ChatStore.appendTurnUsage). The step carries far more than usage —
|
||||
// the narrow param keeps the seam honest about what we consume.
|
||||
onStepFinish?: (step: { usage: LanguageModelUsage }) => void,
|
||||
) {
|
||||
return new ToolLoopAgent({
|
||||
// Model choice + thinking budget live in config.ts (MODELS.teacher).
|
||||
// Reasoning tokens stream as reasoning parts into the UI's collapsible block.
|
||||
@@ -74,6 +83,7 @@ export function createTeacherAgent(model?: LanguageModel, tools?: ToolSet, graph
|
||||
// job tools resolve immediately (fire-and-forget ack). stepCountIs is
|
||||
// just the safety net.
|
||||
stopWhen: stepCountIs(20),
|
||||
...(onStepFinish ? { onStepFinish } : {}),
|
||||
});
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,169 @@
|
||||
/**
|
||||
* Plan-block emission (wave 4, the seam the dock/popup render against):
|
||||
* the PURE builder that turns a plan tree + transcript into the persisted
|
||||
* `data-plan-block` parts the chat renders in flow. Emission happens ONCE,
|
||||
* at persist time (turn-service seam) — blocks are frozen transcript
|
||||
* content (user decision, 2026-09-11), never derived at render:
|
||||
*
|
||||
* - FULL block: attached to the assistant message carrying a successful
|
||||
* plan_init call — the whole DAG, snapshot statuses. Always emitted on
|
||||
* (re)authoring, even an identical rewrite: authoring is the event.
|
||||
* - NEIGHBORHOOD block: attached to the LAST assistant message when the
|
||||
* derived plan state MOVED since the newest existing block — current
|
||||
* shifted or a check passed (done grew). A resolved FAILED check also
|
||||
* moves: done doesn't grow, but the check is no longer out, so current
|
||||
* falls back to the first eligible node — the block reflects that
|
||||
* standing. The as-of state is read from the newest block's FROZEN
|
||||
* SNAPSHOT (not re-derived from the transcript, which the client
|
||||
* re-posts with answered quizzes mutated in place) and compared with
|
||||
* the state derived now — no extra state channel exists.
|
||||
*
|
||||
* Idempotent by construction: a message that already carries a plan-block
|
||||
* part is never given another (persist replays / retry paths are safe).
|
||||
*/
|
||||
|
||||
import type { PlanBlockData, PlanNode, PlanTree } from "./contracts.ts";
|
||||
import { derivePlanStatuses } from "./contracts.ts";
|
||||
import type { UIMessage } from "ai";
|
||||
|
||||
/** The message part the chat renderer consumes (TeacherUIDataParts key
|
||||
* "plan-block" — the SDK strips the `data-` prefix from the name). */
|
||||
export interface PlanBlockPart {
|
||||
type: "data-plan-block";
|
||||
id: string;
|
||||
data: PlanBlockData;
|
||||
}
|
||||
|
||||
/** Nodes/edges of the one-level neighborhood: current + direct parents +
|
||||
* direct children, and the edges among exactly those nodes. */
|
||||
export function neighborhood(tree: PlanTree, current: string): { nodes: PlanNode[]; edges: Array<{ from: string; to: string }> } {
|
||||
const parents = tree.edges.filter((e) => e.to === current).map((e) => e.from);
|
||||
const children = tree.edges.filter((e) => e.from === current).map((e) => e.to);
|
||||
const ids = new Set([current, ...parents, ...children]);
|
||||
return {
|
||||
nodes: tree.nodes.filter((n) => ids.has(n.id)),
|
||||
edges: tree.edges.filter((e) => ids.has(e.from) && ids.has(e.to)),
|
||||
};
|
||||
}
|
||||
|
||||
/** Does this message carry a successful plan_init (the full-block trigger)? */
|
||||
function hasSuccessfulPlanInit(message: UIMessage): boolean {
|
||||
for (const part of message.parts) {
|
||||
if (part.type !== "tool-plan_init") continue;
|
||||
const ok = (part as { output?: { ok?: unknown } }).output?.ok;
|
||||
if (ok === true) return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
/** The plan state as derived from a message slice — the diff key. `checked`
|
||||
* = every node with a tagged quiz out or answered (a new check IS a plan
|
||||
* event: the block's glowing current follows it). */
|
||||
function planState(messages: readonly UIMessage[], tree: PlanTree) {
|
||||
const statuses = derivePlanStatuses(tree, messages);
|
||||
const checked = new Set<string>();
|
||||
for (const m of messages) {
|
||||
for (const part of m.parts) {
|
||||
if (part.type !== "tool-quiz") continue;
|
||||
const nodeId = (part as { input?: { nodeId?: unknown } }).input?.nodeId;
|
||||
if (typeof nodeId === "string" && nodeId !== "") checked.add(nodeId);
|
||||
}
|
||||
}
|
||||
return {
|
||||
current: statuses.current,
|
||||
done: [...statuses.done].sort().join(","),
|
||||
checked: [...checked].sort().join(","),
|
||||
};
|
||||
}
|
||||
|
||||
/** The messages up to and including the newest plan-block carrier — the
|
||||
* "as of last block" slice. Null when no block exists yet. */
|
||||
function sliceAsOfLastBlock(messages: readonly UIMessage[]): UIMessage[] | null {
|
||||
for (let i = messages.length - 1; i >= 0; i--) {
|
||||
if (messages[i]!.parts.some((p) => p.type === "data-plan-block")) return messages.slice(0, i + 1);
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
/** The as-of state read from the newest existing block's FROZEN SNAPSHOT —
|
||||
* not re-derived from the transcript (which the client re-posts with quiz
|
||||
* parts mutated in place: an answered check's part flips from
|
||||
* input-available to output-available, so a re-derivation would read the
|
||||
* post-answer state and the move would be invisible). The block is the
|
||||
* as-of record; that is what frozen snapshots are for. */
|
||||
function stateAsOfLastBlock(messages: readonly UIMessage[]): { current: string | null; done: string } | null {
|
||||
for (let i = messages.length - 1; i >= 0; i--) {
|
||||
for (const part of messages[i]!.parts) {
|
||||
if (part.type !== "data-plan-block") continue;
|
||||
const data = (part as { data?: PlanBlockData }).data;
|
||||
if (!data) return null;
|
||||
return {
|
||||
current: data.current,
|
||||
done: data.nodes
|
||||
.filter((n) => n.status === "done")
|
||||
.map((n) => n.id)
|
||||
.sort()
|
||||
.join(","),
|
||||
};
|
||||
}
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
function snapshot(kind: PlanBlockData["kind"], tree: PlanTree, messages: readonly UIMessage[]): PlanBlockData | null {
|
||||
const statuses = derivePlanStatuses(tree, messages);
|
||||
if (kind === "neighborhood" && statuses.current === null) return null;
|
||||
const scope = kind === "full" || statuses.current === null ? tree : neighborhood(tree, statuses.current);
|
||||
return {
|
||||
kind,
|
||||
at: new Date().toISOString(),
|
||||
current: statuses.current,
|
||||
nodes: scope.nodes.map((n) => ({
|
||||
id: n.id,
|
||||
title: n.title,
|
||||
status: statuses.current === n.id ? "current" : statuses.done.has(n.id) ? "done" : "upcoming",
|
||||
})),
|
||||
edges: scope.edges.map((e) => ({ from: e.from, to: e.to })),
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Enrich a message list with plan-block parts — THE emission seam's pure
|
||||
* half. `messages` is the deduped list persistTurn is about to write;
|
||||
* `tree` is the record's current tree (the plan tool saved it during the
|
||||
* turn). Messages already carrying a block are never touched.
|
||||
*/
|
||||
export function withPlanBlocks(messages: UIMessage[], tree: PlanTree | null | undefined): UIMessage[] {
|
||||
if (!tree || tree.nodes.length === 0 || messages.length === 0) return messages;
|
||||
|
||||
const lastAssistant = [...messages].reverse().find((m) => m.role === "assistant");
|
||||
if (!lastAssistant) return messages;
|
||||
|
||||
// neighborhood needed? the plan state moved since the newest existing block
|
||||
let moved = false;
|
||||
const asOf = stateAsOfLastBlock(messages);
|
||||
if (asOf === null) {
|
||||
moved = planState(messages, tree).checked !== ""; // first block: any plan event ever
|
||||
} else {
|
||||
const now = planState(messages, tree);
|
||||
moved = asOf.current !== now.current || asOf.done !== now.done;
|
||||
}
|
||||
|
||||
const blockId = `pb_${Date.now().toString(36)}_${crypto.randomUUID().slice(0, 8)}`;
|
||||
let fullAttached = false;
|
||||
return messages.map((m) => {
|
||||
if (m.role !== "assistant" || m.parts.some((p) => p.type === "data-plan-block")) return m;
|
||||
if (hasSuccessfulPlanInit(m)) {
|
||||
const data = snapshot("full", tree, messages);
|
||||
if (!data) return m;
|
||||
fullAttached = true;
|
||||
return { ...m, parts: [...m.parts, { type: "data-plan-block", id: `${blockId}_f`, data } satisfies PlanBlockPart] };
|
||||
}
|
||||
if (moved && !fullAttached && m.id === lastAssistant.id) {
|
||||
const data = snapshot("neighborhood", tree, messages);
|
||||
if (!data) return m;
|
||||
return { ...m, parts: [...m.parts, { type: "data-plan-block", id: `${blockId}_n`, data } satisfies PlanBlockPart] };
|
||||
}
|
||||
return m;
|
||||
});
|
||||
}
|
||||
@@ -1,12 +1,15 @@
|
||||
import { createAgentUIStreamResponse, InvalidToolInputError, NoSuchToolError, type UIMessage } from "ai";
|
||||
import { createAgentUIStreamResponse, InvalidToolInputError, NoSuchToolError, type LanguageModelUsage, type UIMessage } from "ai";
|
||||
import { createTeacherAgent, type TeacherAgent } from "../agent/teacher-agent.ts";
|
||||
import { createJobRunner } from "../agent/subagents/runner.ts";
|
||||
import { createExtractionPipeline } from "../agent/subagents/extraction.ts";
|
||||
import { createTeacherTools } from "../agent/tools.ts";
|
||||
import { modelFor } from "../agent/models.ts";
|
||||
import { MODELS } from "../../config.ts";
|
||||
import { claimSteers, deliverJobResult, setJobRespawn, setTurnRunner } from "./jobs.ts";
|
||||
import { stores } from "./stores.ts";
|
||||
import { beginTurn, endTurn } from "./turn-liveness.ts";
|
||||
import { buildTeacherContext, persistTurn, triggerExtraction } from "./turn-service.ts";
|
||||
import type { CallUsage } from "../core/contracts.ts";
|
||||
|
||||
/**
|
||||
* One chat turn: run the teacher agent over the client's message list and
|
||||
@@ -93,6 +96,21 @@ export async function handleChatRequest(req: Request, agent?: TeacherAgent): Pro
|
||||
console.error(`[chat ${chatId}] write-ahead persist failed:`, err);
|
||||
}
|
||||
|
||||
// MONEY CHANNEL (wave 4): one collector per request, fed by the agent's
|
||||
// step callback. Every model call's usage lands in the chat record's
|
||||
// append-only turns log at onFinish — idempotent per response message id,
|
||||
// so abort replays and wake retries append exactly once.
|
||||
const steps: CallUsage[] = [];
|
||||
const collectUsage = (step: { usage: LanguageModelUsage }) => {
|
||||
steps.push({
|
||||
inputTokens: step.usage.inputTokens,
|
||||
outputTokens: step.usage.outputTokens,
|
||||
cacheReadTokens: step.usage.inputTokenDetails?.cacheReadTokens,
|
||||
cacheWriteTokens: step.usage.inputTokenDetails?.cacheWriteTokens,
|
||||
reasoningTokens: step.usage.outputTokenDetails?.reasoningTokens,
|
||||
});
|
||||
};
|
||||
|
||||
try {
|
||||
return await createAgentUIStreamResponse({
|
||||
agent:
|
||||
@@ -106,6 +124,7 @@ export async function handleChatRequest(req: Request, agent?: TeacherAgent): Pro
|
||||
spawn: (job, input) => void jobRunner.runJob(job.id, input),
|
||||
}),
|
||||
buildTeacherContext(stores.graphStore, chatId),
|
||||
collectUsage,
|
||||
),
|
||||
uiMessages: outgoing,
|
||||
// Tool-call validation failures are messages OUR schemas authored (the
|
||||
@@ -121,6 +140,17 @@ export async function handleChatRequest(req: Request, agent?: TeacherAgent): Pro
|
||||
},
|
||||
onFinish: ({ messages, isAborted }) => {
|
||||
persistTurn(stores.chatStore, chatId, messages as UIMessage[], isAborted);
|
||||
// usage log append (idempotent per response message id) — the
|
||||
// stats bar refolds this on the client after the turn settles
|
||||
const response = [...messages].reverse().find((m) => m.role === "assistant");
|
||||
if (response && steps.length > 0) {
|
||||
stores.chatStore.appendTurnUsage(chatId, {
|
||||
messageId: response.id,
|
||||
at: new Date().toISOString(),
|
||||
model: MODELS.teacher.model,
|
||||
steps,
|
||||
});
|
||||
}
|
||||
triggerExtraction(stores, extraction, chatId);
|
||||
// LAST: a poller that sees not-in-flight must read an already-persisted
|
||||
// record — endTurn only after persistTurn has done its job.
|
||||
|
||||
@@ -17,6 +17,7 @@
|
||||
*/
|
||||
|
||||
import type { UIMessage } from "ai";
|
||||
import { withPlanBlocks } from "../core/plan-blocks.ts";
|
||||
import { buildGraphContext } from "../core/graph-context.ts";
|
||||
import { getGraph, pendingQuizAnswers } from "../core/graph.ts";
|
||||
import type { ChatStore, GraphStore } from "../core/ports.ts";
|
||||
@@ -76,7 +77,12 @@ export function withMessageIds(messages: UIMessage[]): UIMessage[] {
|
||||
|
||||
export function persistTurn(chatStore: ChatStore, chatId: string, messages: UIMessage[], isAborted: boolean): void {
|
||||
try {
|
||||
chatStore.saveMessages(chatId, withMessageIds(messages));
|
||||
const deduped = withMessageIds(messages);
|
||||
// Plan-block emission (wave 4): blocks are frozen transcript content —
|
||||
// attached here, once, at the canonical state write. The tree was
|
||||
// saved by the plan tool during the turn; snapshots derive now.
|
||||
const tree = chatStore.getChat(chatId)?.tree;
|
||||
chatStore.saveMessages(chatId, withPlanBlocks(deduped, tree));
|
||||
} catch (err) {
|
||||
console.error(`[chat ${chatId}] persist failed:`, err, { isAborted });
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user