diff --git a/components/chat.tsx b/components/chat.tsx index 274ae5c..0e779fa 100644 --- a/components/chat.tsx +++ b/components/chat.tsx @@ -200,12 +200,47 @@ export function Chat({ useEffect(() => { void syncTurns(); }, [syncTurns]); - // After the learner's own turn ends the log has a fresh entry — refetch. + // After the learner's own turn ends the log has a fresh entry — and the + // persisted transcript may have grown server-authored parts the live + // stream can't carry (plan blocks attach at persist). Adopt the record + // the moment the server agrees the turn is done: endTurn flips ONLY + // after persist, so a not-in-flight record is a complete one — the same + // invariant the wake-adoption poll below rides. Poll briefly (the flip + // trails the client's stream-end by milliseconds), then adopt once. const prevStatusRef = useRef(status); useEffect(() => { - if (prevStatusRef.current !== "ready" && status === "ready") void syncTurns(); + if (prevStatusRef.current !== "ready" && status === "ready") { + // immediate: turns-only refetch — the stats bar updates even while + // the server's persist/endTurn lags the client by milliseconds + void syncTurns(); + // deferred: adopt the full record once the server agrees the turn + // is done (plan blocks and any other persist-authored parts) + let cancelled = false; + let timer: ReturnType | undefined; + const adopt = async (attempt: number) => { + if (cancelled) return; + try { + const res = await fetch(`/api/chats/${chatId}`); + if (res.ok) { + const record = (await res.json()) as { turnInFlight?: boolean }; + if (record.turnInFlight !== true) { + void syncFromServer(); + return; + } + } + } catch { + /* transient — keep polling */ + } + if (attempt < 8 && !cancelled) timer = setTimeout(() => void adopt(attempt + 1), 500); + }; + void adopt(0); + return () => { + cancelled = true; + if (timer) clearTimeout(timer); + }; + } prevStatusRef.current = status; - }, [status, syncTurns]); + }, [status, chatId, syncFromServer, syncTurns]); // ── Central job observability (contracts.ts "THE ASYNC PROTOCOL" §4) ── // ONE poll of /api/jobs drives every job surface (transcript dispatch diff --git a/src/agent/teacher-agent.ts b/src/agent/teacher-agent.ts index 3c0c75b..971db29 100644 --- a/src/agent/teacher-agent.ts +++ b/src/agent/teacher-agent.ts @@ -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 } : {}), }); } diff --git a/src/core/plan-blocks.ts b/src/core/plan-blocks.ts new file mode 100644 index 0000000..48baf7f --- /dev/null +++ b/src/core/plan-blocks.ts @@ -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(); + 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; + }); +} diff --git a/src/server/chat-handler.ts b/src/server/chat-handler.ts index 8f7fa69..4e9809f 100644 --- a/src/server/chat-handler.ts +++ b/src/server/chat-handler.ts @@ -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. diff --git a/src/server/turn-service.ts b/src/server/turn-service.ts index e999669..eca8b1b 100644 --- a/src/server/turn-service.ts +++ b/src/server/turn-service.ts @@ -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 }); } diff --git a/tests/plan-blocks.test.ts b/tests/plan-blocks.test.ts new file mode 100644 index 0000000..d54cdf5 --- /dev/null +++ b/tests/plan-blocks.test.ts @@ -0,0 +1,191 @@ +import assert from "node:assert/strict"; +import { describe, it } from "node:test"; +import type { UIMessage } from "ai"; +import { neighborhood, withPlanBlocks, type PlanBlockPart } from "../src/core/plan-blocks.ts"; +import type { PlanTree } from "../src/core/contracts.ts"; + +/** The 5-node DAG the neighborhood tests slice: + * A → B → C, B → D → E (B's children C,D; A's only child is B). */ +function tree(): PlanTree { + return { + nodes: [ + { id: "a", title: "Foundations" }, + { id: "b", title: "Core" }, + { id: "c", title: "Applied" }, + { id: "d", title: "Deep dive" }, + { id: "e", title: "Frontier" }, + ], + edges: [ + { from: "a", to: "b" }, + { from: "b", to: "c" }, + { from: "b", to: "d" }, + { from: "d", to: "e" }, + ], + currentNodeId: null, + doneNodeIds: [], + version: 1, + }; +} + +function msg(id: string, role: "user" | "assistant", parts: UIMessage["parts"] = []): UIMessage { + return { id, role, parts }; +} + +function blocksOf(m: UIMessage): PlanBlockPart[] { + return m.parts.filter((p): p is PlanBlockPart => p.type === "data-plan-block") as PlanBlockPart[]; +} + +function allBlocks(messages: UIMessage[]): PlanBlockPart[] { + return messages.flatMap((m) => blocksOf(m)); +} + +/** tool-plan_init part in terminal state (the authoring call succeeded). */ +function planInitPart(tree: PlanTree, ok = true) { + return { type: "tool-plan_init", state: "output-available", input: { mermaid: "graph TD" }, output: { ok, tree } } as unknown as UIMessage["parts"][number]; +} + +/** tool-quiz part tagged with a node. `answered` appends the learner's output. */ +function quizPart(nodeId: string, answered?: { correct: boolean }) { + const part = { + type: "tool-quiz", + state: answered ? "output-available" : "input-available", + input: { nodeId, question: "…", options: [] }, + ...(answered ? { output: { correct: answered.correct } } : {}), + } as unknown as UIMessage["parts"][number]; + return part; +} + +describe("neighborhood (pure scope)", () => { + it("includes current + direct parents + direct children, nothing further", () => { + const { nodes, edges } = neighborhood(tree(), "b"); + assert.deepEqual( + nodes.map((n) => n.id).sort(), + ["a", "b", "c", "d"], + ); + assert.deepEqual(edges.map((e) => `${e.from}>${e.to}`).sort(), ["a>b", "b>c", "b>d"]); + }); +}); + +describe("withPlanBlocks", () => { + it("attaches a FULL block to the assistant message with a successful plan_init", () => { + const t = tree(); + const out = withPlanBlocks( + [msg("m1", "user", [{ type: "text", text: "teach me" } as never]), msg("m2", "assistant", [planInitPart(t)])], + t, + ); + const blocks = allBlocks(out); + assert.equal(blocks.length, 1); + assert.equal(blocks[0]!.data.kind, "full"); + assert.equal(blocks[0]!.data.current, "a"); // first prereq-free node + assert.equal(blocks[0]!.data.nodes.length, 5); + assert.equal(blocksOf(out[1]!).length, 1); // on the plan_init message itself + }); + + it("ignores a failed plan_init (ok: false)", () => { + const t = tree(); + const out = withPlanBlocks([msg("m1", "user"), msg("m2", "assistant", [planInitPart(t, false)])], t); + assert.equal(allBlocks(out).length, 0); + }); + + it("attaches a NEIGHBORHOOD block on the last assistant message when a check goes out", () => { + const t = tree(); + const out = withPlanBlocks( + [ + msg("m1", "user"), + msg("m2", "assistant", [quizPart("a", { correct: true })]), // a done earlier + msg("m3", "assistant", [quizPart("b")]), // check on b now out + ], + t, + ); + const blocks = allBlocks(out); + assert.equal(blocks.length, 1); + assert.equal(blocks[0]!.data.kind, "neighborhood"); + assert.equal(blocks[0]!.data.current, "b"); + assert.deepEqual( + blocks[0]!.data.nodes.map((n) => n.id).sort(), + ["a", "b", "c", "d"], // e excluded: grandchild of b via d + ); + assert.equal(blocksOf(out[2]!).length, 1); // on the LAST assistant message + }); + + it("does not re-emit when the plan state has not moved", () => { + const t = tree(); + const turn1 = withPlanBlocks([msg("m1", "user"), msg("m2", "assistant", [quizPart("b")])], t); + assert.equal(allBlocks(turn1).length, 1); + // next turn: plain exchange, plan untouched — the previous block IS the as-of slice + const turn2 = withPlanBlocks( + [...turn1, msg("m3", "user"), msg("m4", "assistant", [{ type: "text", text: "ok" } as never])], + t, + ); + assert.equal(allBlocks(turn2).length, 1); // nothing new + assert.equal(blocksOf(turn2[3]!).length, 0); + }); + + it("re-emits when a check resolves (done grows / current falls back)", () => { + const t = tree(); + const asked = withPlanBlocks([msg("m1", "user"), msg("m2", "assistant", [quizPart("b")])], t); + // learner answers correctly: done grows, no check out → current falls to a + const answered = withPlanBlocks( + [...asked, msg("m3", "user"), msg("m4", "assistant", [{ type: "text", text: "good" } as never])].map((m) => + m.id === "m2" ? { ...m, parts: m.parts.map((p) => (p.type === "tool-quiz" ? quizPart("b", { correct: true }) : p)) } : m, + ), + t, + ); + const blocks = allBlocks(answered); + assert.equal(blocks.length, 2); + const latest = blocks.at(-1)!.data; + assert.equal(latest.kind, "neighborhood"); + assert.deepEqual( + latest.nodes.find((n) => n.id === "b")!.status, + "done", + ); + assert.equal(latest.current, "a"); // b done, a never checked → first eligible + }); + + it("re-emits when a failed check resets the standing (no done growth)", () => { + const t = tree(); + const asked = withPlanBlocks([msg("m1", "user"), msg("m2", "assistant", [quizPart("b")])], t); + // learner answers WRONG: done doesn't grow, but no check is out anymore + // → current falls back to the first eligible node + const failed = withPlanBlocks( + [...asked, msg("m3", "user"), msg("m4", "assistant", [{ type: "text", text: "let's review" } as never])].map((m) => + m.id === "m2" ? { ...m, parts: m.parts.map((p) => (p.type === "tool-quiz" ? quizPart("b", { correct: false }) : p)) } : m, + ), + t, + ); + const blocks = allBlocks(failed); + assert.equal(blocks.length, 2); + const latest = blocks.at(-1)!.data; + assert.equal(latest.current, "a"); + assert.equal(latest.nodes.find((n) => n.id === "b")!.status, "upcoming"); // not done + }); + + it("is idempotent when the same list is persisted twice (retry replay)", () => { + const t = tree(); + const once = withPlanBlocks([msg("m1", "user"), msg("m2", "assistant", [quizPart("b")])], t); + const twice = withPlanBlocks(once, t); + assert.equal(allBlocks(twice).length, 1); + assert.equal(twice[1]!.parts.length, once[1]!.parts.length); + }); + + it("attaches a new FULL block on a plan rewrite, keeping the old block frozen", () => { + const t = tree(); + const authored = withPlanBlocks([msg("m1", "user"), msg("m2", "assistant", [planInitPart(t)])], t); + const rewrite = withPlanBlocks( + [...authored, msg("m3", "user"), msg("m4", "assistant", [planInitPart(t)])], + t, + ); + const blocks = allBlocks(rewrite); + assert.equal(blocks.length, 2); + assert.deepEqual(blocks.map((b) => b.data.kind), ["full", "full"]); + assert.equal(blocksOf(rewrite[3]!).length, 1); // new block on the NEW authoring message + }); + + it("no-ops on null/empty tree and on messages without assistant turns", () => { + assert.equal(withPlanBlocks([msg("m1", "user")], tree()).length, 1); + const empty = withPlanBlocks([msg("m1", "user"), msg("m2", "assistant")], { ...tree(), nodes: [] }); + assert.equal(allBlocks(empty).length, 0); + const noTree = withPlanBlocks([msg("m1", "user"), msg("m2", "assistant")], null); + assert.equal(allBlocks(noTree).length, 0); + }); +});