Swap the server engine to pi-ai/pi-agent-core (Band B+C of docs/pi-engine-swap.md)

The Vercel AI SDK survives only as Band D (useChat client + UIMessage
record types); @openrouter/ai-sdk-provider is uninstalled.

- src/agent/models.ts: PiEngine {model, streamFn} via createModels +
  openrouterProvider subpath; thinkingLevelFor maps config effort none->off
- src/agent/lh-tool.ts: neutral LhTool (zod authoring, optional execute,
  clientQuestion flag); toPiTool converts via z.toJSONSchema and re-runs
  zod inside execute so refinements ride the error-result retry path
- src/server/pi/record-adapter.ts: UIMessage records -> pi Messages
  (dangling tool calls get synthesized interrupted results)
- src/server/pi/ui-stream.ts: SSE bridge emitting AI-SDK-shaped chunks so
  the useChat client parses unchanged (continuation id rule, question
  tools never emit output, tool-input-error for bad authoring)
- teacher-agent.ts: stateless per-turn pi Agents; shouldStopAfterTurn
  stops only when every question-tool call resolved cleanly (errors retry)
- chat-handler.ts: buildAgent(outgoing) factory seam (stateless doctrine);
  assertValidUiMessages owns the retired SDK validation gate
- sub-agents via shared runSubagent (pi-run.ts); extraction stages stop
  after a clean submit; visual-maker keeps its hand-rolled loop over
  direct streamSimple calls; aborted model calls bill no usage
- tests: all model seams ported to pi-ai faux provider (scripted
  FauxResponseSteps, agent factories); handleChatRequest second param
  is now an agent factory

Gates: tsc, lint, build, 305/305. Live GLM smoke (disposable driver,
deleted): prober streaming + persistence, question seam, turn-2
continuation same message id, per-turn real token usage - ALL PASS.
Data wiped at the pivot (chats/graph/jobs, .gitkeeps kept).
This commit is contained in:
amosblomqvist
2026-09-12 02:34:18 +03:00
parent 377d4b6548
commit 2d717fd932
27 changed files with 3557 additions and 769 deletions
+278
View File
@@ -0,0 +1,278 @@
# The pi engine swap — design record
Date: 2026-09-12. Status: PLANNED (facts verified, awaiting go).
Decision (user, 2026-09-12): swap the server agent engine (Band B) and the
UI-stream bridge (Band C) from the Vercel AI SDK to **@earendil-works/pi-ai**
(provider/catalog/streaming layer) + **@earendil-works/pi-agent-core**
(agent loop). The client (Band D: `useChat`, transport, UI kit) **stays** on
the AI SDK. The transcript contract (Band A: UIMessage-shaped records,
`TeacherUIMessage`) **stays** as-is. Model/thinking controls come after the
swap, and the architecture below is shaped so they fall out of pi's catalog
naturally.
Motivation: SDK ambiguity has cost us repeatedly (data-part conventions,
UIMessage merge rules, retry paths, continuation coalescing — all took
forensics to understand). pi's layers are minimal and the user knows them
deeply. After this swap every piece of the turn pipeline is either ours or
pi's, each fully understood.
Verified package facts (all three scouts + npm registry):
- `@earendil-works/pi-ai` + `pi-agent-core` latest **0.85.1**, MIT, ESM-only,
Node `>=22.19.0` (local dev runs 26.7 — fine; Electron wave must ship ≥22).
Plain registry installs, no vendoring. The old prototype
(/Volumes/T7/code/learning-system) ran 0.84.2 cleanly with the same wiring.
- pi-ai = providers/catalog/streaming only: `createModels()`,
`models.streamSimple`, `calculateCost`, `getSupportedThinkingLevels`,
`validateToolArguments`. **No agent loop, no zod, no React, no SSE.**
- pi-agent-core = the loop: `Agent` (steer/followUp/abort/subscribe),
`agentLoop`, `AgentTool` (execute + `terminate`), JSON-Schema validation
that accepts plain JSON Schema (zod `z.toJSONSchema` output works, with
coercion — verified in `validation.js:285`).
## 1. Engine architecture — stateless per-turn Agent, record is the truth
**The record stays the source of truth.** Every turn constructs a *fresh*
`Agent` from the persisted UIMessage record via our `convertToLlm` adapter,
runs it, persists the resulting parts back into the record, discards the
Agent. pi's `steer()`/`followUp()` queues are **in-memory** (plain arrays on
the instance) and die with the process — they are for a *held-live* agent
(pi-coding-agent holds one Agent per session lifetime). lh's architecture is
the opposite: HTTP request-scoped turns, durable wake pipeline, write-ahead
persistence, 409 one-turn guard. Those stay exactly as they are.
Concretely, per POST to `/api/chat`:
1. `beginTurn` / 409 guard / write-ahead persist (unchanged, `chat-handler.ts`).
2. `claimSteers` coalescing (unchanged — server wake messages are appended
into the record *before* the turn, which is what already happens).
3. Build `initialState.messages` = `recordToPiMessages(record)` (§4).
4. `new Agent({ initialState, streamFn, convertToLlm, getApiKey, tools })`.
5. Run to completion (loop runs because question tools terminate early,
see §2; otherwise stopReason `stop`/`length` ends it).
6. Bridge (§3) emits SSE chunks while it runs; events collected for usage.
7. `onFinish` equivalent: `persistTurn` (plan-block splice, unchanged),
`appendTurnUsage`, `triggerExtraction`, `endTurn` (all unchanged).
What we deliberately do **not** adopt:
- Held-live Agents across requests (loses durability guarantees, changes the
409/one-writer model).
- pi's in-loop steering injection (mid-turn steers) — today a wake waits for
the running turn to end (claimSteers at request boundary). In-loop
injection via `getSteeringMessages` is a possible future enhancement, noted
in §8; not in this swap.
- pi's JSONL session machinery (`harness/session/jsonl`) — our record store
already plays that role.
## 2. The quiz/ask pattern — the heart of the swap
Three shapes were found across pi itself and the prototype; lh needs a third
variant because its question tools are **client-answered** and turn-scoped:
- pi's own `ask_question` (interactive-subagents ext): execute writes a
sidecar file and returns; the *session layer* orchestrates the reply as a
future message. Fits a held-live subagent, not our turn model.
- The prototype's plan: **blocking execute** awaiting the learner, graded
verdict as the result. Fits a TUI, not HTTP — would hold a POST open for
the minutes a learner takes to answer.
**lh's variant (the design):** question tools (`quiz`, `ask_user_question`)
are `AgentTool`s whose `execute()` returns *immediately* with a synthetic
result:
```ts
execute: async () => ({
content: [{ type: "text",
text: "Question posed to the learner. Turn ends here; the answer arrives as this call's result next turn." }],
})
```
- **The stop mechanism is the `shouldStopAfterTurn` hook, not
`terminate`.** (Verified in `agent-loop.js`: terminate lives on the
*result*, fires only after the batch executes, and only when *every*
result in the batch sets it — so a quiz+research batch would NOT stop,
the model would keep teaching past the quiz. `terminate` is also shaped
for workflow-final tools, not pause-pending-input; we don't repurpose
it.) Our wrapper's hook: *did this turn's completed assistant message
contain a question-tool call? → return true → the loop stops before any
further model call.* Deterministic for every batch shape; server tools
called in the same batch (research etc.) still execute and stream their
results first, matching today's SDK behavior where a client tool call
ends the turn.
- The synthetic result exists only because pi's loop requires a result for
every executed call to settle into a clean terminal state **within that
turn**. It is *never persisted* — our record keeps the tool part with
**no output** until the client answers, exactly as today. The dead run's
in-memory context (placeholder included) is discarded; the learner's
answer attaches to the *record*, and the next turn's fresh Agent sees
the answered call as a normal tool result via the adapter (§4).
- The learner's answer (dock → `addToolOutput` → next POST carrying the
mutated transcript) re-enters via `recordToPiMessages`: the answered tool
part becomes a `ToolResultMessage` with the original `toolCallId`, so the
model sees the answer as the tool result. This is precisely what pi's
`continue()` documents ("last message must be a user or tool-result
message") — we just feed it through our own record instead of a held Agent.
- **Authoring validation stays zod, stays ours.** `z.toJSONSchema` output
becomes the tool's `parameters` (pi validates structurally + coerces), but
our `superRefine` cross-field rules (correctAnswer∈options — the
twice-fixed wording) don't survive JSON Schema. So the wrapper runs the
zod schema itself at call time, before execution: on failure it pushes the
error back the same way pi's own validation does — an error
`ToolResultMessage` the model sees and retries against (verified: that's
pi's native self-healing path, `prepareToolCall` catch →
`createErrorToolResult`). Identical semantics to today's
`InvalidToolInputError`.
- **Dangling question calls** (turn aborted mid-run leaving an unanswered
quiz, or an aborted LLM call — pi salvages the partial assistant message
with stopReason `aborted`): `recordToPiMessages` must synthesize a
`ToolResultMessage` ("The question was interrupted") for any toolCall
lacking a result *before* the LLM call, or providers reject the request.
This is a hard invariant of the adapter, covered by tests.
## 3. The bridge — pi AgentEvents → UIMessageChunk SSE
We own a ~200-line `src/server/ui-stream-bridge.ts` replacing
`createAgentUIStreamResponse`. Wire format (verified from SDK source):
- Response: `text/event-stream`, chunks as flat JSON
`data: {"type":"text-delta","id":"…","delta":"…"}\n\n`, terminated with
`data: [DONE]\n\n`. Vercel headers not required (client ignores them);
we send `content-type` + `cache-control: no-cache` anyway.
- Every chunk zod-validated client-side (looseObjects); malformed required
fields put the chat in error status. Our chunk construction is typed and
unit-tested against `uiMessageChunkSchema` shapes.
Event mapping (pi event → chunk(s)):
- `agent_start`/`turn_start` → `start` (see id rule) + `start-step`.
- `message_update` carrying pi text events: `text_start/delta/end` →
`text-start`/`text-delta`/`text-end` (stable id per block, e.g.
`${messageId}-t${n}`).
- thinking events → `reasoning-start`/`reasoning-delta`/`reasoning-end`.
- `toolcall_start/delta/end` → `tool-input-start`/`tool-input-delta`
(optional; safe to skip deltas and emit only `tool-input-available` on
end — verified clean path)/`tool-input-available` (part state
`input-available`, static name `tool-${toolName}`).
- `tool_execution_start/end` (server-executed tools only: research, demo,
visual, plan_init, probe_done) → `tool-output-available` with the
serialized result. **Never** emit outputs for question tools (client-owned;
emitting for an unknown callId throws client-side → error status).
- `tool_execution_end` carrying validation-error results → same
`tool-output-available` (error content), matching today's behavior.
- `message_end` → collect usage (§5); `turn_end` → `finish-step`.
- `agent_end` → `finish {finishReason}` — send it always (feeds `onFinish`
fidelity); `stop`/`length` map naturally.
- Errors mid-run: emit `error {errorText}` (does not end the stream; the
client throws it into error status), then close. Unmask
validation/unknown-tool messages like today.
- **Continuation id rule (the load-bearing one, from client source
`makeRequest`/`createStreamingUIMessageState`):** if the transcript's last
message is an *assistant* message, the client's streaming state starts as
a clone of it and the stream must continue that message — so `start`
carries `messageId` = last assistant message's id. Otherwise (fresh turn)
a fresh server-generated id. Violating this duplicates every part of the
prior assistant message (question card twice, duplicated toolCallIds).
- **Abort:** the client's Stop button aborts the fetch → Node `request.signal`
fires. The bridge subscribes: `agent.abort()`, salvage path persists the
partial (pi keeps the partial with stopReason `aborted`), same
isAborted persistence as today. Job/extraction aborts (`abortJob`,
`abortExtraction`) call `agent.abort()` on their registered runs — the
signal threads into `streamFn` and every `execute()`, verified.
## 4. The record adapter — `recordToPiMessages` / message shapes
`convertToLlm` (ours) converts record UIMessages → pi `Message[]`:
- user text message → `UserMessage {role:"user", content, timestamp}`.
Data parts (`data-plan-accepted`, `data-plan-block`) are dropped
(UI-only), exactly as `convertToModelMessages` does today.
- assistant message → `AssistantMessage {role:"assistant", content:[…]}` —
text parts → text blocks; reasoning parts → thinking blocks (+
`thinkingSignature` persisted in the record for Anthropic resume);
tool parts → `ToolCall {id, name, arguments, thoughtSignature?}`.
- answered tool part → after its assistant message, a
`ToolResultMessage {role:"toolResult", toolCallId, toolName, content,
isError?, timestamp}` built from the part's output. Question-tool outputs
become plain text content (the graded-verdict payload, as the model sees
it today).
- Unanswered question toolCall → synthesized interrupted-result (§2).
- `data-plan-accepted` on a user message → text "The plan is accepted —
begin teaching." rides the message (today: `data-plan-accepted` part +
the text part; the model only ever saw the text — identical).
- Prober/teacher seam unchanged: the teacher branch still keys on the
`tool-probe_done` part in the transcript; both agents are constructed per
turn from the same record.
The record format itself does **not** change this wave (Band A untouched);
thinkingSignature support is added to the part schema as one new optional
field (used only when a provider returns signed thinking).
## 5. Usage, cost, and the model catalog
- Usage: subscribe to `message_end`/`turn_end` events →
`AssistantMessage.usage {input, output, cacheRead, cacheWrite, reasoning?,
totalTokens, cost}` → mapped into our `CallUsage`/`TurnUsage` contracts
(field renames only: `reasoningTokens` → `reasoning` etc.). The stats bar
and activity dock are untouched.
- Cost: pi's `Usage.cost` comes pre-computed via `calculateCost` (tiered).
**This wave keeps `core/prices.ts`** (tested folds, synced OpenRouter
table) — no behavior change. Replacing prices.ts with pi's catalog cost
is a noted cleanup for the model-control wave, where the catalog becomes
load-bearing anyway.
- Catalog: `createModels()` + `openrouterProvider()`; our working set
(`z-ai/glm-5.3-flash`) is **registered programmatically** (id, provider,
api, cost/contextWindow/thinkingLevelMap from our data) because it's not
in pi's static catalog (verified; it is in the pi.dev overlay, but we
don't take a runtime pi.dev dependency). `getApiKey: provider → key` from
env (`OPENROUTER_API_KEY`) now; the wave after swaps the source to
`data/auth.json` without touching this seam.
## 6. Sub-agents
All five sub-agents (research runner, demo-builder, visual-maker, extractor,
auditor) port from hand-rolled `generateText`/ToolLoopAgent loops to
`new Agent` or `agentLoop` with the same tools+streamFn wiring. Abort
signals thread through (`execute`'s 3rd arg), usage folds from the same
events, `calculateCost` available for per-stage cost. Mechanical; largest
surface is the extractor's staged loop.
## 7. Test harness
`MockLanguageModelV4` (23 sites) is replaced by scripted pi-side fakes:
- Engine-level: a fake `streamFn` yielding scripted
`AssistantMessageEvent` sequences (text, tool calls, terminal) — the
async-iterator shape is simpler than the SDK's.
- Bridge-level: golden-chunk tests asserting our SSE output against the
verified chunk shapes (continuation id rule, terminate paths, dangling
synthesis, abort salvage).
- Adapter-level: `recordToPiMessages` round-trips (answers, data parts,
aborted quizzes, plan-accepted).
## 8. Waves, gates, and deferred items
One atomic wave (the engine can't ship half-swapped): deps + Next
`serverExternalPackages` for pi-ai/pi-agent-core → adapters (§4) → tools
port (§2, zod kept) → bridge (§3) → chat-handler rewiring → sub-agents (§6)
→ test port (§7) → gates (tsc, full suite, lint, build) → **data wiped at
the pivot** (doctrine; record format is unchanged, so the wipe is doctrine
not necessity) → live GLM smoke: full probe→assessment→handoff→plan→
verdict→teach flow plus abort paths (question mid-stream abort, job abort,
extraction abort), zero console errors, usage correct in the stats bar.
Deferred (explicitly): in-loop steering injection (mid-turn wakes);
prices.ts → pi catalog cost migration; client-band (Band D) ownership
decisions; multi-provider auth beyond OpenRouter (the `getApiKey`/catalog
seams above are the whole preparation); pi.dev overlay refresh.
## 9. Open flags carried into implementation
- Chunk `id`s for text/reasoning blocks must be unique per message and
stable across continuation turns (the client throws on a delta for an
inactive id).
- `tool-input-available` requires `input` as a JSON value; pi gives parsed
arguments on toolcall end.
- Step bookkeeping: pi `turn` ≈ SDK `step`; our per-step boundaries must
keep `step-start` parts where the plan-block splice and phase timelines
expect them (splice keys on tool-plan_init parts — unaffected).
- Next must not bundle pi-ai's provider deps (AWS/Google SDKs) into the
route bundles — `serverExternalPackages`.
- pi-agent-core is single-run per Agent (`prompt()` throws if active):
fresh-Agent-per-turn sidesteps it entirely (and matches our one-writer
turn guard).
+3
View File
@@ -2,6 +2,9 @@ import type { NextConfig } from "next";
const nextConfig: NextConfig = { const nextConfig: NextConfig = {
// Framework-free core lives in src/ (plain TS, imported by routes and CLI alike). // Framework-free core lives in src/ (plain TS, imported by routes and CLI alike).
// The pi engine must stay external: its providers are runtime-loaded Node
// code (undici streams, catalog data) that breaks when bundled into routes.
serverExternalPackages: ["@earendil-works/pi-ai", "@earendil-works/pi-agent-core"],
}; };
export default nextConfig; export default nextConfig;
+1622 -13
View File
File diff suppressed because it is too large Load Diff
+2 -1
View File
@@ -14,8 +14,9 @@
}, },
"dependencies": { "dependencies": {
"@ai-sdk/react": "^4.0.79", "@ai-sdk/react": "^4.0.79",
"@earendil-works/pi-agent-core": "^0.85.1",
"@earendil-works/pi-ai": "^0.85.1",
"@mozilla/readability": "^0.6.0", "@mozilla/readability": "^0.6.0",
"@openrouter/ai-sdk-provider": "^3.0.0",
"@radix-ui/react-use-controllable-state": "^1.2.6", "@radix-ui/react-use-controllable-state": "^1.2.6",
"@rive-app/react-webgl2": "^4.32.1", "@rive-app/react-webgl2": "^4.32.1",
"@streamdown/cjk": "^1.0.3", "@streamdown/cjk": "^1.0.3",
+109
View File
@@ -0,0 +1,109 @@
import { z } from "zod";
import type { AgentTool } from "@earendil-works/pi-agent-core";
/**
* lh's neutral tool definition + the converter to pi-agent-core's AgentTool
* (docs/pi-engine-swap.md §3). The agent-band swap must not re-risk the
* battle-tested zod authoring contracts (the quiz correctAnswer wording,
* the nodeId superRefines) — so tools are AUTHORED exactly as before
* (zod schemas, same execute signatures) in this shape, and converted
* once, centrally, here.
*
* What the converter owns (the semantics the Vercel SDK gave implicitly):
* - parameters: z.toJSONSchema feeds the model-facing declaration (pi
* structurally validates plain JSON Schema itself).
* - refinements: superRefine rules don't survive JSON Schema — the schema
* is re-run inside execute, and a failure THROWS. pi's loop catches
* execute throws and turns them into an error tool result the model
* sees and retries on — the exact InvalidToolInputError semantics.
* - client question tools (quiz / ask_user_question): execute returns a
* synthetic "question posed" placeholder that is NEVER persisted — the
* record keeps the part output-less until the client answers; the
* learner's answer re-enters the next turn as the tool result. The
* turn itself stops via shouldStopAfterTurn (teacher-agent.ts).
*/
export interface LhTool<TInput = any, TResult = any> {
name: string;
description: string;
/** The authoring contract (zod, refinements included). */
schema: z.ZodType<TInput>;
/**
* Server-side execution. Absent + clientQuestion = the client answers.
* Throws on validation failure — the loop converts to an error result.
*/
execute?: (input: TInput, ctx: { toolCallId: string }) => TResult | Promise<TResult>;
/** Client-answered question tool (quiz/ask): a call ends the turn. */
clientQuestion?: boolean;
}
/** Author a tool (the `tool({...})` shape, engine-neutral). */
export function lhTool<TInput, TResult = unknown>(def: LhTool<TInput, TResult>): LhTool<TInput, TResult> {
return def;
}
/** lh's port of the SDK's invalid-input message: what the model retried on. */
export class ToolInputError extends Error {
constructor(
public readonly toolName: string,
public readonly zodError: z.ZodError,
) {
super(
`Invalid tool input for tool "${toolName}":\n${zodError.issues
.map((issue) => {
const path = issue.path.length > 0 ? ` at ${issue.path.join(".")}` : "";
return `- ${issue.message}${path}`;
})
.join("\n")}`,
);
this.name = "ToolInputError";
}
}
/** Convert one lh tool to a pi AgentTool. */
export function toPiTool(def: LhTool<any, any>): AgentTool<any, any> {
return {
name: def.name,
label: def.name,
description: def.description,
parameters: z.toJSONSchema(def.schema as z.ZodType) as never,
execute: async (toolCallId, args) => {
// The authoring check: structural JSON Schema validation already
// ran in the loop's prepare step; this re-run adds the refinements
// (superRefine rules) the JSON Schema can't carry.
const parsed = def.schema.safeParse(args);
if (!parsed.success) throw new ToolInputError(def.name, parsed.error);
if (def.clientQuestion || !def.execute) {
// The synthetic placeholder — in-memory only; the record's tool
// part keeps no output until the learner answers.
return {
content: [{ type: "text", text: "Question posed — the turn ends here; the learner's graded answer arrives as this call's result." }],
details: null,
};
}
const result = await def.execute(parsed.data, { toolCallId });
return {
// The model's view: today's parity — string results pass verbatim,
// objects as JSON.
content: [{ type: "text", text: typeof result === "string" ? result : JSON.stringify(result) }],
details: result,
};
},
};
}
/** Convert a tool registry (name → LhTool) to pi's AgentTool[]. */
export function toPiTools(registry: Record<string, LhTool<any, any>>): AgentTool<any, any>[] {
return Object.entries(registry).map(([name, def]) => toPiTool({ ...def, name }));
}
/** Declaration-only pi Tools (name/description/parameters, no execute) —
* for hand-rolled loops that own tool execution themselves. */
export function piToolDeclarations(registry: Record<string, LhTool<any, any>>) {
return Object.entries(registry).map(([name, def]) => ({
name,
label: name,
description: def.description,
parameters: z.toJSONSchema(def.schema as z.ZodType) as never,
}));
}
+67 -11
View File
@@ -1,17 +1,73 @@
import { openrouter } from "@openrouter/ai-sdk-provider"; import {
import type { LanguageModel } from "ai"; createModels,
type Api,
type Model,
type Models,
type StreamFunction,
} from "@earendil-works/pi-ai";
import { openrouterProvider } from "@earendil-works/pi-ai/providers/openrouter";
import type { ThinkingLevel } from "@earendil-works/pi-agent-core";
import { MODELS, type ModelRole } from "../../config.ts"; import { MODELS, type ModelRole } from "../../config.ts";
/** /**
* The one factory every agent's model comes from. Model choice and thinking * The one engine factory every agent's model comes from (the pi-ai swap,
* budget live in config.ts (MODELS) — the single tuning table; this module * docs/pi-engine-swap.md). Model choice and thinking budget live in
* only wires them to OpenRouter. `effort: "none"` is OpenRouter's off-switch * config.ts (MODELS) — the single tuning table; this module only wires them
* for reasoning tokens; higher levels let the model think longer (and bill * to pi-ai's OpenRouter provider.
* accordingly). *
* An "engine" is the model + the stream function the agent loop runs on —
* the two things a test fake replaces. Production engines resolve the model
* from pi-ai's static catalog (glm-5.3-flash and qwen3.8-max are both in
* 0.85.1's built-ins — no overlay fetch) and stream through the same Models
* collection.
*/ */
export function modelFor(role: ModelRole): LanguageModel {
export interface PiEngine {
model: Model<Api>;
streamFn: StreamFunction;
}
export type { ThinkingLevel };
/** config's effort scale maps 1:1 onto pi's thinking levels (none→off). */
export function thinkingLevelFor(role: ModelRole): ThinkingLevel {
const effort = MODELS[role].effort as string;
return effort === "none" ? "off" : (effort as ThinkingLevel);
}
let collection: ReturnType<typeof createModels> | undefined;
/** The process-wide pi-ai Models collection (OpenRouter catalog). */
export function piModels() {
if (!collection) {
collection = createModels();
collection.setProvider(openrouterProvider());
}
return collection;
}
/** OpenRouter API key (BYOK); tests never read this — they inject engines. */
export function openRouterApiKey(): string | undefined {
return process.env.OPENROUTER_API_KEY || undefined;
}
/** The production engine for one agent role (config.ts MODELS). */
export function piEngineFor(role: ModelRole): PiEngine {
const models: Models = piModels();
const spec = MODELS[role]; const spec = MODELS[role];
return openrouter(spec.model, { const model = models.getModel("openrouter", spec.model);
reasoning: { enabled: true, effort: spec.effort }, if (!model) throw new Error(`OpenRouter catalog has no model ${spec.model} (role ${role})`);
}); const level = thinkingLevelFor(role);
return {
model,
// bound streamSimple: the agent loop forwards state.thinkingLevel as
// `reasoning` (undefined = off on pi-ai's scale); pin the role's level
// so a bare engine can never drift.
streamFn: (requestModel, context, options) =>
models.streamSimple(requestModel, context, {
...options,
reasoning: level === "off" ? undefined : level,
apiKey: options?.apiKey ?? openRouterApiKey(),
}),
};
} }
+23 -22
View File
@@ -1,7 +1,9 @@
import { ToolLoopAgent, stepCountIs, tool, type LanguageModel, type ToolSet } from "ai";
import { z } from "zod"; import { z } from "zod";
import { modelFor } from "../models.ts"; import type { Usage } from "@earendil-works/pi-ai";
import type { DemoResult, UsageStepLike } from "../../core/contracts.ts"; import { lhTool, toPiTools, type LhTool } from "../lh-tool.ts";
import type { PiEngine } from "../models.ts";
import { runSubagent } from "./pi-run.ts";
import type { DemoResult } from "../../core/contracts.ts";
import { getHeadless, type DemoRunResult, type HeadlessApi } from "./headless.ts"; import { getHeadless, type DemoRunResult, type HeadlessApi } from "./headless.ts";
/** /**
@@ -11,7 +13,7 @@ import { getHeadless, type DemoRunResult, type HeadlessApi } from "./headless.ts
* fix → re-run until clean — before anything reaches the lesson. * fix → re-run until clean — before anything reaches the lesson.
* *
* Shape: a ToolLoopAgent with two executable tools (write_demo / run_demo) — * Shape: a ToolLoopAgent with two executable tools (write_demo / run_demo) —
* the model is injectable (tests use MockLanguageModelV4) and the headless * the engine is injectable (tests use a scripted faux engine) and the headless
* API is injectable (tests use a fake; production uses system Chrome). * API is injectable (tests use a fake; production uses system Chrome).
* *
* Failure semantics (decided in the W6 brief): `tested: true` means "the * Failure semantics (decided in the W6 brief): `tested: true` means "the
@@ -26,7 +28,8 @@ export interface DemoBuilderInput {
} }
export interface DemoBuilderOptions { export interface DemoBuilderOptions {
model?: LanguageModel; /** Test seam: scripted engine; production uses the demoBuilder role. */
engine?: PiEngine;
headless?: HeadlessApi; headless?: HeadlessApi;
/** Max agent steps (each write/run/fix is one step). Default 24 ≈ 6 cycles. */ /** Max agent steps (each write/run/fix is one step). Default 24 ≈ 6 cycles. */
maxSteps?: number; maxSteps?: number;
@@ -35,7 +38,7 @@ export interface DemoBuilderOptions {
/** Kills the run mid-model-call (2026-09 real-kill abort semantics). */ /** Kills the run mid-model-call (2026-09 real-kill abort semantics). */
abortSignal?: AbortSignal; abortSignal?: AbortSignal;
/** Money channel: one call per agent step, folded into JobRecord.usage. */ /** Money channel: one call per agent step, folded into JobRecord.usage. */
onUsage?: (usage: UsageStepLike) => void; onUsage?: (usage: Usage) => void;
} }
const INSTRUCTIONS = `You are a demo builder — a specialist sub-agent inside a learning application. You const INSTRUCTIONS = `You are a demo builder — a specialist sub-agent inside a learning application. You
@@ -78,13 +81,14 @@ export function createDemoBuilderTools(
state: BuilderState, state: BuilderState,
headless: HeadlessApi, headless: HeadlessApi,
onActivity?: (line: string) => void, onActivity?: (line: string) => void,
): ToolSet { ): Record<string, LhTool<any, any>> {
return { return {
write_demo: tool({ write_demo: lhTool({
name: "write_demo",
description: description:
"Write the complete self-contained HTML document for the demo (overwrites any previous version). " + "Write the complete self-contained HTML document for the demo (overwrites any previous version). " +
"Always follow with run_demo to test it headlessly.", "Always follow with run_demo to test it headlessly.",
inputSchema: z.object({ schema: z.object({
html: z.string().describe("The complete HTML document: <!doctype html>…</html>, all JS inline."), html: z.string().describe("The complete HTML document: <!doctype html>…</html>, all JS inline."),
}), }),
execute: async ({ html }) => { execute: async ({ html }) => {
@@ -94,11 +98,12 @@ export function createDemoBuilderTools(
return `Saved (${html.length} bytes). Now call run_demo to test it — never deliver untested.`; return `Saved (${html.length} bytes). Now call run_demo to test it — never deliver untested.`;
}, },
}), }),
run_demo: tool({ run_demo: lhTool({
name: "run_demo",
description: description:
"Headlessly execute the current demo in Chrome and report console errors, uncaught errors, and failed requests. " + "Headlessly execute the current demo in Chrome and report console errors, uncaught errors, and failed requests. " +
"Optionally pass a probe expression to assert on page state.", "Optionally pass a probe expression to assert on page state.",
inputSchema: z.object({ schema: z.object({
probe: z probe: z
.string() .string()
.optional() .optional()
@@ -174,12 +179,6 @@ export async function runDemoBuilder(input: DemoBuilderInput, opts: DemoBuilderO
} }
const state: BuilderState = { html: null, testedHtml: null, lastResult: null }; const state: BuilderState = { html: null, testedHtml: null, lastResult: null };
const agent = new ToolLoopAgent({
model: opts.model ?? modelFor("demoBuilder"),
instructions: INSTRUCTIONS,
tools: createDemoBuilderTools(state, headless, opts.onActivity),
stopWhen: stepCountIs(opts.maxSteps ?? 24),
});
const prompt = [ const prompt = [
`Demo title: ${title}`, `Demo title: ${title}`,
@@ -191,13 +190,15 @@ export async function runDemoBuilder(input: DemoBuilderInput, opts: DemoBuilderO
let finalText = ""; let finalText = "";
try { try {
const result = await agent.generate({ const result = await runSubagent(toPiTools(createDemoBuilderTools(state, headless, opts.onActivity)), {
role: "demoBuilder",
engine: opts.engine,
systemPrompt: INSTRUCTIONS,
prompt, prompt,
maxTurns: opts.maxSteps ?? 24,
abortSignal: opts.abortSignal, abortSignal: opts.abortSignal,
onStepFinish: ({ usage }) => opts.onUsage?.(usage), onUsage: (usage) => opts.onUsage?.(usage),
onToolExecutionStart: ({ toolCall }) => { onToolStart: (toolName) => opts.onActivity?.(toolName),
opts.onActivity?.(toolCall.toolName);
},
}); });
finalText = result.text; finalText = result.text;
} catch (err) { } catch (err) {
+114 -75
View File
@@ -1,6 +1,9 @@
import { ToolLoopAgent, stepCountIs, tool, type LanguageModel, type UIMessage } from "ai"; import type { UIMessage } from "ai";
import type { Usage } from "@earendil-works/pi-ai";
import { z } from "zod"; import { z } from "zod";
import { modelFor } from "../models.ts"; import { lhTool, toPiTools } from "../lh-tool.ts";
import type { PiEngine } from "../models.ts";
import { runSubagent } from "./pi-run.ts";
import type { import type {
AuditResult, AuditResult,
CallUsage, CallUsage,
@@ -207,19 +210,50 @@ transcript). The auditor rejects anything unbacked. When in doubt, report; the a
Your entire deliverable is ONE submit_findings call.`; Your entire deliverable is ONE submit_findings call.`;
function createExtractor(model: LanguageModel) { /** One extractor run: the whole deliverable is ONE zod-validated submit
return new ToolLoopAgent({ * call (a pi agent purely so the model can retry a schema-invalid call).
model, * Returns the submitted findings — captured from the execute seam, which
instructions: EXTRACTOR_INSTRUCTIONS, * only fires on a schema-valid call. */
tools: { async function runExtractor(
submit_findings: tool({ engine: PiEngine | undefined,
prompt: string,
signal: AbortSignal | undefined,
onStep: (usage: Usage, model: string, step: number, finishReason: string, tools: string) => void,
): Promise<SessionFindings> {
let submitted: SessionFindings | undefined;
const result = await runSubagent(
toPiTools({
submit_findings: lhTool({
name: "submit_findings",
description: "Submit the session findings.", description: "Submit the session findings.",
inputSchema: findingsSchema, schema: findingsSchema,
execute: async (findings) => findings, execute: async (findings) => {
}), submitted = findings;
return findings;
}, },
stopWhen: stepCountIs(4), }),
}); }),
{
role: "extractor",
engine,
systemPrompt: EXTRACTOR_INSTRUCTIONS,
prompt,
maxTurns: 4,
abortSignal: signal,
onAssistantMessage: (message, step) =>
onStep(
message.usage,
message.model,
step,
message.stopReason,
message.content.filter((c) => c.type === "toolCall").map((c) => (c as { name: string }).name).join(", ") || "(none)",
),
shouldStop: (_message, toolResults) => toolResults.some((r) => r.toolName === "submit_findings" && !r.isError),
},
);
void result;
if (!submitted) throw new Error("the extractor never submitted findings");
return submitted;
} }
// ─── the auditor ─────────────────────────────────────────────────────────── // ─── the auditor ───────────────────────────────────────────────────────────
@@ -290,19 +324,47 @@ function graphView(nodes: GraphNode[]): string {
return lines.join("\n"); return lines.join("\n");
} }
function createAuditor(model: LanguageModel) { /** One auditor run — same single-submit shape as the extractor. */
return new ToolLoopAgent({ async function runAuditor(
model, engine: PiEngine | undefined,
instructions: AUDITOR_INSTRUCTIONS, prompt: string,
tools: { signal: AbortSignal | undefined,
submit_review: tool({ onStep: (usage: Usage, model: string, step: number, finishReason: string, tools: string) => void,
): Promise<AuditResult> {
let submitted: AuditResult | undefined;
const result = await runSubagent(
toPiTools({
submit_review: lhTool({
name: "submit_review",
description: "Submit the verified findings subset (with identity verdicts) + the dropped list.", description: "Submit the verified findings subset (with identity verdicts) + the dropped list.",
inputSchema: auditSchema, schema: auditSchema,
execute: async (review) => review, execute: async (review) => {
}), submitted = review;
return review;
}, },
stopWhen: stepCountIs(4), }),
}); }),
{
role: "auditor",
engine,
systemPrompt: AUDITOR_INSTRUCTIONS,
prompt,
maxTurns: 4,
abortSignal: signal,
onAssistantMessage: (message, step) =>
onStep(
message.usage,
message.model,
step,
message.stopReason,
message.content.filter((c) => c.type === "toolCall").map((c) => (c as { name: string }).name).join(", ") || "(none)",
),
shouldStop: (_message, toolResults) => toolResults.some((r) => r.toolName === "submit_review" && !r.isError),
},
);
void result;
if (!submitted) throw new Error("the auditor never submitted a review");
return submitted;
} }
// ─── the pipeline status registry ───────────────────────────────────────────── // ─── the pipeline status registry ─────────────────────────────────────────────
@@ -346,10 +408,9 @@ export function abortExtraction(chatId: string): boolean {
// ─── the pipeline ─────────────────────────────────────────────────────────── // ─── the pipeline ───────────────────────────────────────────────────────────
export interface ExtractionModels { export interface ExtractionModels {
/** The extractor's model (tests inject a MockLanguageModelV4). */ /** Test seam: a scripted engine per stage (production uses the roles). */
extractor?: LanguageModel; extractor?: PiEngine;
/** The auditor's model. */ auditor?: PiEngine;
auditor?: LanguageModel;
} }
export interface ExtractionPipelineDeps extends GraphDeps, ExtractionModels {} export interface ExtractionPipelineDeps extends GraphDeps, ExtractionModels {}
@@ -362,8 +423,6 @@ export interface ExtractionPipeline {
export function createExtractionPipeline(deps: ExtractionPipelineDeps): ExtractionPipeline { export function createExtractionPipeline(deps: ExtractionPipelineDeps): ExtractionPipeline {
const { chatStore, graphStore } = deps; const { chatStore, graphStore } = deps;
const extractorAgent = createExtractor(deps.extractor ?? modelFor("extractor"));
const auditorAgent = createAuditor(deps.auditor ?? modelFor("auditor"));
const inFlight = new Set<string>(); const inFlight = new Set<string>();
function setPhase(chatId: string, phase: ExtractionPhase): void { function setPhase(chatId: string, phase: ExtractionPhase): void {
@@ -385,22 +444,17 @@ export function createExtractionPipeline(deps: ExtractionPipelineDeps): Extracti
run.usage[stage]!.steps.push(usage); run.usage[stage]!.steps.push(usage);
} }
/** The SDK's per-step usage, flattened to the persisted CallUsage fields /** pi's per-step usage, flattened to the persisted CallUsage fields
* (the money channel's capture shape, contracts/chat.ts). */ * (the money channel's capture shape, contracts/chat.ts). */
function callUsage(usage: { function callUsage(usage: Usage): CallUsage {
inputTokens: number | undefined;
outputTokens: number | undefined;
inputTokenDetails?: { cacheReadTokens?: number; cacheWriteTokens?: number };
outputTokenDetails?: { reasoningTokens?: number };
}): CallUsage {
// Optional-chained: the SDK types the detail objects as required, but
// providers (and mocks) may leave them unset.
return { return {
inputTokens: usage.inputTokens, inputTokens: usage.input,
outputTokens: usage.outputTokens, outputTokens: usage.output,
cacheReadTokens: usage.inputTokenDetails?.cacheReadTokens, cacheReadTokens: usage.cacheRead,
cacheWriteTokens: usage.inputTokenDetails?.cacheWriteTokens, cacheWriteTokens: usage.cacheWrite,
reasoningTokens: usage.outputTokenDetails?.reasoningTokens, // Optional in pi's Usage (providers that never emit reasoning
// tokens leave it unset) — persisted as 0, never dropped.
reasoningTokens: usage.reasoning ?? 0,
}; };
} }
@@ -413,23 +467,17 @@ export function createExtractionPipeline(deps: ExtractionPipelineDeps): Extracti
console.info(`[extraction ${chatId}] extractor starting (${view.split("\n").length} transcript lines)`); console.info(`[extraction ${chatId}] extractor starting (${view.split("\n").length} transcript lines)`);
const t0 = performance.now(); const t0 = performance.now();
const extraction = await extractorAgent.generate({ const findings = await runExtractor(
prompt: `Transcript:\n\n${view}`, deps.extractor,
abortSignal: signal, `Transcript:\n\n${view}`,
onStepFinish: (step) => { signal,
recordStep(chatId, "extractor", callUsage(step.usage), step.model.modelId); (usage, model, step, finishReason, tools) => {
console.info( recordStep(chatId, "extractor", callUsage(usage), model);
`[extraction] extractor step ${step.stepNumber}: ${step.finishReason}, tools: ${step.toolCalls.map((c) => c.toolName).join(", ") || "(none)"}`, console.info(`[extraction] extractor step ${step}: ${finishReason}, tools: ${tools}`);
);
}, },
}); );
for (const part of extraction.content) {
if (part.type === "tool-result" && part.toolName === "submit_findings") {
console.info(`[extraction ${chatId}] extractor done in ${Math.round(performance.now() - t0)}ms`); console.info(`[extraction ${chatId}] extractor done in ${Math.round(performance.now() - t0)}ms`);
return part.output as SessionFindings; return findings;
}
}
throw new Error("the extractor never submitted findings");
} }
async function audit(chatId: string, findings: SessionFindings): Promise<AuditResult> { async function audit(chatId: string, findings: SessionFindings): Promise<AuditResult> {
@@ -445,25 +493,16 @@ export function createExtractionPipeline(deps: ExtractionPipelineDeps): Extracti
); );
const t0 = performance.now(); const t0 = performance.now();
const review = await auditorAgent.generate({ const auditResult = await runAuditor(
prompt: deps.auditor,
`Transcript:\n\n${view}\n\nProposed findings:\n\n${JSON.stringify(findings, null, 2)}\n\n` + `Transcript:\n\n${view}\n\nProposed findings:\n\n${JSON.stringify(findings, null, 2)}\n\n` +
`Current knowledge graph (${graph.nodes.length} nodes):\n\n${graphView(graph.nodes)}`, `Current knowledge graph (${graph.nodes.length} nodes):\n\n${graphView(graph.nodes)}`,
abortSignal: signal, signal,
onStepFinish: (step) => { (usage, model, step, finishReason, tools) => {
recordStep(chatId, "auditor", callUsage(step.usage), step.model.modelId); recordStep(chatId, "auditor", callUsage(usage), model);
console.info( console.info(`[extraction] auditor step ${step}: ${finishReason}, tools: ${tools}`);
`[extraction] auditor step ${step.stepNumber}: ${step.finishReason}, tools: ${step.toolCalls.map((c) => c.toolName).join(", ") || "(none)"}`,
);
}, },
}); );
let auditResult: AuditResult | null = null;
for (const part of review.content) {
if (part.type === "tool-result" && part.toolName === "submit_review") {
auditResult = part.output as AuditResult;
}
}
if (!auditResult) throw new Error("the auditor never submitted a review");
if (auditResult.dropped.length > 0) { if (auditResult.dropped.length > 0) {
console.info( console.info(
`[extraction ${chatId}] audit dropped ${auditResult.dropped.length}: ` + `[extraction ${chatId}] audit dropped ${auditResult.dropped.length}: ` +
+130
View File
@@ -0,0 +1,130 @@
import { Agent, type AgentEvent, type AgentMessage, type AgentTool } from "@earendil-works/pi-agent-core";
import type { AssistantMessage, ToolResultMessage, Usage } from "@earendil-works/pi-ai";
import { openRouterApiKey, piEngineFor, thinkingLevelFor, type PiEngine } from "../models.ts";
import type { ModelRole } from "../../../config.ts";
/**
* The shared sub-agent runner (docs/pi-engine-swap.md §6): one helper every
* background agent (researcher, demo-builder, visual-maker, extractor,
* auditor) runs on — a stateless pi Agent built per call, driven to
* completion, with the seams the job runner and the money channel need.
*
* Error semantics mirror the old ToolLoopAgent.generate: provider/loop
* failures THROW (pi encodes them in the stream; this helper surfaces
* them), and an abort throws an AbortError so every existing
* `signal.aborted` early-return keeps working unchanged.
*/
export interface SubagentRunOptions {
/** The role whose engine (model + thinking level) this run uses. */
role: ModelRole;
/** Test seam: a scripted engine replaces the role's production one. */
engine?: PiEngine;
systemPrompt: string;
/** Initial transcript (continuation semantics: ends user / toolResult). */
messages?: AgentMessage[];
/** Fresh-run prompt text (mutually exclusive with a continuation). */
prompt?: string;
/** Max model calls (the old stepCountIs safety net). */
maxTurns: number;
/** Abort mid-model-call (the real-kill semantics). */
abortSignal?: AbortSignal;
/** Money channel: one call per model request. */
onUsage?: (usage: Usage) => void;
/** Activity line per tool execution (job record progress). */
onToolStart?: (toolName: string, args: unknown) => void;
/** Per completed assistant message (extraction's step log). */
onAssistantMessage?: (message: AssistantMessage, step: number) => void;
/** Extra stop rule (extractor: stop after a clean submit). */
shouldStop?: (message: AssistantMessage, toolResults: ToolResultMessage[]) => boolean;
}
export interface SubagentRunResult {
/** The messages this run added (assistant + tool results). */
newMessages: AgentMessage[];
/** The final assistant text (the deliverable). */
text: string;
}
export function subagentEngine(role: ModelRole, override?: PiEngine): PiEngine {
return override ?? piEngineFor(role);
}
export async function runSubagent(tools: AgentTool<any, any>[], opts: SubagentRunOptions): Promise<SubagentRunResult> {
const engine = subagentEngine(opts.role, opts.engine);
let turns = 0;
const agent = new Agent({
initialState: {
systemPrompt: opts.systemPrompt,
model: engine.model,
thinkingLevel: thinkingLevelFor(opts.role),
tools,
messages: opts.messages ?? [],
},
streamFn: engine.streamFn,
getApiKey: () => openRouterApiKey(),
shouldStopAfterTurn: (context) => {
turns++;
return turns >= opts.maxTurns || opts.shouldStop?.(context.message, context.toolResults) === true;
},
});
const onAbort = () => agent.abort();
opts.abortSignal?.addEventListener("abort", onAbort, { once: true });
const before = opts.messages?.length ?? 0;
let step = 0;
let errorMessage: string | undefined;
let aborted = false;
try {
agent.subscribe((event: AgentEvent) => {
switch (event.type) {
case "message_end":
if ((event.message as { role?: string }).role === "assistant") {
const m = event.message as AssistantMessage;
if (m.stopReason === "aborted") {
// The call was cancelled — no completed step, no billed
// usage (the abort path asserts zero steps).
aborted = true;
break;
}
step++;
opts.onUsage?.(m.usage);
opts.onAssistantMessage?.(m, step);
if (m.stopReason === "error") errorMessage = m.errorMessage;
}
break;
case "tool_execution_start":
opts.onToolStart?.(event.toolName, event.args);
break;
}
});
if (opts.prompt !== undefined) await agent.prompt({
role: "user",
content: opts.prompt,
timestamp: Date.now(),
});
else await agent.continue();
} finally {
opts.abortSignal?.removeEventListener("abort", onAbort);
}
if (aborted || opts.abortSignal?.aborted) {
const err = new Error("aborted");
err.name = "AbortError";
throw err;
}
if (errorMessage !== undefined) throw new Error(errorMessage);
const newMessages = agent.state.messages.slice(before);
const assistants = newMessages.filter((m) => (m as { role?: string }).role === "assistant") as AssistantMessage[];
const last = assistants[assistants.length - 1];
const text = last
? last.content
.filter((c): c is { type: "text"; text: string } => c.type === "text")
.map((c) => c.text)
.join("")
: "";
return { newMessages, text };
}
+30 -11
View File
@@ -1,6 +1,9 @@
import { ToolLoopAgent, stepCountIs, type LanguageModel } from "ai"; import type { AgentTool } from "@earendil-works/pi-agent-core";
import type { Usage } from "@earendil-works/pi-ai";
import { webTools } from "../tools/web.ts"; import { webTools } from "../tools/web.ts";
import { modelFor } from "../models.ts"; import { toPiTools } from "../lh-tool.ts";
import type { PiEngine } from "../models.ts";
import { runSubagent } from "./pi-run.ts";
/** /**
* The researcher sub-agent (vision §2.5): a SEPARATE ToolLoopAgent instance * The researcher sub-agent (vision §2.5): a SEPARATE ToolLoopAgent instance
@@ -9,7 +12,7 @@ import { modelFor } from "../models.ts";
* *
* System prompt adapted from the pi researcher agent * System prompt adapted from the pi researcher agent
* (~/.pi/agent/extensions/interactive-subagents/agents/researcher.md). * (~/.pi/agent/extensions/interactive-subagents/agents/researcher.md).
* The model is injectable so tests run the loop with a MockLanguageModelV4. * The engine is injectable so tests run the loop with a scripted faux engine.
*/ */
const INSTRUCTIONS = `You are a research specialist. You work for a teacher agent inside a learning const INSTRUCTIONS = `You are a research specialist. You work for a teacher agent inside a learning
@@ -57,13 +60,29 @@ Numbered findings with inline source citations:
## Gaps ## Gaps
What couldn't be answered. Suggested next steps.`; What couldn't be answered. Suggested next steps.`;
export function createResearcherAgent(model?: LanguageModel) { export interface ResearcherRunOptions {
return new ToolLoopAgent({ engine?: PiEngine;
model: model ?? modelFor("researcher"), abortSignal?: AbortSignal;
instructions: INSTRUCTIONS, onUsage?: (usage: Usage) => void;
tools: webTools, onToolStart?: (toolName: string, args: unknown) => void;
stopWhen: stepCountIs(12),
});
} }
export type ResearcherAgent = ReturnType<typeof createResearcherAgent>; /**
* Run the researcher to completion: isolated context, web tools only, final
* text = the brief. The engine is injectable so tests run the loop scripted.
*/
export async function runResearcherAgent(prompt: string, opts: ResearcherRunOptions = {}): Promise<string> {
const { text } = await runSubagent(toPiTools(webTools) as AgentTool<any, any>[], {
role: "researcher",
prompt,
engine: opts.engine,
systemPrompt: INSTRUCTIONS,
maxTurns: 12,
abortSignal: opts.abortSignal,
onUsage: opts.onUsage,
onToolStart: opts.onToolStart,
});
return text;
}
export type ResearcherAgent = ReturnType<typeof runResearcherAgent>;
+24 -25
View File
@@ -1,7 +1,9 @@
import { isToolUIPart, type LanguageModel, type UIMessage } from "ai"; import { isToolUIPart, type UIMessage } from "ai";
import type { PiEngine } from "../models.ts";
import { import {
accumulateJobUsage, accumulateJobUsage,
buildSubagentSteer, buildSubagentSteer,
piUsageToStep,
type DemoResult, type DemoResult,
type DemoToolInput, type DemoToolInput,
type JobKind, type JobKind,
@@ -18,7 +20,7 @@ import { registerJobAbort, unregisterJobAbort } from "./aborts.ts";
import { MODELS } from "../../../config.ts"; import { MODELS } from "../../../config.ts";
import type { ChatStore, JobStore } from "../../core/ports.ts"; import type { ChatStore, JobStore } from "../../core/ports.ts";
import { runDemoBuilder } from "./demo-builder.ts"; import { runDemoBuilder } from "./demo-builder.ts";
import { createResearcherAgent } from "./researcher.ts"; import { runResearcherAgent } from "./researcher.ts";
import { runVisualMaker } from "./visual-maker.ts"; import { runVisualMaker } from "./visual-maker.ts";
/** /**
@@ -42,13 +44,13 @@ import { runVisualMaker } from "./visual-maker.ts";
/** /**
* Test seam: the spawn wiring has no injection point at the tool layer, so * 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. * tests swap in a scripted pi engine here to keep job tests off the network.
* Undefined in production. (The jobs spawn through the module-level KIND * Undefined in production. (The jobs spawn through the module-level KIND
* runners, which read these at call time.) * runners, which read these at call time.)
*/ */
export const researchModelOverride: { current?: LanguageModel } = {}; export const researchModelOverride: { current?: PiEngine } = {};
export const demoModelOverride: { current?: LanguageModel } = {}; export const demoModelOverride: { current?: PiEngine } = {};
export const visualModelOverride: { current?: LanguageModel } = {}; export const visualModelOverride: { current?: PiEngine } = {};
/** The dispatch-time tool input (from src/agent/tools/jobs.ts schemas). */ /** The dispatch-time tool input (from src/agent/tools/jobs.ts schemas). */
type JobToolInput = ResearchToolInput | DemoToolInput | VisualToolInput; type JobToolInput = ResearchToolInput | DemoToolInput | VisualToolInput;
@@ -232,7 +234,7 @@ export function createJobRunner(deps: JobRunnerDeps): JobRunner {
async function runResearchJob( async function runResearchJob(
jobId: string, jobId: string,
model: LanguageModel | undefined, model: PiEngine | undefined,
passedInput: JobToolInput | undefined, passedInput: JobToolInput | undefined,
loop: LoopContext, loop: LoopContext,
): Promise<void> { ): Promise<void> {
@@ -253,28 +255,25 @@ export function createJobRunner(deps: JobRunnerDeps): JobRunner {
.join("\n\n"); .join("\n\n");
const fetchedUrls: string[] = []; const fetchedUrls: string[] = [];
const agent = createResearcherAgent(model);
try { try {
jobStore.updateJob(jobId, { activity: "researching" }); jobStore.updateJob(jobId, { activity: "researching" });
const result = await agent.generate({ const brief = await runResearcherAgent(prompt, {
prompt, engine: model,
abortSignal: loop.abortSignal, abortSignal: loop.abortSignal,
// Money channel: each agent step is one model call — fold its // Money channel: each agent step is one model call — fold its
// usage into the record as it completes (progress write). // usage into the record as it completes (progress write).
onStepFinish: ({ usage }) => loop.onUsage(usage), onUsage: (usage) => loop.onUsage(piUsageToStep(usage)),
onToolExecutionStart: ({ toolCall }) => { onToolStart: (toolName, args) => {
const input = toolCall.input as { query?: string; url?: string } | undefined; const input = args as { query?: string; url?: string } | undefined;
if (toolCall.toolName === "web_fetch" && typeof input?.url === "string") { if (toolName === "web_fetch" && typeof input?.url === "string") {
fetchedUrls.push(input.url); fetchedUrls.push(input.url);
} }
jobStore.updateJob(jobId, { activity: activityLine(toolCall.toolName, input) }); jobStore.updateJob(jobId, { activity: activityLine(toolName, input) });
}, },
}); });
if (!brief.trim()) throw new Error("the researcher produced no text");
const brief = result.text.trim(); const research: ResearchResult = { brief: brief.trim(), sources: parseSources(brief.trim(), fetchedUrls) };
if (!brief) throw new Error("the researcher produced no text");
const research: ResearchResult = { brief, sources: parseSources(brief, fetchedUrls) };
await finishJob(jobId, research); await finishJob(jobId, research);
} catch (err) { } catch (err) {
// Real-kill abort (2026-09): the abort path flipped the record and // Real-kill abort (2026-09): the abort path flipped the record and
@@ -291,7 +290,7 @@ export function createJobRunner(deps: JobRunnerDeps): JobRunner {
async function runDemoJob( async function runDemoJob(
jobId: string, jobId: string,
model: LanguageModel | undefined, model: PiEngine | undefined,
passedInput: JobToolInput | undefined, passedInput: JobToolInput | undefined,
loop: LoopContext, loop: LoopContext,
): Promise<void> { ): Promise<void> {
@@ -308,9 +307,9 @@ export function createJobRunner(deps: JobRunnerDeps): JobRunner {
const result = await runDemoBuilder( const result = await runDemoBuilder(
{ title, brief }, { title, brief },
{ {
model, engine: model,
abortSignal: loop.abortSignal, abortSignal: loop.abortSignal,
onUsage: loop.onUsage, onUsage: (usage) => loop.onUsage(piUsageToStep(usage)),
onActivity: (line) => jobStore.updateJob(jobId, { activity: line }), onActivity: (line) => jobStore.updateJob(jobId, { activity: line }),
}, },
); );
@@ -327,7 +326,7 @@ export function createJobRunner(deps: JobRunnerDeps): JobRunner {
async function runVisualJob( async function runVisualJob(
jobId: string, jobId: string,
model: LanguageModel | undefined, model: PiEngine | undefined,
passedInput: JobToolInput | undefined, passedInput: JobToolInput | undefined,
loop: LoopContext, loop: LoopContext,
): Promise<void> { ): Promise<void> {
@@ -343,9 +342,9 @@ export function createJobRunner(deps: JobRunnerDeps): JobRunner {
const result = await runVisualMaker( const result = await runVisualMaker(
{ brief }, { brief },
{ {
model, engine: model,
abortSignal: loop.abortSignal, abortSignal: loop.abortSignal,
onUsage: loop.onUsage, onUsage: (usage) => loop.onUsage(piUsageToStep(usage)),
onActivity: (line) => jobStore.updateJob(jobId, { activity: line }), onActivity: (line) => jobStore.updateJob(jobId, { activity: line }),
}, },
); );
+100 -47
View File
@@ -1,7 +1,8 @@
import { generateText, tool, type LanguageModel, type ModelMessage, type ToolResultPart } from "ai"; import type { AssistantMessage, Message, ToolResultMessage, Usage } from "@earendil-works/pi-ai";
import { z } from "zod"; import { z } from "zod";
import { modelFor } from "../models.ts"; import { lhTool, piToolDeclarations, type LhTool } from "../lh-tool.ts";
import type { UsageStepLike, VisualResult } from "../../core/contracts.ts"; import { piEngineFor, thinkingLevelFor, type PiEngine } from "../models.ts";
import type { VisualResult } from "../../core/contracts.ts";
import { getHeadless, type HeadlessApi } from "./headless.ts"; import { getHeadless, type HeadlessApi } from "./headless.ts";
/** /**
@@ -31,7 +32,8 @@ export interface VisualMakerInput {
} }
export interface VisualMakerOptions { export interface VisualMakerOptions {
model?: LanguageModel; /** Test seam: scripted engine; production uses the visualMaker role. */
engine?: PiEngine;
headless?: HeadlessApi; headless?: HeadlessApi;
/** Max model rounds (write/render cycles + final answer). Default 8. */ /** Max model rounds (write/render cycles + final answer). Default 8. */
maxIterations?: number; maxIterations?: number;
@@ -39,9 +41,9 @@ export interface VisualMakerOptions {
onActivity?: (line: string) => void; onActivity?: (line: string) => void;
/** Kills the run mid-model-call (2026-09 real-kill abort semantics). */ /** Kills the run mid-model-call (2026-09 real-kill abort semantics). */
abortSignal?: AbortSignal; abortSignal?: AbortSignal;
/** Money channel: one call per loop iteration (generateText), folded /** Money channel: one call per loop iteration (streamSimple), folded
* into JobRecord.usage. */ * into JobRecord.usage. */
onUsage?: (usage: UsageStepLike) => void; onUsage?: (usage: Usage) => void;
} }
const INSTRUCTIONS = `You are a visual maker — a diagram author + renderer for precise, spatial figures. You const INSTRUCTIONS = `You are a visual maker — a diagram author + renderer for precise, spatial figures. You
@@ -86,25 +88,28 @@ Your FINAL message: 1-2 sentences of notes for the teacher (what the figure show
NEVER the SVG source itself. If you genuinely cannot make a correct picture of the brief, say exactly NEVER the SVG source itself. If you genuinely cannot make a correct picture of the brief, say exactly
that in your final message and why.`; that in your final message and why.`;
export function createVisualTools() { export function createVisualTools(): Record<string, LhTool<any, any>> {
return { return {
write_svg: tool({ write_svg: lhTool({
name: "write_svg",
description: "Write the complete SVG source for the figure (overwrites any previous version).", description: "Write the complete SVG source for the figure (overwrites any previous version).",
inputSchema: z.object({ schema: z.object({
svg: z.string().describe('Complete <svg xmlns="…">…</svg> markup.'), svg: z.string().describe('Complete <svg xmlns="…">…</svg> markup.'),
}), }),
}), }),
edit_svg: tool({ edit_svg: lhTool({
name: "edit_svg",
description: description:
"Replace one exact spot in the current SVG: old_text must appear EXACTLY ONCE and is replaced by new_text. Call write_svg first. For larger changes, write_svg the whole document.", "Replace one exact spot in the current SVG: old_text must appear EXACTLY ONCE and is replaced by new_text. Call write_svg first. For larger changes, write_svg the whole document.",
inputSchema: z.object({ schema: z.object({
old_text: z.string().describe("The exact text to replace — must match exactly once in the current SVG."), old_text: z.string().describe("The exact text to replace — must match exactly once in the current SVG."),
new_text: z.string().describe("The replacement text."), new_text: z.string().describe("The replacement text."),
}), }),
}), }),
render_svg: tool({ render_svg: lhTool({
name: "render_svg",
description: "Render the current SVG to a PNG and show it to you for inspection. No arguments.", description: "Render the current SVG to a PNG and show it to you for inspection. No arguments.",
inputSchema: z.object({}), schema: z.object({}),
}), }),
}; };
} }
@@ -125,6 +130,9 @@ export async function runVisualMaker(input: VisualMakerInput, opts: VisualMakerO
} }
const tools = createVisualTools(); const tools = createVisualTools();
const toolDeclarations = piToolDeclarations(tools);
const engine = opts.engine ?? piEngineFor("visualMaker");
const reasoning = thinkingLevelFor("visualMaker");
const prompt = [ const prompt = [
"Brief for ONE precise figure (from the teacher, verbatim — it is everything you know):", "Brief for ONE precise figure (from the teacher, verbatim — it is everything you know):",
brief, brief,
@@ -132,7 +140,7 @@ export async function runVisualMaker(input: VisualMakerInput, opts: VisualMakerO
"Author the SVG, render it, look at it, iterate until correct, then finish with your notes.", "Author the SVG, render it, look at it, iterate until correct, then finish with your notes.",
].join("\n"); ].join("\n");
let messages: ModelMessage[] = [{ role: "user", content: prompt }]; let messages: Message[] = [{ role: "user", content: prompt, timestamp: Date.now() }];
let currentSvg: string | null = null; let currentSvg: string | null = null;
/** The svg shown in the most recent render (null = never rendered). */ /** The svg shown in the most recent render (null = never rendered). */
let renderedSvg: string | null = null; let renderedSvg: string | null = null;
@@ -140,20 +148,20 @@ export async function runVisualMaker(input: VisualMakerInput, opts: VisualMakerO
try { try {
for (let iteration = 0; iteration < maxIterations; iteration++) { for (let iteration = 0; iteration < maxIterations; iteration++) {
const result = await generateText({ // One model call per iteration (streamSimple, no agent loop — this
model: opts.model ?? modelFor("visualMaker"), // loop owns tool execution and the image-in-the-loop injection).
system: INSTRUCTIONS, const result = await streamOnce(engine, INSTRUCTIONS, messages, toolDeclarations, reasoning, opts.abortSignal);
messages,
tools,
abortSignal: opts.abortSignal,
});
// Money channel: each iteration is one model call. // Money channel: each iteration is one model call.
opts.onUsage?.(result.usage); opts.onUsage?.(result.usage);
messages = [...messages, ...result.response.messages]; messages = [...messages, result];
const toolCalls = result.toolCalls; const toolCalls = result.content.filter((c): c is Extract<typeof c, { type: "toolCall" }> => c.type === "toolCall");
if (toolCalls.length === 0) { if (toolCalls.length === 0) {
finalText = result.text.trim(); finalText = result.content
.filter((c): c is { type: "text"; text: string } => c.type === "text")
.map((c) => c.text)
.join("")
.trim();
// Doctrine gate: never deliver an unlooked SVG. If the current SVG // Doctrine gate: never deliver an unlooked SVG. If the current SVG
// changed since its last render (or was never rendered), show it now // changed since its last render (or was never rendered), show it now
// and make the model look before accepting the finish. // and make the model look before accepting the finish.
@@ -169,41 +177,48 @@ export async function runVisualMaker(input: VisualMakerInput, opts: VisualMakerO
renderedSvg = currentSvg; renderedSvg = currentSvg;
messages = [ messages = [
...messages, ...messages,
{ role: "user", content: [{ type: "file", data: png, mediaType: "image/png" }, { type: "text", text: UNLOOKED_FINISH_MESSAGE }] }, {
role: "user",
content: [
{ type: "text", text: UNLOOKED_FINISH_MESSAGE },
{ type: "image", data: png.toString("base64"), mimeType: "image/png" },
],
timestamp: Date.now(),
},
]; ];
continue; continue;
} }
return finish(currentSvg, finalText); return finish(currentSvg, finalText);
} }
const toolResults: ToolResultPart[] = []; const toolResults: ToolResultMessage[] = [];
let renderPng: Buffer | null = null; // appended as the image message AFTER the tool results let renderPng: Buffer | null = null; // appended as the image message AFTER the tool results
for (const call of toolCalls) { for (const call of toolCalls) {
if (call.toolName === "write_svg") { if (call.name === "write_svg") {
const svg = (call.input as { svg?: unknown }).svg; const svg = (call.arguments as { svg?: unknown }).svg;
if (typeof svg !== "string" || !svg.includes("<svg")) { if (typeof svg !== "string" || !svg.includes("<svg")) {
toolResults.push(toolText(call.toolCallId, call.toolName, "ERROR: expected a complete <svg>…</svg> document in the `svg` field.")); toolResults.push(toolText(call.id, call.name, "ERROR: expected a complete <svg>…</svg> document in the `svg` field."));
continue; continue;
} }
currentSvg = svg; currentSvg = svg;
opts.onActivity?.(`write_svg: ${(svg.length / 1024).toFixed(1)} KB`); opts.onActivity?.(`write_svg: ${(svg.length / 1024).toFixed(1)} KB`);
toolResults.push(toolText(call.toolCallId, call.toolName, `Saved (${svg.length} bytes). Call render_svg to look at it.`)); toolResults.push(toolText(call.id, call.name, `Saved (${svg.length} bytes). Call render_svg to look at it.`));
} else if (call.toolName === "edit_svg") { } else if (call.name === "edit_svg") {
const edit = call.input as { old_text?: unknown; new_text?: unknown }; const edit = call.arguments as { old_text?: unknown; new_text?: unknown };
if (!currentSvg) { if (!currentSvg) {
toolResults.push(toolText(call.toolCallId, call.toolName, "ERROR: nothing to edit — call write_svg first.")); toolResults.push(toolText(call.id, call.name, "ERROR: nothing to edit — call write_svg first."));
continue; continue;
} }
if (typeof edit.old_text !== "string" || typeof edit.new_text !== "string" || !edit.old_text) { if (typeof edit.old_text !== "string" || typeof edit.new_text !== "string" || !edit.old_text) {
toolResults.push(toolText(call.toolCallId, call.toolName, "ERROR: expected old_text (non-empty) and new_text strings.")); toolResults.push(toolText(call.id, call.name, "ERROR: expected old_text (non-empty) and new_text strings."));
continue; continue;
} }
const occurrences = currentSvg.split(edit.old_text).length - 1; const occurrences = currentSvg.split(edit.old_text).length - 1;
if (occurrences === 0) { if (occurrences === 0) {
toolResults.push( toolResults.push(
toolText( toolText(
call.toolCallId, call.id,
call.toolName, call.name,
`ERROR: old_text not found. Closest line for reference:\n${snippetAround(currentSvg, edit.old_text)}`, `ERROR: old_text not found. Closest line for reference:\n${snippetAround(currentSvg, edit.old_text)}`,
), ),
); );
@@ -211,16 +226,16 @@ export async function runVisualMaker(input: VisualMakerInput, opts: VisualMakerO
} }
if (occurrences > 1) { if (occurrences > 1) {
toolResults.push( toolResults.push(
toolText(call.toolCallId, call.toolName, `ERROR: old_text appears ${occurrences} times — it must match EXACTLY ONCE. Include more surrounding context to make it unique.`), toolText(call.id, call.name, `ERROR: old_text appears ${occurrences} times — it must match EXACTLY ONCE. Include more surrounding context to make it unique.`),
); );
continue; continue;
} }
currentSvg = currentSvg.replace(edit.old_text, edit.new_text); currentSvg = currentSvg.replace(edit.old_text, edit.new_text);
opts.onActivity?.("edit_svg: applied"); opts.onActivity?.("edit_svg: applied");
toolResults.push(toolText(call.toolCallId, call.toolName, "OK — replaced 1 occurrence. Call render_svg to look at the edited SVG.")); toolResults.push(toolText(call.id, call.name, "OK — replaced 1 occurrence. Call render_svg to look at the edited SVG."));
} else if (call.toolName === "render_svg") { } else if (call.name === "render_svg") {
if (!currentSvg) { if (!currentSvg) {
toolResults.push(toolText(call.toolCallId, call.toolName, "ERROR: nothing to render — call write_svg first.")); toolResults.push(toolText(call.id, call.name, "ERROR: nothing to render — call write_svg first."));
continue; continue;
} }
try { try {
@@ -228,18 +243,18 @@ export async function runVisualMaker(input: VisualMakerInput, opts: VisualMakerO
renderPng = png; renderPng = png;
renderedSvg = currentSvg; renderedSvg = currentSvg;
opts.onActivity?.("render_svg: inspecting render"); opts.onActivity?.("render_svg: inspecting render");
toolResults.push(toolText(call.toolCallId, call.toolName, "Rendered — the PNG follows in the next message. Look at it.")); toolResults.push(toolText(call.id, call.name, "Rendered — the PNG follows in the next message. Look at it."));
} catch (err) { } catch (err) {
const msg = err instanceof Error ? err.message : String(err); const msg = err instanceof Error ? err.message : String(err);
opts.onActivity?.("render_svg: render error"); opts.onActivity?.("render_svg: render error");
toolResults.push(toolText(call.toolCallId, call.toolName, `RENDER ERROR (the SVG may be malformed): ${msg}. Fix it with write_svg and render again.`)); toolResults.push(toolText(call.id, call.name, `RENDER ERROR (the SVG may be malformed): ${msg}. Fix it with write_svg and render again.`));
} }
} else { } else {
toolResults.push(toolText(call.toolCallId, call.toolName, `ERROR: unknown tool "${call.toolName}". Use write_svg, edit_svg, or render_svg.`)); toolResults.push(toolText(call.id, call.name, `ERROR: unknown tool "${call.name}". Use write_svg, edit_svg, or render_svg.`));
} }
} }
// Message order: assistant tool-calls → tool results → render image. // Message order: assistant tool-calls → tool results → render image.
messages = [...messages, { role: "tool", content: toolResults }]; messages = [...messages, ...toolResults];
if (renderPng) { if (renderPng) {
messages = [ messages = [
...messages, ...messages,
@@ -247,8 +262,9 @@ export async function runVisualMaker(input: VisualMakerInput, opts: VisualMakerO
role: "user", role: "user",
content: [ content: [
{ type: "text", text: LOOK_MESSAGE }, { type: "text", text: LOOK_MESSAGE },
{ type: "file", data: renderPng, mediaType: "image/png" }, { type: "image", data: renderPng.toString("base64"), mimeType: "image/png" },
], ],
timestamp: Date.now(),
}, },
]; ];
} }
@@ -265,8 +281,45 @@ export async function runVisualMaker(input: VisualMakerInput, opts: VisualMakerO
return { svg: "", notes: `FAILED: iteration limit (${maxIterations}) reached without a verified SVG.` }; return { svg: "", notes: `FAILED: iteration limit (${maxIterations}) reached without a verified SVG.` };
} }
function toolText(toolCallId: string, toolName: string, value: string): ToolResultPart { function toolText(toolCallId: string, toolName: string, value: string): ToolResultMessage {
return { type: "tool-result", toolCallId, toolName, output: { type: "text", value } }; return {
role: "toolResult",
toolCallId,
toolName,
content: [{ type: "text", text: value }],
isError: value.startsWith("ERROR:") || value.startsWith("RENDER ERROR"),
timestamp: Date.now(),
};
}
/** One raw model call (no agent loop): drain the event stream to the final
* AssistantMessage; provider errors and aborts throw so the loop's catch
* path keeps its old generateText semantics. */
async function streamOnce(
engine: PiEngine,
systemPrompt: string,
messages: Message[],
tools: ReturnType<typeof piToolDeclarations>,
reasoning: ReturnType<typeof thinkingLevelFor>,
abortSignal?: AbortSignal,
): Promise<AssistantMessage> {
const stream = engine.streamFn(engine.model, { systemPrompt, messages, tools }, {
reasoning: reasoning === "off" ? undefined : reasoning,
apiKey: process.env.OPENROUTER_API_KEY || undefined,
signal: abortSignal,
} as never);
for await (const event of stream) {
if (event.type === "done") return event.message;
if (event.type === "error") {
if (event.reason === "aborted") {
const err = new Error("aborted");
err.name = "AbortError";
throw err;
}
throw new Error(event.error.errorMessage ?? "the model call failed");
}
}
throw new Error("the model stream ended without a result");
} }
/** A 3-line numbered context window around the closest fuzzy match for a /** A 3-line numbered context window around the closest fuzzy match for a
+105 -46
View File
@@ -1,8 +1,11 @@
import * as fs from "node:fs"; import * as fs from "node:fs";
import * as path from "node:path"; import * as path from "node:path";
import { ToolLoopAgent, stepCountIs, tool, type LanguageModel, type LanguageModelUsage, type ToolSet } from "ai"; import { Agent } from "@earendil-works/pi-agent-core";
import type { AgentEvent, AgentMessage, AgentTool } from "@earendil-works/pi-agent-core";
import type { Usage } from "@earendil-works/pi-ai";
import { PATHS } from "../../config.ts"; import { PATHS } from "../../config.ts";
import { modelFor } from "./models.ts"; import { MODELS } from "../../config.ts";
import { piEngineFor, thinkingLevelFor, openRouterApiKey, type PiEngine } from "./models.ts";
import { import {
askUserQuestionTool, askUserQuestionTool,
createProberTools, createProberTools,
@@ -10,17 +13,20 @@ import {
type ProbeDoneContext, type ProbeDoneContext,
type TeacherToolContext, type TeacherToolContext,
} from "./tools.ts"; } from "./tools.ts";
import { toPiTools, type LhTool } from "./lh-tool.ts";
import { quizInputObject, quizInputProblems } from "../core/question-core.ts"; import { quizInputObject, quizInputProblems } from "../core/question-core.ts";
import { QUESTION_TOOLS } from "../core/pending.ts";
/** /**
* The two teaching-side agents (docs/probe-teach-split.md): the PROBER and * The two teaching-side agents (docs/probe-teach-split.md): the PROBER and
* the TEACHER. Same harness (agent loop, tool calling, streaming, * the TEACHER. Same transcript, different prompts and tool sets; the
* multi-step — the few lines ToolLoopAgent provides), same transcript, but * chat-handler constructs ONE per request off the transcript's seam
* different prompts and different tool sets; the chat-handler constructs * (probe_done = teacher) — the STATELESS per-turn engine (docs/
* one per request off the transcript's seam (probe_done = teacher). * pi-engine-swap.md §2): the persisted record is the source of truth, the
* agent is rebuilt from it every POST and discarded at turn end.
* *
* Models are injectable so tests can pass a MockLanguageModelV4 and drive * Engines are injectable so tests can pass a scripted faux engine and drive
* deterministic scripted lessons (see tests/chat-flow.test.ts). * deterministic lessons (see tests/chat-flow.test.ts).
* *
* System prompts are COMPOSED per agent: system-prompt.md (the shared base: * System prompts are COMPOSED per agent: system-prompt.md (the shared base:
* philosophy, accuracy, quiz construction, question tools, formatting) + * philosophy, accuracy, quiz construction, question tools, formatting) +
@@ -62,9 +68,81 @@ const composedInstructions: Record<AgentKind, string> = {
// Money channel (wave 4): the chat-handler passes a per-turn collector; // 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 // every model call's usage lands in the chat record's turns log
// (ChatStore.appendTurnUsage). The step carries far more than usage — // (ChatStore.appendTurnUsage). The narrow param keeps the seam honest.
// the narrow param keeps the seam honest about what we consume. export type StepUsageCollector = (usage: Usage) => void;
export type StepUsageCollector = (step: { usage: LanguageModelUsage }) => void;
/** Per-request construction options (the handler wires all of these). */
export interface AgentBuildOptions {
/** Model + stream function — tests inject a scripted faux engine. */
engine?: PiEngine;
/** The initial transcript, as pi messages (recordToPiMessages). */
messages?: AgentMessage[];
/** Money channel: one call per model request. */
onUsage?: StepUsageCollector;
/** Safety net (the old stepCountIs). Default 20 turns. */
maxTurns?: number;
/** Tests may pass a bare tool registry to drive deterministic flows. */
tools?: Record<string, LhTool<any, any>>;
}
/**
* The turn's stop rule (docs/pi-engine-swap.md §2):
* - a question tool call (quiz / ask_user_question) that RESOLVED — its
* synthetic "question posed" result is not an error — ends the run: the
* learner's answer arrives as the next POST. A call that FAILED authoring
* validation (error result) must NOT stop the run — the model retries
* within the same turn, exactly the InvalidToolInputError semantics.
* - the maxTurns safety net replaces stepCountIs(20).
*/
function makeStopper(maxTurns: number) {
let turns = 0;
return (context: {
message: { content: Array<{ type: string; id?: string; name?: string }> };
toolResults: Array<{ toolCallId: string; isError?: boolean }>;
}): boolean => {
turns++;
if (turns >= maxTurns) return true;
const called = new Set(
context.message.content
.filter((c) => c.type === "toolCall" && c.name !== undefined && (QUESTION_TOOLS as readonly string[]).includes(c.name!))
.map((c) => c.id!),
);
if (called.size === 0) return false;
// Stop only when every question call of this turn resolved cleanly.
for (const result of context.toolResults) {
if (called.has(result.toolCallId) && result.isError) return false;
}
return true;
};
}
function baseAgentOptions(kind: AgentKind, opts: AgentBuildOptions, tools: AgentTool<any, any>[]) {
const engine = opts.engine ?? piEngineFor(kind);
const stopper = makeStopper(opts.maxTurns ?? 20);
const agent = new Agent({
initialState: {
systemPrompt: composedInstructions[kind],
model: engine.model,
thinkingLevel: thinkingLevelFor(kind),
tools,
messages: opts.messages ?? [],
},
streamFn: engine.streamFn,
getApiKey: () => openRouterApiKey(),
shouldStopAfterTurn: stopper,
});
if (opts.onUsage) {
agent.subscribe((event: AgentEvent) => {
if (event.type === "message_end" && event.message.role === "assistant" && "usage" in event.message) {
// A cancelled call is not a billed step (partial/estimated
// usage on an aborted message is not money spent).
if (event.message.stopReason === "aborted") return;
opts.onUsage!(event.message.usage);
}
});
}
return agent;
}
/** /**
* The PROBER (agent A): quiz (no nodeId) + ask + probe_done, lean on * The PROBER (agent A): quiz (no nodeId) + ask + probe_done, lean on
@@ -73,19 +151,8 @@ export type StepUsageCollector = (step: { usage: LanguageModelUsage }) => void;
* handoff wake (see tools/probe.ts) and is the transcript seam the * handoff wake (see tools/probe.ts) and is the transcript seam the
* handler keys teacher construction on. * handler keys teacher construction on.
*/ */
export function createProberAgent(ctx: ProbeDoneContext, model?: LanguageModel, onStepFinish?: StepUsageCollector) { export function createProberAgent(ctx: ProbeDoneContext, opts: AgentBuildOptions = {}) {
return new ToolLoopAgent({ return baseAgentOptions("prober", opts, toPiTools(createProberTools(ctx)));
// Model choice + thinking budget live in config.ts (MODELS.prober).
// Reasoning tokens stream as reasoning parts into the UI's collapsible block.
model: model ?? modelFor("prober"),
instructions: composedInstructions.prober,
tools: createProberTools(ctx),
// Question tools are client-side (a call ends the server turn);
// probe_done resolves immediately (it queues the handoff and acks).
// stepCountIs is just the safety net.
stopWhen: stepCountIs(20),
...(onStepFinish ? { onStepFinish } : {}),
});
} }
/** /**
@@ -99,40 +166,27 @@ export function createTeacherAgent(
// wiring): without ctx the tools fall back to the validation-free // wiring): without ctx the tools fall back to the validation-free
// question set. The handler ALWAYS passes ctx (the composition root). // question set. The handler ALWAYS passes ctx (the composition root).
ctx?: TeacherToolContext, ctx?: TeacherToolContext,
model?: LanguageModel, opts: AgentBuildOptions = {},
onStepFinish?: StepUsageCollector,
// Tests may pass a bare tools set to drive deterministic flows; when
// present it overrides the ctx-composed registry.
tools?: ToolSet,
) { ) {
return new ToolLoopAgent({ const tools = opts.tools ?? (ctx ? createTeacherTools(ctx) : bareQuestionTools());
// Model choice + thinking budget live in config.ts (MODELS.teacher). return baseAgentOptions("teacher", opts, toPiTools(tools));
model: model ?? modelFor("teacher"),
instructions: composedInstructions.teacher,
// Tools are per-request (they close over the injected ctx — chat-
// handler passes the composition root's stores + spawn seam).
tools: tools ?? (ctx ? createTeacherTools(ctx) : bareQuestionTools()),
// Question tools are client-side (a call ends the server turn);
// job tools resolve immediately (fire-and-forget ack). stepCountIs is
// just the safety net.
stopWhen: stepCountIs(20),
...(onStepFinish ? { onStepFinish } : {}),
});
} }
/** The validation-free teach question set (tests, bare agents): quiz with /** The validation-free teach question set (tests, bare agents): quiz with
* the shared authoring refinement but NO per-session nodeId check (no * the shared authoring refinement but NO per-session nodeId check (no
* record to read), plus ask. */ * record to read), plus ask. */
export function bareQuestionTools(): ToolSet { export function bareQuestionTools(): Record<string, LhTool<any, any>> {
return { return {
quiz: tool({ quiz: {
name: "quiz",
description: "Ask the learner a graded question.", description: "Ask the learner a graded question.",
inputSchema: quizInputObject.superRefine((input, add) => { clientQuestion: true,
schema: quizInputObject.superRefine((input, add) => {
for (const problem of quizInputProblems(input)) { for (const problem of quizInputProblems(input)) {
add.addIssue({ code: "custom", message: problem }); add.addIssue({ code: "custom", message: problem });
} }
}), }),
}), },
ask_user_question: askUserQuestionTool, ask_user_question: askUserQuestionTool,
}; };
} }
@@ -140,3 +194,8 @@ export function bareQuestionTools(): ToolSet {
export type ProberAgent = ReturnType<typeof createProberAgent>; export type ProberAgent = ReturnType<typeof createProberAgent>;
export type TeacherAgent = ReturnType<typeof createTeacherAgent>; export type TeacherAgent = ReturnType<typeof createTeacherAgent>;
export type TeachingAgent = ProberAgent | TeacherAgent; export type TeachingAgent = ProberAgent | TeacherAgent;
/** The model id a role's engine resolves to (observability). */
export function modelIdFor(kind: AgentKind): string {
return MODELS[kind].model;
}
+17 -11
View File
@@ -1,4 +1,4 @@
import { tool, type ToolSet } from "ai"; import { lhTool, type LhTool } from "./lh-tool.ts";
import { import {
askInputSchema, askInputSchema,
probeQuizInputSchema, probeQuizInputSchema,
@@ -33,17 +33,21 @@ export type { ProbeDoneContext } from "./tools/probe.ts";
* only. * only.
*/ */
export const askUserQuestionTool = tool({ export const askUserQuestionTool: LhTool<any, any> = {
name: "ask_user_question",
clientQuestion: true,
description: description:
"Ask the learner a single question and pause until they answer. NOT graded — there is no correct answer. " + "Ask the learner a single question and pause until they answer. NOT graded — there is no correct answer. " +
"Use when requirements are ambiguous, preferences are needed, a decision would materially affect the lesson, " + "Use when requirements are ambiguous, preferences are needed, a decision would materially affect the lesson, " +
"or you need an open-ended response. Ask exactly one question per tool call; after calling, end your turn and wait.", "or you need an open-ended response. Ask exactly one question per tool call; after calling, end your turn and wait.",
inputSchema: askInputSchema, schema: askInputSchema,
}); };
// ─── the prober's question tools (no nodeId — no map exists yet) ───────────── // ─── the prober's question tools (no nodeId — no map exists yet) ─────────────
export const probeQuizTool = tool({ export const probeQuizTool: LhTool<any, any> = {
name: "quiz",
clientQuestion: true,
description: description:
"Ask the learner a GRADED question with a known correct answer, then instantly grade and give feedback. " + "Ask the learner a GRADED question with a known correct answer, then instantly grade and give feedback. " +
"Unlike ask_user_question (which collects preferences/decisions with no right answer), quiz always has a " + "Unlike ask_user_question (which collects preferences/decisions with no right answer), quiz always has a " +
@@ -55,11 +59,11 @@ export const probeQuizTool = tool({
"free-text note to ANY answer; it reaches you only when non-empty. No free-text answers — for non-graded " + "free-text note to ANY answer; it reaches you only when non-empty. No free-text answers — for non-graded " +
"questions use ask_user_question instead. Ask exactly one question per tool call; after calling, end your " + "questions use ask_user_question instead. Ask exactly one question per tool call; after calling, end your " +
"turn and wait for the graded result.", "turn and wait for the graded result.",
inputSchema: probeQuizInputSchema, schema: probeQuizInputSchema,
}); };
/** The prober's tool set: probe quiz (no nodeId) + ask + probe_done. */ /** The prober's tool set: probe quiz (no nodeId) + ask + probe_done. */
export function createProberTools(ctx: ProbeDoneContext): ToolSet { export function createProberTools(ctx: ProbeDoneContext): Record<string, LhTool<any, any>> {
return { return {
quiz: probeQuizTool, quiz: probeQuizTool,
ask_user_question: askUserQuestionTool, ask_user_question: askUserQuestionTool,
@@ -96,7 +100,9 @@ export function teachQuizToolFor(ctx: Pick<PlanToolContext, "chatId" | "chatStor
}); });
} }
}); });
return tool({ return lhTool({
name: "quiz",
clientQuestion: true,
description: description:
"Ask the learner a GRADED question with a known correct answer, then instantly grade and give feedback. " + "Ask the learner a GRADED question with a known correct answer, then instantly grade and give feedback. " +
"Unlike ask_user_question (which collects preferences/decisions with no right answer), quiz always has a " + "Unlike ask_user_question (which collects preferences/decisions with no right answer), quiz always has a " +
@@ -108,7 +114,7 @@ export function teachQuizToolFor(ctx: Pick<PlanToolContext, "chatId" | "chatStor
"An always-present optional note field lets the learner attach a free-text note to ANY answer; it reaches you " + "An always-present optional note field lets the learner attach a free-text note to ANY answer; it reaches you " +
"only when non-empty. No free-text answers — for non-graded questions use ask_user_question instead. " + "only when non-empty. No free-text answers — for non-graded questions use ask_user_question instead. " +
"Ask exactly one question per tool call; after calling, end your turn and wait for the graded result.", "Ask exactly one question per tool call; after calling, end your turn and wait for the graded result.",
inputSchema, schema: inputSchema,
}); });
} }
@@ -126,7 +132,7 @@ export interface TeacherToolContext extends PlanToolContext {
* Per-request teacher tool registry: teach questions (validated nodeId) + * Per-request teacher tool registry: teach questions (validated nodeId) +
* plan tools + job tools, all closing over the ctx. * plan tools + job tools, all closing over the ctx.
*/ */
export function createTeacherTools(ctx: TeacherToolContext): ToolSet { export function createTeacherTools(ctx: TeacherToolContext): Record<string, LhTool<any, any>> {
return { return {
quiz: teachQuizToolFor(ctx), quiz: teachQuizToolFor(ctx),
ask_user_question: askUserQuestionTool, ask_user_question: askUserQuestionTool,
+11 -8
View File
@@ -1,5 +1,5 @@
import { tool, type ToolSet } from "ai";
import { z } from "zod"; import { z } from "zod";
import { lhTool, type LhTool } from "../lh-tool.ts";
import type { JobRecord } from "../../core/contracts.ts"; import type { JobRecord } from "../../core/contracts.ts";
import type { JobStore } from "../../core/ports.ts"; import type { JobStore } from "../../core/ports.ts";
@@ -103,21 +103,24 @@ const VISUAL_DESCRIPTION =
* context (chatId + jobStore + the spawn seam), so concurrent chats never * context (chatId + jobStore + the spawn seam), so concurrent chats never
* share state. chat-handler wires it each turn; tests wire their own. * share state. chat-handler wires it each turn; tests wire their own.
*/ */
export function jobTools(ctx: JobToolContext): ToolSet { export function jobTools(ctx: JobToolContext): Record<string, LhTool<any, any>> {
return { return {
research: tool({ research: lhTool({
name: "research",
description: RESEARCH_DESCRIPTION, description: RESEARCH_DESCRIPTION,
inputSchema: researchInputSchema, schema: researchInputSchema,
execute: async (input, { toolCallId }) => dispatch(ctx, "research", toolCallId, RESEARCH_ACK, input), execute: async (input, { toolCallId }) => dispatch(ctx, "research", toolCallId, RESEARCH_ACK, input),
}), }),
demo: tool({ demo: lhTool({
name: "demo",
description: DEMO_DESCRIPTION, description: DEMO_DESCRIPTION,
inputSchema: demoInputSchema, schema: demoInputSchema,
execute: async (input, { toolCallId }) => dispatch(ctx, "demo", toolCallId, DEMO_ACK, input), execute: async (input, { toolCallId }) => dispatch(ctx, "demo", toolCallId, DEMO_ACK, input),
}), }),
visual: tool({ visual: lhTool({
name: "visual",
description: VISUAL_DESCRIPTION, description: VISUAL_DESCRIPTION,
inputSchema: visualInputSchema, schema: visualInputSchema,
execute: async (input, { toolCallId }) => dispatch(ctx, "visual", toolCallId, VISUAL_ACK, input), execute: async (input, { toolCallId }) => dispatch(ctx, "visual", toolCallId, VISUAL_ACK, input),
}), }),
}; };
+5 -4
View File
@@ -1,5 +1,5 @@
import { tool, type ToolSet } from "ai";
import { z } from "zod"; import { z } from "zod";
import { lhTool, type LhTool } from "../lh-tool.ts";
import type { PlanToolResult, PlanTree } from "../../core/contracts.ts"; import type { PlanToolResult, PlanTree } from "../../core/contracts.ts";
import type { ChatStore } from "../../core/ports.ts"; import type { ChatStore } from "../../core/ports.ts";
import { createTree, emptyTree, PlanTreeError } from "../../core/plan-tree.ts"; import { createTree, emptyTree, PlanTreeError } from "../../core/plan-tree.ts";
@@ -37,16 +37,17 @@ export interface PlanToolContext {
* latest state straight from the transcript's tool outputs (the plan * latest state straight from the transcript's tool outputs (the plan
* window and the plan blocks). * window and the plan blocks).
*/ */
export function planTools(ctx: PlanToolContext): ToolSet { export function planTools(ctx: PlanToolContext): Record<string, LhTool<any, any>> {
return { return {
plan_init: tool({ plan_init: lhTool({
name: "plan_init",
description: description:
"Write or REWRITE this session's plan map — the dependency graph the learner follows live in the plan window. " + "Write or REWRITE this session's plan map — the dependency graph the learner follows live in the plan window. " +
"Call once right after the probe concludes, to present the map; re-call whenever the map is wrong (the new block replaces the old tree). " + "Call once right after the probe concludes, to present the map; re-call whenever the map is wrong (the new block replaces the old tree). " +
"A mermaid graph block (any direction): one node per line (`slug[Short Title]` — slugs are permanent lowercase ids: a-z, 0-9, hyphens/underscores), " + "A mermaid graph block (any direction): one node per line (`slug[Short Title]` — slugs are permanent lowercase ids: a-z, 0-9, hyphens/underscores), " +
"`-->` edges meaning 'from must be understood before to' (chains fine). 5–15 nodes: a map, not a syllabus — one reasoning step per node, compress. " + "`-->` edges meaning 'from must be understood before to' (chains fine). 5–15 nodes: a map, not a syllabus — one reasoning step per node, compress. " +
"Every check you fire names the node it assesses via `nodeId`; where the lesson stands derives from those tags.", "Every check you fire names the node it assesses via `nodeId`; where the lesson stands derives from those tags.",
inputSchema: z.object({ schema: z.object({
mermaid: z mermaid: z
.string() .string()
.min(1) .min(1)
+5 -4
View File
@@ -1,5 +1,5 @@
import { tool, type ToolSet } from "ai";
import { z } from "zod"; import { z } from "zod";
import { lhTool, type LhTool } from "../lh-tool.ts";
import type { ProbeDoneInput, ProbeDoneResult } from "../../core/contracts.ts"; import type { ProbeDoneInput, ProbeDoneResult } from "../../core/contracts.ts";
import type { UIMessage } from "ai"; import type { UIMessage } from "ai";
import { findPendingToolCall } from "../../core/pending.ts"; import { findPendingToolCall } from "../../core/pending.ts";
@@ -44,9 +44,10 @@ function probeConcluded(messages: UIMessage[]): boolean {
* The first probe_done in the transcript is also the construction seam: * The first probe_done in the transcript is also the construction seam:
* from that message on, every request builds the TEACHER agent. * from that message on, every request builds the TEACHER agent.
*/ */
export function probeDoneTool(ctx: ProbeDoneContext): ToolSet { export function probeDoneTool(ctx: ProbeDoneContext): Record<string, LhTool<any, any>> {
return { return {
probe_done: tool({ probe_done: lhTool({
name: "probe_done",
description: description:
"Conclude the probing phase and hand off to the teaching phase. Call this ONCE, when the probe is " + "Conclude the probing phase and hand off to the teaching phase. Call this ONCE, when the probe is " +
"complete: every strand the lesson depends on is bracketed (a floor he reliably gets right AND a " + "complete: every strand the lesson depends on is bracketed (a floor he reliably gets right AND a " +
@@ -55,7 +56,7 @@ export function probeDoneTool(ctx: ProbeDoneContext): ToolSet {
"(floor, ceiling, misconceptions, what to build from). It streams to the learner live and becomes " + "(floor, ceiling, misconceptions, what to build from). It streams to the learner live and becomes " +
"the handoff artifact the teacher plans from. Refused while a question is still open, and after " + "the handoff artifact the teacher plans from. Refused while a question is still open, and after " +
"the probe is already concluded. After it succeeds, say your closing sentence and end your turn.", "the probe is already concluded. After it succeeds, say your closing sentence and end your turn.",
inputSchema: z.object({ schema: z.object({
assessment: z assessment: z
.object({ .object({
strands: z strands: z
+7 -5
View File
@@ -1,5 +1,5 @@
import { tool } from "ai";
import { z } from "zod"; import { z } from "zod";
import { lhTool } from "../lh-tool.ts";
import { Readability } from "@mozilla/readability"; import { Readability } from "@mozilla/readability";
import { parseHTML } from "linkedom"; import { parseHTML } from "linkedom";
import TurndownService from "turndown"; import TurndownService from "turndown";
@@ -137,11 +137,12 @@ async function fetchAndExtract(url: string): Promise<string> {
// ── tool registry ──────────────────────────────────────────────────────────── // ── tool registry ────────────────────────────────────────────────────────────
export const webTools = { export const webTools = {
web_search: tool({ web_search: lhTool({
name: "web_search",
description: description:
"Search the web via Google Custom Search. One search per call — use one call per search angle. " + "Search the web via Google Custom Search. One search per call — use one call per search angle. " +
"Returns an array of results with title, url, and snippet.", "Returns an array of results with title, url, and snippet.",
inputSchema: z.object({ schema: z.object({
query: z.string().describe("The search query (a normal search string, phrases included)"), query: z.string().describe("The search query (a normal search string, phrases included)"),
count: z.number().int().min(1).max(10).optional().describe("Number of results (default 5, max 10)"), count: z.number().int().min(1).max(10).optional().describe("Number of results (default 5, max 10)"),
}), }),
@@ -161,11 +162,12 @@ export const webTools = {
}, },
}), }),
web_fetch: tool({ web_fetch: lhTool({
name: "web_fetch",
description: description:
"Fetch a web page and extract its readable content as markdown (Readability). " + "Fetch a web page and extract its readable content as markdown (Readability). " +
"Handles HTML pages and plain text; JS-rendered pages fall back to a reader service.", "Handles HTML pages and plain text; JS-rendered pages fall back to a reader service.",
inputSchema: z.object({ schema: z.object({
url: z.string().describe("The URL to fetch"), url: z.string().describe("The URL to fetch"),
}), }),
execute: async ({ url }): Promise<string> => { execute: async ({ url }): Promise<string> => {
+18
View File
@@ -68,6 +68,24 @@ export interface UsageStepLike {
outputTokenDetails?: { reasoningTokens?: number } | undefined; outputTokenDetails?: { reasoningTokens?: number } | undefined;
} }
/** pi-ai's Usage shape (input/output/cacheRead/cacheWrite/reasoning?) →
* UsageStepLike. Structural — no pi import here; the engine band maps at
* its boundaries. */
export function piUsageToStep(usage: {
input: number;
output: number;
cacheRead: number;
cacheWrite: number;
reasoning?: number;
}): UsageStepLike {
return {
inputTokens: usage.input,
outputTokens: usage.output,
inputTokenDetails: { cacheReadTokens: usage.cacheRead, cacheWriteTokens: usage.cacheWrite },
outputTokenDetails: usage.reasoning !== undefined ? { reasoningTokens: usage.reasoning } : undefined,
};
}
/** Fold one model call's usage into the run's running total (pure). A /** Fold one model call's usage into the run's running total (pure). A
* field neither the total nor the step reported stays undefined; once * field neither the total nor the step reported stays undefined; once
* reported, later unreported calls contribute 0. The `steps` counter is * reported, later unreported calls contribute 0. The `steps` counter is
+64 -31
View File
@@ -1,18 +1,21 @@
import { createAgentUIStreamResponse, InvalidToolInputError, NoSuchToolError, type LanguageModelUsage, type UIMessage } from "ai"; import type { UIMessage } from "ai";
import type { Usage } from "@earendil-works/pi-ai";
import { createProberAgent, createTeacherAgent, type TeachingAgent } from "../agent/teacher-agent.ts"; import { createProberAgent, createTeacherAgent, type TeachingAgent } from "../agent/teacher-agent.ts";
import { createJobRunner } from "../agent/subagents/runner.ts"; import { createJobRunner } from "../agent/subagents/runner.ts";
import { createExtractionPipeline } from "../agent/subagents/extraction.ts"; import { createExtractionPipeline } from "../agent/subagents/extraction.ts";
import { MODELS } from "../../config.ts"; import { MODELS } from "../../config.ts";
import { claimSteers, deliverJobResult, setJobRespawn, setTurnRunner } from "./jobs.ts"; import { claimSteers, deliverJobResult, 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 { stores } from "./stores.ts";
import { beginTurn, endTurn } from "./turn-liveness.ts"; import { beginTurn, endTurn } from "./turn-liveness.ts";
import { persistTurn, triggerExtraction, withMessageIds } from "./turn-service.ts"; import { persistTurn, triggerExtraction, withMessageIds } from "./turn-service.ts";
import type { CallUsage } from "../core/contracts.ts"; import type { CallUsage } from "../core/contracts.ts";
/** /**
* One chat turn: run the teacher agent over the client's message list and * One chat turn: run the teaching agent over the client's message list and
* stream the UI message stream back. Framework-free (plain Request/Response) * stream the UI message stream back. Framework-free (plain Request/Response)
* so tests can drive it directly with a scripted mock model — the Next route * so tests can drive it directly with a scripted faux engine — the Next route
* handler in app/api/chat is a one-line wrapper. * handler in app/api/chat is a one-line wrapper.
* *
* This module is the TURN COMPOSITION ROOT: it wires the storage ports, the * This module is the TURN COMPOSITION ROOT: it wires the storage ports, the
@@ -21,6 +24,11 @@ import type { CallUsage } from "../core/contracts.ts";
* THE SAME COMPOSITION serves both entry points: client requests and the * THE SAME COMPOSITION serves both entry points: client requests and the
* wake pipeline's server-side steers (src/server/jobs.ts) — one code path * wake pipeline's server-side steers (src/server/jobs.ts) — one code path
* for write-ahead persistence, turn liveness, and the turn services. * for write-ahead persistence, turn liveness, and the turn services.
*
* Engine note (docs/pi-engine-swap.md): the agent is STATELESS per turn —
* built fresh from the persisted record (recordToPiMessages), run by the
* bridge (pi/ui-stream.ts), discarded at turn end. The record is the only
* source of truth.
*/ */
const jobRunner = createJobRunner({ const jobRunner = createJobRunner({
@@ -48,7 +56,27 @@ setTurnRunner(async (chatId, messages) => {
}); });
setJobRespawn((jobId) => void jobRunner.runJob(jobId)); setJobRespawn((jobId) => void jobRunner.runJob(jobId));
export async function handleChatRequest(req: Request, agent?: TeachingAgent): Promise<Response> { /** The message-shape gate the SDK's UIMessageValidator owned: a malformed
* body must fail BEFORE the stream exists — the path where no onFinish will
* ever run (the catch clears the turn registry so the chat never locks). */
function assertValidUiMessages(messages: UIMessage[]): void {
for (const m of messages) {
if (m.role !== "user" && m.role !== "assistant" && m.role !== "system") {
throw new Error(`invalid message role: ${String((m as { role?: unknown }).role)}`);
}
if (!Array.isArray(m.parts)) throw new Error("message has no parts array");
for (const p of m.parts) {
if (p.type === "text" && typeof (p as { text?: unknown }).text !== "string") {
throw new Error("text part has no text");
}
}
}
}
export async function handleChatRequest(
req: Request,
buildAgent?: (outgoing: UIMessage[]) => TeachingAgent,
): Promise<Response> {
let body: { chatId?: unknown; messages?: unknown }; let body: { chatId?: unknown; messages?: unknown };
try { try {
body = await req.json(); body = await req.json();
@@ -95,19 +123,20 @@ export async function handleChatRequest(req: Request, agent?: TeachingAgent): Pr
} }
// MONEY CHANNEL (wave 4): one collector per request, fed by the agent's // 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 // message_end events. Every model call's usage lands in the chat record's
// append-only turns log at onFinish — idempotent per turn id, so a // append-only turns log at onFinish — idempotent per turn id, so a
// replayed persist appends once while continuation turns (which share a // replayed persist appends once while continuation turns (which share a
// response message id — the SDK coalesces them) each keep their usage. // response message id) each keep their usage.
const turnUsageId = crypto.randomUUID(); const turnUsageId = crypto.randomUUID();
const steps: CallUsage[] = []; const steps: CallUsage[] = [];
const collectUsage = (step: { usage: LanguageModelUsage }) => { const collectUsage = (usage: Usage) => {
steps.push({ steps.push({
inputTokens: step.usage.inputTokens, inputTokens: usage.input,
outputTokens: step.usage.outputTokens, outputTokens: usage.output,
cacheReadTokens: step.usage.inputTokenDetails?.cacheReadTokens, cacheReadTokens: usage.cacheRead,
cacheWriteTokens: step.usage.inputTokenDetails?.cacheWriteTokens, cacheWriteTokens: usage.cacheWrite,
reasoningTokens: step.usage.outputTokenDetails?.reasoningTokens, // Optional in pi's Usage — persisted as 0, never dropped.
reasoningTokens: usage.reasoning ?? 0,
}); });
}; };
@@ -119,7 +148,12 @@ export async function handleChatRequest(req: Request, agent?: TeachingAgent): Pr
// steers, usage, extraction) — only the agent differs. // steers, usage, extraction) — only the agent differs.
const teaching = outgoing.some((m) => m.parts.some((p) => p.type === "tool-probe_done")); const teaching = outgoing.some((m) => m.parts.some((p) => p.type === "tool-probe_done"));
const modelId = teaching ? MODELS.teacher.model : MODELS.prober.model; const modelId = teaching ? MODELS.teacher.model : MODELS.prober.model;
const constructed: TeachingAgent = teaching // Tests inject a factory: the agent is stateless per turn (docs/
// pi-engine-swap.md §2), so each request builds a fresh one off the
// transcript — the factory mirrors that seam exactly.
const constructed: TeachingAgent = buildAgent
? buildAgent(outgoing)
: teaching
? createTeacherAgent( ? createTeacherAgent(
{ {
chatId, chatId,
@@ -127,8 +161,7 @@ export async function handleChatRequest(req: Request, agent?: TeachingAgent): Pr
jobStore: stores.jobStore, jobStore: stores.jobStore,
spawn: (job, input) => void jobRunner.runJob(job.id, input), spawn: (job, input) => void jobRunner.runJob(job.id, input),
}, },
undefined, { messages: recordToPiMessages(outgoing), onUsage: collectUsage },
collectUsage,
) )
: createProberAgent( : createProberAgent(
{ {
@@ -136,32 +169,32 @@ export async function handleChatRequest(req: Request, agent?: TeachingAgent): Pr
chatStore: stores.chatStore, chatStore: stores.chatStore,
deliver: (steer) => void deliverJobResult(steer), deliver: (steer) => void deliverJobResult(steer),
}, },
undefined, { messages: recordToPiMessages(outgoing), onUsage: collectUsage },
collectUsage,
); );
// A test-injected agent gets no usage collector (tests assert the
// transcript, not the money channel); when it's ours, usage flows through
// the factory's subscription (wired above via onUsage).
try { try {
return await createAgentUIStreamResponse({ assertValidUiMessages(outgoing);
agent: agent ?? constructed, return createPiTurnResponse({
agent: constructed,
uiMessages: outgoing, uiMessages: outgoing,
// Tool-call validation failures are messages OUR schemas authored (the signal: req.signal,
// quiz authoring contract, unknown tools) — safe and useful in the // Stream-level failures are provider/transport internals — masked,
// transcript: they are exactly what the teacher saw and retried on. The // like the SDK defaults. Tool-input validation messages never
// tool loop delivers them as plain strings (what the model received) or // route through here: they ride the error-result path and reach the
// as the original error object; expose both, and keep everything else // transcript directly (the retry context the model saw).
// (provider/transport internals) masked, as the SDK defaults.
onError: (error) => { onError: (error) => {
if (typeof error === "string") return error; console.error(`[chat ${chatId}] turn error:`, error);
if (InvalidToolInputError.isInstance(error) || NoSuchToolError.isInstance(error)) return error.message;
return "An error occurred."; return "An error occurred.";
}, },
onFinish: ({ messages, isAborted }) => { onFinish: (messages, isAborted) => {
persistTurn(stores.chatStore, chatId, messages as UIMessage[], isAborted); persistTurn(stores.chatStore, chatId, messages as UIMessage[], isAborted);
// usage log append (idempotent per turn id) — the stats bar // usage log append (idempotent per turn id) — the stats bar
// refolds this on the client after the turn settles. messageId // refolds this on the client after the turn settles. messageId
// comes from the id-FILLED list: onFinish messages carry no ids // comes from the id-FILLED list: the bridge's response message
// (withMessageIds fills copies in persistTurn, not these); the // already carries its id; the re-fill here is idempotent.
// re-fill here is idempotent, so the key is stable.
const filled = withMessageIds(messages as UIMessage[]); const filled = withMessageIds(messages as UIMessage[]);
const response = [...filled].reverse().find((m) => m.role === "assistant"); const response = [...filled].reverse().find((m) => m.role === "assistant");
if (response && steps.length > 0) { if (response && steps.length > 0) {
+163
View File
@@ -0,0 +1,163 @@
import type { UIMessage } from "ai";
import type { AgentMessage } from "@earendil-works/pi-agent-core";
import type { AssistantMessage, Message, ToolResultMessage } from "@earendil-works/pi-ai";
/**
* The record → pi-ai transcript adapter (docs/pi-engine-swap.md §4). The
* persisted UIMessage record IS the source of truth; every turn builds a
* fresh agent from it (the stateless per-turn decision). This replaces the
* SDK's convertToModelMessages with an owned, explicit mapping:
*
* - user text parts → UserMessage text content; data parts are dropped
* (the plan-accept seam rides its text twin — the client always sends
* "The plan is accepted — begin teaching." alongside the marker).
* - assistant text/reasoning/tool parts → one AssistantMessage whose
* content blocks mirror the part order; tool parts become ToolCall
* blocks (output-error parts keep their rawInput).
* - answered tool parts (output-available) → ToolResultMessage after the
* assistant message, JSON-encoded like the model saw it live.
* - errored tool parts (output-error) → an isError ToolResultMessage
* carrying the validation message (the model's retry context).
* - DANGLING calls (input-available, never answered — an aborted turn)
* → a synthesized interrupted ToolResultMessage. THE HARD INVARIANT:
* an assistant message with tool calls must be followed by tool
* results, or providers reject the request outright.
*/
const INTERRUPTED = "The question was interrupted — the learner never answered; continue without holding a result for it.";
function textOf(part: { text?: string }): string {
return typeof part.text === "string" ? part.text : "";
}
/** Encode a tool output the way the model saw it live (string verbatim, object JSON). */
function toolResultText(output: unknown): string {
return typeof output === "string" ? output : JSON.stringify(output);
}
function zeroUsage() {
return {
input: 0,
output: 0,
cacheRead: 0,
cacheWrite: 0,
totalTokens: 0,
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
};
}
export function recordToPiMessages(messages: UIMessage[]): Message[] {
const out: AgentMessage[] = [];
for (const message of messages) {
if (message.role === "user") {
const texts = message.parts
.filter((p) => p.type === "text")
.map((p) => textOf(p as { text?: string }))
.filter((t) => t.length > 0);
if (texts.length === 0) continue; // pure data-part marker, no text
out.push({
role: "user",
content: texts.map((text) => ({ type: "text" as const, text })),
timestamp: Date.now(),
});
} else if (message.role === "assistant") {
const content: AssistantMessage["content"] = [];
const results: ToolResultMessage[] = [];
let answered = false;
for (const part of message.parts) {
if (part.type === "text") {
const text = textOf(part);
if (text) content.push({ type: "text", text });
} else if (part.type === "reasoning") {
const text = textOf(part);
if (text) content.push({ type: "thinking", thinking: text });
} else if (part.type === "step-start") {
// presentation-only; the pi loop re-derives turn boundaries
} else if (part.type.startsWith("tool-")) {
const tool = part.type.slice("tool-".length);
const p = part as unknown as {
toolCallId: string;
state: string;
input?: unknown;
rawInput?: unknown;
output?: unknown;
errorText?: string;
};
// The call block (any state — the model authored it).
if (p.state === "input-available" || p.state === "input-streaming" || p.state === "output-available") {
content.push({
type: "toolCall",
id: p.toolCallId,
name: tool,
arguments: (p.input ?? {}) as Record<string, unknown>,
});
} else if (p.state === "output-error") {
content.push({
type: "toolCall",
id: p.toolCallId,
name: tool,
arguments: (p.rawInput ?? p.input ?? {}) as Record<string, unknown>,
});
}
// The result, if any.
if (p.state === "output-available") {
results.push({
role: "toolResult",
toolCallId: p.toolCallId,
toolName: tool,
content: [{ type: "text", text: toolResultText(p.output) }],
isError: false,
timestamp: Date.now(),
});
answered = true;
} else if (p.state === "output-error") {
results.push({
role: "toolResult",
toolCallId: p.toolCallId,
toolName: tool,
content: [{ type: "text", text: p.errorText ?? "The tool call failed validation." }],
isError: true,
timestamp: Date.now(),
});
answered = true;
}
}
// data-* parts (plan blocks, artifacts, markers) never reach the model.
}
if (content.length === 0) continue; // presentation-only message
out.push({
role: "assistant",
content,
api: "openai-completions",
provider: "openrouter",
model: "replayed",
usage: zeroUsage(),
stopReason: content[content.length - 1]!.type === "toolCall" ? "toolUse" : "stop",
timestamp: Date.now(),
});
// Dangling calls: synthesize the interrupted result so providers
// never see an unanswered tool call (the hard invariant).
if (!answered) {
for (const block of content) {
if (block.type === "toolCall" && !results.some((r) => r.toolCallId === block.id)) {
results.push({
role: "toolResult",
toolCallId: block.id,
toolName: block.name,
content: [{ type: "text", text: INTERRUPTED }],
isError: false,
timestamp: Date.now(),
});
}
}
}
out.push(...results);
}
// role "system" never appears in records (system prompt is composed separately).
}
return out as Message[];
}
+346
View File
@@ -0,0 +1,346 @@
import type { Agent, AgentEvent } from "@earendil-works/pi-agent-core";
import type { UIMessage } from "ai";
import { QUESTION_TOOLS } from "../../core/pending.ts";
/**
* The UI-stream bridge (docs/pi-engine-swap.md §5): runs one agent turn and
* re-emits it as an AI-SDK-shaped UIMessageStream SSE response — the wire
* protocol `useChat` (Band D, unchanged) parses:
*
* data: {"type":"text-delta", …}\n\n … data: [DONE]
*
* Load-bearing protocol rules (verified against ai@7's client):
* - The `start` chunk's messageId MUST be the transcript's last message id
* when it is an assistant message (the continuation id rule) — a fresh id
* there duplicates every prior part client-side.
* - Question tools (quiz / ask_user_question): emit tool-input-available and
* NEVER tool-output-available (the client throws on an output for a call it
* owns). The synthetic "question posed" result lives in memory only; the
* record's part stays output-less until the learner answers.
* - Errored calls (authoring validation) surface as tool-input-error — the
* exact part shape the SDK produced for InvalidToolInputError (the
* output-error state the chat flow asserts).
* - The persisted message is assembled from the SAME parts the client
* assembles (the processUIMessageStream shapes), so every downstream
* consumer (phases, plan blocks, pending, extraction) is unchanged.
*/
/** SSE chunk — the flat JSON line shape the client's zod schema accepts. */
type Chunk = Record<string, unknown> & { type: string };
export interface PiTurnSpec {
/** The per-turn agent (stateless engine — built from the record). */
agent: Agent;
/** The outgoing request messages (write-ahead persisted already). */
uiMessages: UIMessage[];
/** Request abort signal — aborts the agent mid-run. */
signal?: AbortSignal;
/** The handler's error-unmasking policy (validation messages stay readable). */
onError?: (error: unknown) => string;
/** The canonical end-of-turn seam: persist + usage + extraction + endTurn. */
onFinish: (messages: UIMessage[], isAborted: boolean) => void;
}
/** Map a pi stop reason to the finish chunk's enum. */
function finishReasonOf(stopReason: string | undefined): string {
switch (stopReason) {
case "stop":
return "stop";
case "length":
return "length";
case "toolUse":
return "tool-calls";
case "error":
return "error";
default:
return "other";
}
}
function resultErrorText(result: unknown): string {
const content = (result as { content?: Array<{ type: string; text?: string }> } | undefined)?.content;
const text = content?.find((c) => c.type === "text")?.text;
return text ?? "The tool call failed.";
}
/** Loose part shapes the bridge assembles (the client's part shapes). */
interface TextPart {
type: "text";
text: string;
state: string;
}
interface ReasoningPart {
type: "reasoning";
id: string;
text: string;
state: string;
}
interface ToolPart {
type: string;
toolCallId: string;
state: string;
input?: unknown;
rawInput?: unknown;
output?: unknown;
errorText?: string;
}
export function createPiTurnResponse(spec: PiTurnSpec): Response {
const { agent, uiMessages, onFinish } = spec;
// The streaming target: continuation (same id, parts appended to a copy
// of the existing message) or a fresh assistant message.
const last = uiMessages[uiMessages.length - 1];
const continuing = last?.role === "assistant";
const messageId =
(continuing && last.id) || `msg_${Date.now().toString(36)}_${crypto.randomUUID().slice(0, 8)}`;
const baseParts: UIMessage["parts"] = continuing ? [...last.parts] : [];
const parts: UIMessage["parts"] = [...baseParts];
// Stable, collision-free part ids across continuation runs (reasoning
// parts persist their id; text ids only need per-run uniqueness).
let textN = baseParts.filter((p) => p.type === "text").length;
let reasoningN = baseParts.filter((p) => p.type === "reasoning").length;
const textIds = new Map<number, string>(); // contentIndex → chunk id
const reasoningIds = new Map<number, string>();
const toolInputs = new Map<string, { toolName: string; input: unknown }>();
/** The streaming text/reasoning parts of THIS run, keyed by chunk id. */
const activeParts = new Map<string, TextPart | ReasoningPart>();
/** Tool parts created this run, keyed by toolCallId. */
const newToolParts = new Map<string, ToolPart>();
let settled = false;
let sawAbort = false;
let finishReason: string | undefined;
let errorMessage: string | undefined;
const stream = new ReadableStream<Uint8Array>({
start(controller) {
const encoder = new TextEncoder();
// Swallow enqueue on a cancelled stream (the client went away): the
// run must still settle — onFinish persists the partial turn.
const send = (chunk: Chunk) => {
if (settled) return;
try {
controller.enqueue(encoder.encode(`data: ${JSON.stringify(chunk)}\n\n`));
} catch {
/* reader cancelled — stop emitting, keep running */
}
};
const close = () => {
if (settled) return;
settled = true;
try {
controller.enqueue(encoder.encode("data: [DONE]\n\n"));
controller.close();
} catch {
/* already cancelled */
}
};
send({ type: "start", messageId });
const onAbort = () => {
sawAbort = true;
agent.abort();
};
spec.signal?.addEventListener("abort", onAbort, { once: true });
const settleWith = () => {
const assembled: UIMessage = { id: messageId, role: "assistant", parts };
// The SDK's onFinish shape: the input copy stays, the extended
// response copy appends (persistTurn's withMessageIds dedupe
// keeps the last occurrence).
onFinish([...uiMessages, assembled], sawAbort || spec.signal?.aborted === true);
};
const handle = (event: AgentEvent) => {
switch (event.type) {
case "turn_start":
send({ type: "start-step" });
parts.push({ type: "step-start" });
break;
case "message_update": {
const streamEvent = event.assistantMessageEvent;
switch (streamEvent.type) {
case "text_start": {
const id = `${messageId}-t${textN++}`;
textIds.set(streamEvent.contentIndex, id);
const part: TextPart = { type: "text", text: "", state: "streaming" };
activeParts.set(id, part);
parts.push(part as UIMessage["parts"][number]);
send({ type: "text-start", id });
break;
}
case "text_delta": {
const id = textIds.get(streamEvent.contentIndex);
const part = id !== undefined ? (activeParts.get(id) as TextPart | undefined) : undefined;
if (part) part.text += streamEvent.delta;
if (id !== undefined) send({ type: "text-delta", id, delta: streamEvent.delta });
break;
}
case "text_end": {
const id = textIds.get(streamEvent.contentIndex);
const part = id !== undefined ? (activeParts.get(id) as TextPart | undefined) : undefined;
if (part) part.state = "done";
activeParts.delete(id ?? "");
if (id !== undefined) send({ type: "text-end", id });
break;
}
case "thinking_start": {
const id = `${messageId}-r${reasoningN++}`;
reasoningIds.set(streamEvent.contentIndex, id);
const part: ReasoningPart = { type: "reasoning", id, text: "", state: "streaming" };
activeParts.set(id, part);
parts.push(part as UIMessage["parts"][number]);
send({ type: "reasoning-start", id });
break;
}
case "thinking_delta": {
const id = reasoningIds.get(streamEvent.contentIndex);
const part = id !== undefined ? (activeParts.get(id) as ReasoningPart | undefined) : undefined;
if (part) part.text += streamEvent.delta;
if (id !== undefined) send({ type: "reasoning-delta", id, delta: streamEvent.delta });
break;
}
case "thinking_end": {
const id = reasoningIds.get(streamEvent.contentIndex);
const part = id !== undefined ? (activeParts.get(id) as ReasoningPart | undefined) : undefined;
if (part) part.state = "done";
activeParts.delete(id ?? "");
if (id !== undefined) send({ type: "reasoning-end", id });
break;
}
case "toolcall_start": {
const block = (event.message as { content: any[] }).content[streamEvent.contentIndex];
if (block?.type !== "toolCall") break;
toolInputs.set(block.id, { toolName: block.name, input: undefined });
const part: ToolPart = {
type: `tool-${block.name}`,
toolCallId: block.id,
state: "input-streaming",
input: undefined,
};
newToolParts.set(block.id, part);
parts.push(part as unknown as UIMessage["parts"][number]);
send({ type: "tool-input-start", toolCallId: block.id, toolName: block.name });
break;
}
case "toolcall_delta": {
const block = (event.message as { content: any[] }).content[streamEvent.contentIndex];
if (block?.type !== "toolCall") break;
send({
type: "tool-input-delta",
toolCallId: block.id,
inputTextDelta: streamEvent.delta,
});
break;
}
case "toolcall_end": {
const toolCall = streamEvent.toolCall;
toolInputs.set(toolCall.id, { toolName: toolCall.name, input: toolCall.arguments });
const part = newToolParts.get(toolCall.id);
if (part) {
part.state = "input-available";
part.input = toolCall.arguments;
}
send({
type: "tool-input-available",
toolCallId: toolCall.id,
toolName: toolCall.name,
input: toolCall.arguments,
});
break;
}
}
break;
}
case "tool_execution_end": {
const known = toolInputs.get(event.toolCallId);
const toolName = event.toolName || known?.toolName || "";
const input = known?.input;
const part = newToolParts.get(event.toolCallId);
if (event.isError) {
const errorText = resultErrorText(event.result);
if (part) {
part.state = "output-error";
part.rawInput = input;
part.errorText = errorText;
}
send({
type: "tool-input-error",
toolCallId: event.toolCallId,
toolName,
input,
errorText,
});
} else if ((QUESTION_TOOLS as readonly string[]).includes(toolName)) {
// Client-owned: the answer arrives via addToolOutput;
// the synthetic placeholder result never reaches wire
// or record.
break;
} else {
const output = (event.result as { details?: unknown } | undefined)?.details;
if (part) {
part.state = "output-available";
part.output = output;
}
send({ type: "tool-output-available", toolCallId: event.toolCallId, output });
}
break;
}
case "turn_end":
send({ type: "finish-step" });
break;
case "message_end":
if (event.message && (event.message as { role?: string }).role === "assistant") {
const m = event.message as { stopReason?: string; errorMessage?: string };
finishReason = m.stopReason;
if (m.stopReason === "aborted") sawAbort = true;
if (m.stopReason === "error") errorMessage = m.errorMessage;
}
break;
case "agent_end": {
spec.signal?.removeEventListener("abort", onAbort);
if (errorMessage !== undefined) {
send({ type: "error", errorText: spec.onError?.(errorMessage) ?? errorMessage });
}
send({ type: "finish", finishReason: finishReasonOf(finishReason) });
close();
settleWith();
break;
}
}
};
agent.subscribe(handle);
void (async () => {
try {
await agent.continue();
} catch (err) {
// A loop-level failure (not a provider error — those ride the
// stream as stopReason "error"): surface, persist, close.
spec.signal?.removeEventListener("abort", onAbort);
send({ type: "error", errorText: spec.onError?.(err) ?? "An error occurred." });
send({ type: "finish", finishReason: "error" });
close();
settleWith();
}
})();
},
cancel() {
// The client went away (reload/abort): kill the model call. The
// run still settles — agent_end fires and onFinish persists.
agent.abort();
},
});
return new Response(stream, {
headers: {
"Content-Type": "text/event-stream",
"Cache-Control": "no-cache, no-transform",
"X-Accel-Buffering": "no",
},
});
}
+91 -117
View File
@@ -1,28 +1,38 @@
/** /**
* Scripted integration test — the full teacher loop with zero tokens: * Scripted integration test — the full teacher loop with zero tokens:
* *
* learner opens chat → mock model calls `quiz` → server streams + persists * learner opens chat → faux model calls `quiz` → server streams + persists
* → client grades and sends the answer back as the tool output * → client grades and sends the answer back as the tool output
* → mock model reacts → transcript persisted. * → faux model reacts → transcript persisted.
* *
* Drives handleChatRequest directly (plain Request/Response — no Next server * Drives handleChatRequest directly (plain Request/Response — no Next server
* needed) with a MockLanguageModelV4. This is the AI SDK equivalent of the * needed) with a scripted pi engine (pi-ai's faux provider, docs/
* old pi dispatch contract test: the quiz bridge is gone, the tool part IS * pi-engine-swap.md §7): the quiz bridge is gone, the tool part IS the
* the pending state, and the persisted transcript is the resume record. * pending state, and the persisted transcript is the resume record.
*/ */
import assert from "node:assert/strict"; import assert from "node:assert/strict";
import * as fs from "node:fs"; import * as fs from "node:fs";
import * as path from "node:path"; import * as path from "node:path";
import { after, describe, it } from "node:test"; import { after, describe, it } from "node:test";
import { MockLanguageModelV4, simulateReadableStream } from "ai/test"; import {
import type { LanguageModelV4StreamPart } from "@ai-sdk/provider"; createFauxCore,
fauxAssistantMessage,
fauxText,
fauxThinking,
fauxToolCall,
type FauxResponseStep,
} from "@earendil-works/pi-ai";
import { createTeacherAgent } from "../src/agent/teacher-agent.ts"; import { createTeacherAgent } from "../src/agent/teacher-agent.ts";
import { handleChatRequest } from "../src/server/chat-handler.ts"; import { handleChatRequest } from "../src/server/chat-handler.ts";
import { isTurnRunning } from "../src/server/turn-liveness.ts"; import { isTurnRunning } from "../src/server/turn-liveness.ts";
import { recordToPiMessages } from "../src/server/pi/record-adapter.ts";
import { chatStore } from "../src/store/chats.ts"; import { chatStore } from "../src/store/chats.ts";
import { buildAnswer, coerceCorrectAnswer, normalizeOptions, resolveCorrect, type QuestionData } from "../src/core/question-core.ts"; import { buildAnswer, coerceCorrectAnswer, normalizeOptions, resolveCorrect, type QuestionData } from "../src/core/question-core.ts";
import { findPendingToolCall } from "../src/core/pending.ts"; import { findPendingToolCall } from "../src/core/pending.ts";
import type { PiEngine } from "../src/agent/models.ts";
import type { TeachingAgent } from "../src/agent/teacher-agent.ts";
import type { UIMessage } from "ai";
import { PATHS } from "../config.ts"; import { PATHS } from "../config.ts";
const QUIZ_INPUT = { const QUIZ_INPUT = {
@@ -36,52 +46,41 @@ const QUIZ_INPUT = {
explanation: "Mercury is the innermost planet.", explanation: "Mercury is the innermost planet.",
}; };
function usage() { /** Turn 1: a bit of reasoning, then the quiz tool call. Turn 2: plain text.
return { * Each response factory records the prompt it received (the transcript the
inputTokens: { total: 10, noCache: 10, cacheRead: 0, cacheWrite: 0 }, * model saw) — the old doStream capture, engine-shaped. */
outputTokens: { total: 5, text: 5, reasoning: 0 }, function scriptedTeacher(steps: (calls: string[][]) => FauxResponseStep[]): {
}; engine: PiEngine;
calls: string[][];
} {
const calls: string[][] = [];
const faux = createFauxCore({});
faux.setResponses(steps(calls));
return { engine: { model: faux.getModel(), streamFn: faux.streamSimple }, calls };
} }
/** Turn 1: a bit of reasoning, then the quiz tool call. Turn 2: plain text. */ /** The default two-turn lesson: probe quiz, then react to the answer. */
function scriptedModel(calls: string[][]) { function defaultLesson(calls: string[][]): FauxResponseStep[] {
let turn = 0; return [
return new MockLanguageModelV4({ (context) => {
doStream: async (options) => { calls.push(context.messages.map((m) => JSON.stringify((m as { content: unknown }).content)));
const promptTexts = options.prompt.map((m) => JSON.stringify(m.content)); return fauxAssistantMessage(
calls.push(promptTexts); [fauxThinking("Probe the basics first."), fauxToolCall("quiz", QUIZ_INPUT, { id: "call-1" })],
turn++; { stopReason: "toolUse" },
const chunks: LanguageModelV4StreamPart[] = );
turn === 1 },
? [ (context) => {
{ type: "stream-start", warnings: [] }, calls.push(context.messages.map((m) => JSON.stringify((m as { content: unknown }).content)));
{ type: "reasoning-start", id: "r1" }, return fauxAssistantMessage([fauxText("Right — Mercury it is.")]);
{ type: "reasoning-delta", id: "r1", delta: "Probe the basics first." },
{ type: "reasoning-end", id: "r1" },
// Same shape real providers use: input streams as deltas,
// then the completed call. (This is what makes the client's
// 'Writing a question…' indicator possible.)
{ type: "tool-input-start", id: "call-1", toolName: "quiz" },
{ type: "tool-input-delta", id: "call-1", delta: JSON.stringify(QUIZ_INPUT) },
{ type: "tool-input-end", id: "call-1" },
{
type: "tool-call",
toolCallId: "call-1",
toolName: "quiz",
input: JSON.stringify(QUIZ_INPUT),
}, },
{ type: "finish", finishReason: { unified: "tool-calls", raw: "tool_calls" }, usage: usage() },
]
: [
{ type: "stream-start", warnings: [] },
{ type: "text-start", id: "t1" },
{ type: "text-delta", id: "t1", delta: "Right — Mercury it is." },
{ type: "text-end", id: "t1" },
{ type: "finish", finishReason: { unified: "stop", raw: "stop" }, usage: usage() },
]; ];
return { stream: simulateReadableStream({ chunks, initialDelayInMs: 0, chunkDelayInMs: 0 }) }; }
},
}); /** The handler seam (stateless per turn): each request builds a fresh agent
* off the transcript it is about to run — the factory tests inject. */
function agentFactory(engine: PiEngine) {
return (outgoing: UIMessage[]): TeachingAgent =>
createTeacherAgent(undefined, { engine, messages: recordToPiMessages(outgoing) });
} }
async function drain(res: Response): Promise<string> { async function drain(res: Response): Promise<string> {
@@ -112,14 +111,13 @@ after(() => {
describe("teacher chat flow (scripted model)", () => { describe("teacher chat flow (scripted model)", () => {
it("quiz tool call persists as a pending part; the answer continues the lesson", async () => { it("quiz tool call persists as a pending part; the answer continues the lesson", async () => {
const calls: string[][] = []; const { engine, calls } = scriptedTeacher(defaultLesson);
const agent = createTeacherAgent(undefined, scriptedModel(calls));
const chat = chatStore.createChat("Test: solar system"); const chat = chatStore.createChat("Test: solar system");
createdChats.push(chat.id); createdChats.push(chat.id);
// ── Turn 1: learner opens with a topic; teacher probes with a quiz ── // ── Turn 1: learner opens with a topic; teacher probes with a quiz ──
const firstUser = { id: "u1", role: "user", parts: [{ type: "text", text: "teach me the solar system" }] }; const firstUser = { id: "u1", role: "user", parts: [{ type: "text", text: "teach me the solar system" }] };
await drain(await handleChatRequest(post(chat.id, [firstUser]), agent)); await drain(await handleChatRequest(post(chat.id, [firstUser]), agentFactory(engine)));
const saved1 = chatStore.getChat(chat.id)!; const saved1 = chatStore.getChat(chat.id)!;
assert.equal(saved1.messages.length, 2, "user message + assistant response persisted"); assert.equal(saved1.messages.length, 2, "user message + assistant response persisted");
@@ -128,12 +126,12 @@ describe("teacher chat flow (scripted model)", () => {
const reasoningPart = assistantMsg.parts.find((p) => p.type === "reasoning"); const reasoningPart = assistantMsg.parts.find((p) => p.type === "reasoning");
assert.ok(reasoningPart && reasoningPart.type === "reasoning", "reasoning part streamed + persisted"); assert.ok(reasoningPart && reasoningPart.type === "reasoning", "reasoning part streamed + persisted");
assert.match(reasoningPart.text, /Probe the basics/); assert.match((reasoningPart as { text?: string }).text!, /Probe the basics/);
const quizPart = assistantMsg.parts.find((p) => p.type === "tool-quiz"); const quizPart = assistantMsg.parts.find((p) => p.type === "tool-quiz");
assert.ok(quizPart && quizPart.type === "tool-quiz", "quiz tool part persisted"); assert.ok(quizPart && quizPart.type === "tool-quiz", "quiz tool part persisted");
assert.equal(quizPart.state, "input-available", "unanswered quiz = pending state (resume record)"); assert.equal((quizPart as { state?: string }).state, "input-available", "unanswered quiz = pending state (resume record)");
assert.deepEqual(quizPart.input, QUIZ_INPUT); assert.deepEqual((quizPart as { input?: unknown }).input, QUIZ_INPUT);
// The pending-question gate sees it. // The pending-question gate sees it.
const pending = findPendingToolCall(saved1.messages); const pending = findPendingToolCall(saved1.messages);
@@ -166,7 +164,7 @@ describe("teacher chat flow (scripted model)", () => {
} }
: m, : m,
); );
await drain(await handleChatRequest(post(chat.id, messages2), agent)); await drain(await handleChatRequest(post(chat.id, messages2), agentFactory(engine)));
// ── The teacher reacts — as a CONTINUATION of the same assistant message // ── The teacher reacts — as a CONTINUATION of the same assistant message
// (the last client message is the assistant message carrying the tool // (the last client message is the assistant message carrying the tool
@@ -176,14 +174,14 @@ describe("teacher chat flow (scripted model)", () => {
const finalParts = saved2.messages[1]!.parts; const finalParts = saved2.messages[1]!.parts;
const textPart = finalParts.find((p) => p.type === "text"); const textPart = finalParts.find((p) => p.type === "text");
assert.ok(textPart && textPart.type === "text", "reaction text appended after the tool part"); assert.ok(textPart && textPart.type === "text", "reaction text appended after the tool part");
assert.match(textPart.text, /Mercury it is/); assert.match((textPart as { text?: string }).text!, /Mercury it is/);
// The answered quiz part is still there, now with its output. // The answered quiz part is still there, now with its output.
const answeredQuiz = finalParts.find((p) => p.type === "tool-quiz"); const answeredQuiz = finalParts.find((p) => p.type === "tool-quiz");
assert.ok(answeredQuiz && answeredQuiz.type === "tool-quiz" && answeredQuiz.state === "output-available"); assert.ok(answeredQuiz && answeredQuiz.type === "tool-quiz" && (answeredQuiz as { state?: string }).state === "output-available");
// The model actually received the graded answer as a tool result. // The model actually received the graded answer as a tool result.
assert.equal(calls.length, 2, "model saw exactly two prompts"); assert.equal(calls.length, 2, "model saw exactly two prompts");
const toolResultSent = calls[1]!.some((c) => c.includes("tool-result") || c.includes("Mercury")); const toolResultSent = calls[1]!.some((c) => c.includes("toolResult") || c.includes("Mercury"));
assert.ok(toolResultSent, "continuation prompt carries the tool result to the model"); assert.ok(toolResultSent, "continuation prompt carries the tool result to the model");
// No pending question remains. // No pending question remains.
@@ -200,11 +198,11 @@ describe("teacher chat flow (scripted model)", () => {
// 'input-streaming', which the stream processor sets from the // 'input-streaming', which the stream processor sets from the
// tool-input-start chunk. If a transport change ever stopped emitting // tool-input-start chunk. If a transport change ever stopped emitting
// it, the indicator would silently die — pin the protocol. // it, the indicator would silently die — pin the protocol.
const agent = createTeacherAgent(undefined, scriptedModel([])); const { engine } = scriptedTeacher(defaultLesson);
const chat = chatStore.createChat("Test: protocol pin"); const chat = chatStore.createChat("Test: protocol pin");
createdChats.push(chat.id); createdChats.push(chat.id);
const firstUser = { id: "u1", role: "user", parts: [{ type: "text", text: "quiz me" }] }; const firstUser = { id: "u1", role: "user", parts: [{ type: "text", text: "quiz me" }] };
const body = await drain(await handleChatRequest(post(chat.id, [firstUser]), agent)); const body = await drain(await handleChatRequest(post(chat.id, [firstUser]), agentFactory(engine)));
const startIdx = body.indexOf('"type":"tool-input-start"'); const startIdx = body.indexOf('"type":"tool-input-start"');
const availIdx = body.indexOf('"type":"tool-input-available"'); const availIdx = body.indexOf('"type":"tool-input-available"');
assert.ok(startIdx !== -1, "tool-input-start chunk present in the stream"); assert.ok(startIdx !== -1, "tool-input-start chunk present in the stream");
@@ -213,20 +211,19 @@ describe("teacher chat flow (scripted model)", () => {
}); });
it("persists even when the stream is aborted (resume keeps partial state)", async () => { it("persists even when the stream is aborted (resume keeps partial state)", async () => {
const calls: string[][] = []; const { engine } = scriptedTeacher(defaultLesson);
const agent = createTeacherAgent(undefined, scriptedModel(calls));
const chat = chatStore.createChat("Test: abort"); const chat = chatStore.createChat("Test: abort");
createdChats.push(chat.id); createdChats.push(chat.id);
const req = post(chat.id, [{ id: "u1", role: "user", parts: [{ type: "text", text: "quiz me" }] }]); const req = post(chat.id, [{ id: "u1", role: "user", parts: [{ type: "text", text: "quiz me" }] }]);
const res = await handleChatRequest(req, agent); const res = await handleChatRequest(req, agentFactory(engine));
assert.equal(res.status, 200); assert.equal(res.status, 200);
// Read one chunk then cancel — simulates the browser closing mid-stream. // Read one chunk then cancel — simulates the browser closing mid-stream.
const reader = res.body!.getReader(); const reader = res.body!.getReader();
await reader.read(); await reader.read();
await reader.cancel(); await reader.cancel();
// Give onFinish a tick to fire. // Give the aborted run a tick to settle (agent_end → onFinish).
await new Promise((r) => setTimeout(r, 50)); await new Promise((r) => setTimeout(r, 200));
const saved = chatStore.getChat(chat.id)!; const saved = chatStore.getChat(chat.id)!;
assert.ok(saved.messages.length >= 1, "partial transcript persisted"); assert.ok(saved.messages.length >= 1, "partial transcript persisted");
}); });
@@ -234,8 +231,8 @@ describe("teacher chat flow (scripted model)", () => {
it("an ungradeable quiz is rejected at the call boundary, the error reaches the model, and the model retries", async () => { it("an ungradeable quiz is rejected at the call boundary, the error reaches the model, and the model retries", async () => {
// The authoring contract (question-core superRefine): a quiz whose // The authoring contract (question-core superRefine): a quiz whose
// correctAnswer matches no option value must NEVER reach the dock — // correctAnswer matches no option value must NEVER reach the dock —
// the SDK validates client-side tool inputs against the schema, marks // toPiTool re-runs the zod schema at the call boundary, marks the
// the call invalid, and returns the validation error to the teacher as // call invalid, and returns the validation error to the teacher as
// a tool result. The loop re-prompts the model, which retries in the // a tool result. The loop re-prompts the model, which retries in the
// same turn. This pins that whole mechanism (the live Nietzsche bug). // same turn. This pins that whole mechanism (the live Nietzsche bug).
const BAD_QUIZ_INPUT = { const BAD_QUIZ_INPUT = {
@@ -245,49 +242,28 @@ describe("teacher chat flow (scripted model)", () => {
explanation: "Genealogy diagnoses, it does not convert.", explanation: "Genealogy diagnoses, it does not convert.",
}; };
const calls: string[][] = []; const steps = (calls: string[][]): FauxResponseStep[] => [
let step = 0; (context) => {
const retryModel = new MockLanguageModelV4({ calls.push(context.messages.map((m) => JSON.stringify((m as { content: unknown }).content)));
doStream: async (options) => { return fauxAssistantMessage(
calls.push(options.prompt.map((m) => JSON.stringify(m.content))); [fauxToolCall("quiz", BAD_QUIZ_INPUT, { id: "call-bad" })],
step++; { stopReason: "toolUse" },
const chunks: LanguageModelV4StreamPart[] = );
step === 1
? [
{ type: "stream-start", warnings: [] },
{ type: "tool-input-start", id: "call-bad", toolName: "quiz" },
{ type: "tool-input-delta", id: "call-bad", delta: JSON.stringify(BAD_QUIZ_INPUT) },
{ type: "tool-input-end", id: "call-bad" },
{
type: "tool-call",
toolCallId: "call-bad",
toolName: "quiz",
input: JSON.stringify(BAD_QUIZ_INPUT),
}, },
{ type: "finish", finishReason: { unified: "tool-calls", raw: "tool_calls" }, usage: usage() }, (context) => {
] calls.push(context.messages.map((m) => JSON.stringify((m as { content: unknown }).content)));
: [ return fauxAssistantMessage(
{ type: "stream-start", warnings: [] }, [fauxToolCall("quiz", QUIZ_INPUT, { id: "call-good" })],
{ type: "tool-input-start", id: "call-good", toolName: "quiz" }, { stopReason: "toolUse" },
{ type: "tool-input-delta", id: "call-good", delta: JSON.stringify(QUIZ_INPUT) }, );
{ type: "tool-input-end", id: "call-good" },
{
type: "tool-call",
toolCallId: "call-good",
toolName: "quiz",
input: JSON.stringify(QUIZ_INPUT),
}, },
{ type: "finish", finishReason: { unified: "tool-calls", raw: "tool_calls" }, usage: usage() },
]; ];
return { stream: simulateReadableStream({ chunks, initialDelayInMs: 0, chunkDelayInMs: 0 }) }; const { engine, calls } = scriptedTeacher(steps);
},
});
const agent = createTeacherAgent(undefined, retryModel);
const chat = chatStore.createChat("Test: bad quiz retry"); const chat = chatStore.createChat("Test: bad quiz retry");
createdChats.push(chat.id); createdChats.push(chat.id);
const firstUser = { id: "u1", role: "user", parts: [{ type: "text", text: "quiz me" }] }; const firstUser = { id: "u1", role: "user", parts: [{ type: "text", text: "quiz me" }] };
await drain(await handleChatRequest(post(chat.id, [firstUser]), agent)); await drain(await handleChatRequest(post(chat.id, [firstUser]), agentFactory(engine)));
// The model was re-prompted (the error did not end the turn)… // The model was re-prompted (the error did not end the turn)…
assert.equal(calls.length, 2, "the invalid call triggered a second model prompt"); assert.equal(calls.length, 2, "the invalid call triggered a second model prompt");
@@ -330,36 +306,34 @@ describe("chat store", () => {
describe("mid-turn navigation (turn liveness + write-ahead)", () => { describe("mid-turn navigation (turn liveness + write-ahead)", () => {
it("one turn per chat: a concurrent request is 409'd; the next one after finish is accepted", async () => { it("one turn per chat: a concurrent request is 409'd; the next one after finish is accepted", async () => {
const calls: string[][] = []; const { engine } = scriptedTeacher(defaultLesson);
const agent = createTeacherAgent(undefined, scriptedModel(calls));
const chat = chatStore.createChat("Test: turn guard"); const chat = chatStore.createChat("Test: turn guard");
createdChats.push(chat.id); createdChats.push(chat.id);
const firstUser = { id: "u1", role: "user", parts: [{ type: "text", text: "teach me the solar system" }] }; const firstUser = { id: "u1", role: "user", parts: [{ type: "text", text: "teach me the solar system" }] };
const res1 = await handleChatRequest(post(chat.id, [firstUser]), agent); // NOT drained → in flight const res1 = await handleChatRequest(post(chat.id, [firstUser]), agentFactory(engine)); // NOT drained → in flight
assert.equal(isTurnRunning(chat.id), true, "turn registered as in flight"); assert.equal(isTurnRunning(chat.id), true, "turn registered as in flight");
// A stale client (remount racing a live turn) tries to start a second turn. // A stale client (remount racing a live turn) tries to start a second turn.
const res2 = await handleChatRequest(post(chat.id, [firstUser]), agent); const res2 = await handleChatRequest(post(chat.id, [firstUser]), agentFactory(engine));
assert.equal(res2.status, 409, "concurrent turn rejected — last-writer-wins persist would clobber it"); assert.equal(res2.status, 409, "concurrent turn rejected — last-writer-wins persist would clobber it");
await drain(res1); // finish → onFinish → endTurn await drain(res1); // finish → onFinish → endTurn
assert.equal(isTurnRunning(chat.id), false, "registry cleared at finish"); assert.equal(isTurnRunning(chat.id), false, "registry cleared at finish");
const res3 = await handleChatRequest(post(chat.id, [firstUser]), agent); const res3 = await handleChatRequest(post(chat.id, [firstUser]), agentFactory(engine));
assert.equal(res3.status, 200, "the next turn is accepted after the previous finished"); assert.equal(res3.status, 200, "the next turn is accepted after the previous finished");
await drain(res3); await drain(res3);
}); });
it("write-ahead: the learner's answer is durable BEFORE the turn finishes", async () => { it("write-ahead: the learner's answer is durable BEFORE the turn finishes", async () => {
const calls: string[][] = []; const { engine } = scriptedTeacher(defaultLesson);
const agent = createTeacherAgent(undefined, scriptedModel(calls));
const chat = chatStore.createChat("Test: write-ahead"); const chat = chatStore.createChat("Test: write-ahead");
createdChats.push(chat.id); createdChats.push(chat.id);
// Turn 1: the teacher asks the quiz (pending part in the record). // Turn 1: the teacher asks the quiz (pending part in the record).
const firstUser = { id: "u1", role: "user", parts: [{ type: "text", text: "teach me the solar system" }] }; const firstUser = { id: "u1", role: "user", parts: [{ type: "text", text: "teach me the solar system" }] };
await drain(await handleChatRequest(post(chat.id, [firstUser]), agent)); await drain(await handleChatRequest(post(chat.id, [firstUser]), agentFactory(engine)));
const saved1 = chatStore.getChat(chat.id)!; const saved1 = chatStore.getChat(chat.id)!;
const assistantMsg = saved1.messages[1]!; const assistantMsg = saved1.messages[1]!;
@@ -390,11 +364,11 @@ describe("mid-turn navigation (turn liveness + write-ahead)", () => {
// NOT drained: the teacher turn is streaming, but the answer is already // NOT drained: the teacher turn is streaming, but the answer is already
// on disk — a remount (or crash) mid-turn sees the answered quiz, never // on disk — a remount (or crash) mid-turn sees the answered quiz, never
// the stale pre-answer question. // the stale pre-answer question.
const res = await handleChatRequest(post(chat.id, messages2), agent); const res = await handleChatRequest(post(chat.id, messages2), agentFactory(engine));
const midTurn = chatStore.getChat(chat.id)!; const midTurn = chatStore.getChat(chat.id)!;
const midQuiz = midTurn.messages[1]!.parts.find((p) => p.type === "tool-quiz"); const midQuiz = midTurn.messages[1]!.parts.find((p) => p.type === "tool-quiz");
assert.ok(midQuiz && midQuiz.type === "tool-quiz"); assert.ok(midQuiz && midQuiz.type === "tool-quiz");
assert.equal(midQuiz.state, "output-available", "write-ahead persisted the answer before finish"); assert.equal((midQuiz as { state?: string }).state, "output-available", "write-ahead persisted the answer before finish");
await drain(res); // finish → persist the full turn → endTurn await drain(res); // finish → persist the full turn → endTurn
const saved2 = chatStore.getChat(chat.id)!; const saved2 = chatStore.getChat(chat.id)!;
@@ -405,17 +379,17 @@ describe("mid-turn navigation (turn liveness + write-ahead)", () => {
it("a turn that fails to start (malformed messages) clears the registry — no lockout", async () => { it("a turn that fails to start (malformed messages) clears the registry — no lockout", async () => {
const chat = chatStore.createChat("Test: turn failed to start"); const chat = chatStore.createChat("Test: turn failed to start");
createdChats.push(chat.id); createdChats.push(chat.id);
// A text part with no text: fails UIMessage validation before the stream // A text part with no text: fails message validation before the stream
// exists — the path where no onFinish will ever run. // exists — the path where no onFinish will ever run.
const bad = [{ id: "m1", role: "user", parts: [{ type: "text" }] }]; const bad = [{ id: "m1", role: "user", parts: [{ type: "text" }] }];
const res = await handleChatRequest(post(chat.id, bad), createTeacherAgent(undefined, scriptedModel([]))); const res = await handleChatRequest(post(chat.id, bad), agentFactory(scriptedTeacher(defaultLesson).engine));
assert.equal(res.status, 500, "turn-start failure is a 500, not a hang"); assert.equal(res.status, 500, "turn-start failure is a 500, not a hang");
assert.equal(isTurnRunning(chat.id), false, "registry cleared when the stream never started"); assert.equal(isTurnRunning(chat.id), false, "registry cleared when the stream never started");
// And the chat is NOT locked: the next (valid) turn is accepted. // And the chat is NOT locked: the next (valid) turn is accepted.
const ok = await handleChatRequest( const ok = await handleChatRequest(
post(chat.id, [{ id: "m1", role: "user", parts: [{ type: "text", text: "teach me" }] }]), post(chat.id, [{ id: "m1", role: "user", parts: [{ type: "text", text: "teach me" }] }]),
createTeacherAgent(undefined, scriptedModel([])), agentFactory(scriptedTeacher(defaultLesson).engine),
); );
assert.equal(ok.status, 200); assert.equal(ok.status, 200);
await drain(ok); await drain(ok);
+40 -73
View File
@@ -27,8 +27,7 @@ import assert from "node:assert/strict";
import * as fs from "node:fs"; import * as fs from "node:fs";
import * as path from "node:path"; import * as path from "node:path";
import { after, describe, it } from "node:test"; import { after, describe, it } from "node:test";
import { MockLanguageModelV4, simulateReadableStream } from "ai/test"; import { createFauxCore, fauxAssistantMessage, fauxToolCall, type FauxResponseStep } from "@earendil-works/pi-ai";
import type { LanguageModelV4StreamPart } from "@ai-sdk/provider";
import { renderToStaticMarkup } from "react-dom/server"; import { renderToStaticMarkup } from "react-dom/server";
import { PATHS } from "../config.ts"; import { PATHS } from "../config.ts";
import { import {
@@ -54,7 +53,10 @@ import { createTeacherTools } from "../src/agent/tools.ts";
import { chatStore } from "../src/store/chats.ts"; import { chatStore } from "../src/store/chats.ts";
import { jobStore } from "../src/store/jobs.ts"; import { jobStore } from "../src/store/jobs.ts";
import { handleChatRequest } from "../src/server/chat-handler.ts"; import { handleChatRequest } from "../src/server/chat-handler.ts";
import { createTeacherAgent } from "../src/agent/teacher-agent.ts"; import { createTeacherAgent, type TeachingAgent } from "../src/agent/teacher-agent.ts";
import { recordToPiMessages } from "../src/server/pi/record-adapter.ts";
import type { PiEngine } from "../src/agent/models.ts";
import type { UIMessage } from "ai";
/** /**
* A fake deliver seam + a runner wired like the chat-handler's, per test: * A fake deliver seam + a runner wired like the chat-handler's, per test:
@@ -104,37 +106,22 @@ after(() => {
visualModelOverride.current = undefined; visualModelOverride.current = undefined;
}); });
function usage() { /** A scripted pi engine (faux provider): the response queue replaces the old
return { * doGenerate turn counter — one queued response per model call (tool-call
inputTokens: { total: 10, noCache: 10, cacheRead: 0, cacheWrite: 0 }, * turn or final text), consumed in order. */
outputTokens: { total: 5, text: 5, reasoning: 0 }, function scriptedModel(turns: Array<{ toolCalls?: Array<{ toolCallId: string; toolName: string; input: unknown }>; text?: string }>): PiEngine {
}; const steps: FauxResponseStep[] = turns.map((step) => {
}
/** A model that answers a scripted sequence of doGenerate turns (tool calls or final text). */
function scriptedModel(turns: Array<{ toolCalls?: Array<{ toolCallId: string; toolName: string; input: unknown }>; text?: string }>): MockLanguageModelV4 {
let turn = -1;
return new MockLanguageModelV4({
doGenerate: async () => {
turn++;
const step = turns[turn];
if (!step) throw new Error(`scripted model exhausted after ${turns.length} turns`);
if (step.toolCalls) { if (step.toolCalls) {
return { return fauxAssistantMessage(
content: step.toolCalls.map((c) => ({ type: "tool-call" as const, toolCallId: c.toolCallId, toolName: c.toolName, input: JSON.stringify(c.input) })), step.toolCalls.map((c) => fauxToolCall(c.toolName, c.input as Record<string, unknown>, { id: c.toolCallId })),
finishReason: { unified: "tool-calls" as const, raw: "tool-calls" }, { stopReason: "toolUse" },
usage: usage(), );
warnings: [],
};
} }
return { return fauxAssistantMessage(step.text ?? "");
content: [{ type: "text" as const, text: step.text ?? "" }],
finishReason: { unified: "stop" as const, raw: "stop" },
usage: usage(),
warnings: [],
};
},
}); });
const faux = createFauxCore({ tokenSize: { min: 64, max: 64 } });
faux.setResponses(steps);
return { model: faux.getModel(), streamFn: faux.streamSimple };
} }
// ─── fake headless ─────────────────────────────────────────────────────────── // ─── fake headless ───────────────────────────────────────────────────────────
@@ -424,7 +411,7 @@ describe("demo builder (fake model + fake headless)", () => {
const result = await runDemoBuilder( const result = await runDemoBuilder(
{ title: "Damped oscillator", brief: "A mass on a spring with a damping slider." }, { title: "Damped oscillator", brief: "A mass on a spring with a damping slider." },
{ model, headless, onActivity: (line) => activities.push(line) }, { engine: model, headless, onActivity: (line) => activities.push(line) },
); );
assert.equal(result.title, "Damped oscillator"); assert.equal(result.title, "Damped oscillator");
@@ -442,7 +429,7 @@ describe("demo builder (fake model + fake headless)", () => {
{ toolCalls: [{ toolCallId: "w1", toolName: "write_demo", input: { html: GOOD_HTML } }] }, { toolCalls: [{ toolCallId: "w1", toolName: "write_demo", input: { html: GOOD_HTML } }] },
{ text: "Done." }, { text: "Done." },
]); ]);
const result = await runDemoBuilder({ title: "T", brief: "b" }, { model, headless }); const result = await runDemoBuilder({ title: "T", brief: "b" }, { engine: model, headless });
assert.equal(result.html, GOOD_HTML); assert.equal(result.html, GOOD_HTML);
assert.equal(result.tested, true, "tested by the final verification gate"); assert.equal(result.tested, true, "tested by the final verification gate");
assert.equal(headless.demoRuns.length, 1, "the gate ran the demo once itself"); assert.equal(headless.demoRuns.length, 1, "the gate ran the demo once itself");
@@ -454,7 +441,7 @@ describe("demo builder (fake model + fake headless)", () => {
{ toolCalls: [{ toolCallId: "w1", toolName: "write_demo", input: { html: BROKEN_HTML } }] }, { toolCalls: [{ toolCallId: "w1", toolName: "write_demo", input: { html: BROKEN_HTML } }] },
{ text: "It works, trust me." }, { text: "It works, trust me." },
]); ]);
const result = await runDemoBuilder({ title: "T", brief: "b" }, { model, headless }); const result = await runDemoBuilder({ title: "T", brief: "b" }, { engine: model, headless });
assert.equal(result.html, ""); assert.equal(result.html, "");
assert.equal(result.tested, true, "the loop ran"); assert.equal(result.tested, true, "the loop ran");
assert.match(result.notes ?? "", /^FAILED:/); assert.match(result.notes ?? "", /^FAILED:/);
@@ -463,14 +450,14 @@ describe("demo builder (fake model + fake headless)", () => {
it("fails fast with a clear message when no Chrome is available (never skips testing)", async () => { it("fails fast with a clear message when no Chrome is available (never skips testing)", async () => {
const model = scriptedModel([{ text: "should never be reached" }]); const model = scriptedModel([{ text: "should never be reached" }]);
const result = await runDemoBuilder({ title: "T", brief: "b" }, { model, headless: noChromeHeadless() }); const result = await runDemoBuilder({ title: "T", brief: "b" }, { engine: model, headless: noChromeHeadless() });
assert.equal(result.html, ""); assert.equal(result.html, "");
assert.match(result.notes ?? "", /^FAILED: .*Chrome/); assert.match(result.notes ?? "", /^FAILED: .*Chrome/);
}); });
it("fails when the model never writes a demo", async () => { it("fails when the model never writes a demo", async () => {
const model = scriptedModel([{ text: "I cannot build this." }]); const model = scriptedModel([{ text: "I cannot build this." }]);
const result = await runDemoBuilder({ title: "T", brief: "b" }, { model, headless: fakeHeadless() }); const result = await runDemoBuilder({ title: "T", brief: "b" }, { engine: model, headless: fakeHeadless() });
assert.equal(result.html, ""); assert.equal(result.html, "");
assert.match(result.notes ?? "", /^FAILED: the builder never wrote a demo/); assert.match(result.notes ?? "", /^FAILED: the builder never wrote a demo/);
}); });
@@ -489,7 +476,7 @@ describe("visual maker (fake model + fake headless)", () => {
{ text: "A number line from -3 to 3 with the root marked at 0." }, { text: "A number line from -3 to 3 with the root marked at 0." },
]); ]);
const activities: string[] = []; const activities: string[] = [];
const result = await runVisualMaker({ brief: "Number line, root at 0." }, { model, headless, onActivity: (l) => activities.push(l) }); const result = await runVisualMaker({ brief: "Number line, root at 0." }, { engine: model, headless, onActivity: (l) => activities.push(l) });
assert.equal(result.svg, SVG_1); assert.equal(result.svg, SVG_1);
assert.equal(result.notes, "A number line from -3 to 3 with the root marked at 0."); assert.equal(result.notes, "A number line from -3 to 3 with the root marked at 0.");
@@ -504,7 +491,7 @@ describe("visual maker (fake model + fake headless)", () => {
{ text: "Done." }, // premature finish — no render yet { text: "Done." }, // premature finish — no render yet
{ text: "Looked at it — correct." }, // after the forced look { text: "Looked at it — correct." }, // after the forced look
]); ]);
const result = await runVisualMaker({ brief: "b" }, { model, headless }); const result = await runVisualMaker({ brief: "b" }, { engine: model, headless });
assert.equal(result.svg, SVG_1, "delivered after the forced look"); assert.equal(result.svg, SVG_1, "delivered after the forced look");
assert.equal(result.notes, "Looked at it — correct."); assert.equal(result.notes, "Looked at it — correct.");
assert.equal(headless.svgRenders.length, 1, "the gate rendered the svg once itself"); assert.equal(headless.svgRenders.length, 1, "the gate rendered the svg once itself");
@@ -522,7 +509,7 @@ describe("visual maker (fake model + fake headless)", () => {
{ toolCalls: [{ toolCallId: "r2", toolName: "render_svg", input: {} }] }, { toolCalls: [{ toolCallId: "r2", toolName: "render_svg", input: {} }] },
{ text: "Moved the start of the line." }, { text: "Moved the start of the line." },
]); ]);
const result = await runVisualMaker({ brief: "b" }, { model, headless }); const result = await runVisualMaker({ brief: "b" }, { engine: model, headless });
assert.equal(result.svg, SVG_1.replace('<line x1="10"', '<line x1="20"'), "the edit was applied to the delivered svg"); assert.equal(result.svg, SVG_1.replace('<line x1="10"', '<line x1="20"'), "the edit was applied to the delivered svg");
assert.equal(result.notes, "Moved the start of the line."); assert.equal(result.notes, "Moved the start of the line.");
@@ -536,7 +523,7 @@ describe("visual maker (fake model + fake headless)", () => {
{ toolCalls: [{ toolCallId: "e0", toolName: "edit_svg", input: { old_text: "svg", new_text: "x" } }] }, { toolCalls: [{ toolCallId: "e0", toolName: "edit_svg", input: { old_text: "svg", new_text: "x" } }] },
{ text: "Cannot proceed." }, { text: "Cannot proceed." },
]); ]);
const result = await runVisualMaker({ brief: "b" }, { model, headless }); const result = await runVisualMaker({ brief: "b" }, { engine: model, headless });
assert.equal(result.svg, "", "the failed edit produced no deliverable"); assert.equal(result.svg, "", "the failed edit produced no deliverable");
assert.match(result.notes ?? "", /never wrote an SVG/); assert.match(result.notes ?? "", /never wrote an SVG/);
}); });
@@ -550,20 +537,20 @@ describe("visual maker (fake model + fake headless)", () => {
{ toolCalls: [{ toolCallId: "r2", toolName: "render_svg", input: {} }] }, { toolCalls: [{ toolCallId: "r2", toolName: "render_svg", input: {} }] },
{ text: "Fixed and verified." }, { text: "Fixed and verified." },
]); ]);
const result = await runVisualMaker({ brief: "b" }, { model, headless }); const result = await runVisualMaker({ brief: "b" }, { engine: model, headless });
assert.equal(result.svg, SVG_1); assert.equal(result.svg, SVG_1);
assert.equal(result.notes, "Fixed and verified."); assert.equal(result.notes, "Fixed and verified.");
}); });
it("fails when the model never writes an svg", async () => { it("fails when the model never writes an svg", async () => {
const model = scriptedModel([{ text: "This idea is purely relational — belongs in mermaid." }]); const model = scriptedModel([{ text: "This idea is purely relational — belongs in mermaid." }]);
const result = await runVisualMaker({ brief: "b" }, { model, headless: fakeHeadless() }); const result = await runVisualMaker({ brief: "b" }, { engine: model, headless: fakeHeadless() });
assert.equal(result.svg, ""); assert.equal(result.svg, "");
assert.match(result.notes ?? "", /^FAILED: the maker never wrote an SVG/); assert.match(result.notes ?? "", /^FAILED: the maker never wrote an SVG/);
}); });
it("fails with a clear message when no Chrome is available", async () => { it("fails with a clear message when no Chrome is available", async () => {
const result = await runVisualMaker({ brief: "b" }, { model: scriptedModel([{ text: "x" }]), headless: noChromeHeadless() }); const result = await runVisualMaker({ brief: "b" }, { engine: scriptedModel([{ text: "x" }]), headless: noChromeHeadless() });
assert.equal(result.svg, ""); assert.equal(result.svg, "");
assert.match(result.notes ?? "", /^FAILED: .*Chrome/); assert.match(result.notes ?? "", /^FAILED: .*Chrome/);
}); });
@@ -685,30 +672,10 @@ describe("demo/visual job runners", () => {
createdChats.push(chat.id); createdChats.push(chat.id);
// A teacher that calls demo on turn 1, then continues in the SAME turn. // A teacher that calls demo on turn 1, then continues in the SAME turn.
let turn = 0; const teacher = scriptedModel([
const teacher = new MockLanguageModelV4({ { toolCalls: [{ toolCallId: "call-d2", toolName: "demo", input: DEMO_INPUT }] },
doStream: async () => { { text: "Here it is." },
turn++; ]);
const chunks: LanguageModelV4StreamPart[] =
turn === 1
? [
{ type: "stream-start", warnings: [] },
{ type: "tool-input-start", id: "call-d2", toolName: "demo" },
{ type: "tool-input-delta", id: "call-d2", delta: JSON.stringify(DEMO_INPUT) },
{ type: "tool-input-end", id: "call-d2" },
{ type: "tool-call", toolCallId: "call-d2", toolName: "demo", input: JSON.stringify(DEMO_INPUT) },
{ type: "finish", finishReason: { unified: "tool-calls", raw: "tool_calls" }, usage: usage() },
]
: [
{ type: "stream-start", warnings: [] },
{ type: "text-start", id: "t2" },
{ type: "text-delta", id: "t2", delta: "Here it is." },
{ type: "text-end", id: "t2" },
{ type: "finish", finishReason: { unified: "stop", raw: "stop" }, usage: usage() },
];
return { stream: simulateReadableStream({ chunks, initialDelayInMs: 0, chunkDelayInMs: 0 }) };
},
});
// The overrides must be live BEFORE the request — dispatch fires the // The overrides must be live BEFORE the request — dispatch fires the
// runner mid-stream, inside the tool execute. // runner mid-stream, inside the tool execute.
@@ -731,17 +698,17 @@ describe("demo/visual job runners", () => {
messages: [{ id: "u1", role: "user", parts: [{ type: "text", text: "teach me compounding" }] }], messages: [{ id: "u1", role: "user", parts: [{ type: "text", text: "teach me compounding" }] }],
}), }),
}), }),
createTeacherAgent( (outgoing: UIMessage[]): TeachingAgent =>
undefined, createTeacherAgent(undefined, {
teacher, engine: teacher,
undefined, messages: recordToPiMessages(outgoing),
createTeacherTools({ tools: createTeacherTools({
chatId: chat.id, chatId: chat.id,
chatStore, chatStore,
jobStore, jobStore,
spawn: (job, input) => void runner.runJob(job.id, input), spawn: (job, input) => void runner.runJob(job.id, input),
}), }),
), }),
), ),
); );
+64 -87
View File
@@ -5,7 +5,7 @@
* *
* transcriptView — the compact view both LLM stages see (pure): * transcriptView — the compact view both LLM stages see (pure):
* message-id-tagged verbatim text + quiz lines, NO tree * message-id-tagged verbatim text + quiz lines, NO tree
* pipeline — MockLanguageModelV4 drives both agents: findings out, * pipeline — scripted faux engines drive both agents: findings out,
* audit verdict (with linkTo identity) in, deterministic * audit verdict (with linkTo identity) in, deterministic
* merge applied; failure paths checkpoint nothing; * merge applied; failure paths checkpoint nothing;
* single-flight per chat * single-flight per chat
@@ -22,8 +22,9 @@
import assert from "node:assert/strict"; import assert from "node:assert/strict";
import * as fs from "node:fs"; import * as fs from "node:fs";
import { after, describe, it } from "node:test"; import { after, describe, it } from "node:test";
import { MockLanguageModelV4 } from "ai/test"; import { createFauxCore, fauxAssistantMessage, fauxToolCall, type FauxResponseStep } from "@earendil-works/pi-ai";
import type { UIMessage } from "ai"; import type { UIMessage } from "ai";
import type { PiEngine } from "../src/agent/models.ts";
import type { import type {
AuditResult, AuditResult,
GraphEdge, GraphEdge,
@@ -42,11 +43,14 @@ import { PATHS } from "../config.ts";
const createdChats: string[] = []; const createdChats: string[] = [];
function usage() { /** A scripted pi engine (faux provider, docs/pi-engine-swap.md §7): the
return { * response queue replaces the old doGenerate turn counter. Usage is
inputTokens: { total: 10, noCache: 10, cacheRead: 0, cacheWrite: 0 }, * ESTIMATED by faux from the prompt/response text (deterministic, > 0) —
outputTokens: { total: 5, text: 5, reasoning: 0 }, * the tests assert structure and positivity, not flat mock numbers. */
}; function fauxEngine(steps: FauxResponseStep[]): PiEngine {
const faux = createFauxCore({});
faux.setResponses(steps);
return { model: faux.getModel(), streamFn: faux.streamSimple };
} }
/** An in-memory GraphStore (core/ports.ts) — this suite never touches data/graph. */ /** An in-memory GraphStore (core/ports.ts) — this suite never touches data/graph. */
@@ -163,68 +167,46 @@ const FINDINGS: SessionFindings = {
attribution: { quiz: { "probe-1": "p1", "check-1": "p2" } }, attribution: { quiz: { "probe-1": "p1", "check-1": "p2" } },
}; };
/** A model whose ONE tool call is `name` with the given input; then text. /** An engine whose ONE response is the `name` tool call with the given
* `onCall` (optional) fires at the moment of the first generate — the * input — the whole deliverable; the pipeline stops after a clean submit
* observation point for mid-run assertions. */ * (shouldStop), so no trailing text step exists anymore. `onCall` fires at
function submitModel(name: string, input: unknown, onCall?: () => void): MockLanguageModelV4 { * the moment of the model call — the observation point for mid-run
let turn = 0; * assertions. */
return new MockLanguageModelV4({ function submitModel(name: string, input: unknown, onCall?: () => void): PiEngine {
doGenerate: async () => { return fauxEngine([
turn++; (context) => {
if (turn === 1) {
onCall?.(); onCall?.();
return { void context;
content: [ return fauxAssistantMessage(
{ [fauxToolCall(name, input as Record<string, unknown>, { id: `${name}-call` })],
type: "tool-call" as const, { stopReason: "toolUse" },
toolCallId: `${name}-call`, );
toolName: name,
input: JSON.stringify(input),
}, },
], ]);
finishReason: { unified: "tool-calls" as const, raw: "tool-calls" }, }
usage: usage(),
warnings: [], /** An engine whose model call never resolves until its abort signal fires
}; * — mimicking a cancelled provider request (the abort path). Like a real
* provider, an ALREADY-aborted signal resolves immediately (the abort can
* win the race before the call attaches its listener); the faux stream
* then sees the aborted signal and ends the message as aborted. */
function hangingModel(): PiEngine {
return fauxEngine([
(_context, options) =>
new Promise((resolve) => {
const sig = options?.signal;
if (sig?.aborted) {
resolve(fauxAssistantMessage("late"));
return;
} }
return { sig?.addEventListener("abort", () => resolve(fauxAssistantMessage("late")), { once: true });
content: [{ type: "text" as const, text: "submitted" }],
finishReason: { unified: "stop" as const, raw: "stop" },
usage: usage(),
warnings: [],
};
},
});
}
/** A model that hangs until its abort signal fires, then rejects —
* mimicking a cancelled provider request (the abort path). Like a real
* provider, an ALREADY-aborted signal rejects immediately (the abort
* can win the race before doGenerate attaches its listener). */
function hangingModel(): MockLanguageModelV4 {
const abortReject = (reject: (reason: Error) => void) =>
reject(Object.assign(new Error("The operation was aborted"), { name: "AbortError" }));
return new MockLanguageModelV4({
doGenerate: async (options): Promise<never> => {
await new Promise<never>((_, reject) => {
if (options.abortSignal?.aborted) abortReject(reject);
else options.abortSignal?.addEventListener("abort", () => abortReject(reject));
});
throw new Error("unreachable");
},
});
}
/** A model that never calls a tool (the failure path). */
function silentModel(): MockLanguageModelV4 {
return new MockLanguageModelV4({
doGenerate: async () => ({
content: [{ type: "text" as const, text: "I refuse to use tools" }],
finishReason: { unified: "stop" as const, raw: "stop" },
usage: usage(),
warnings: [],
}), }),
}); ]);
}
/** An engine that never calls a tool (the failure path). */
function silentModel(): PiEngine {
return fauxEngine([fauxAssistantMessage("I refuse to use tools"), fauxAssistantMessage("…")]);
} }
// ─── transcriptView (pure) ─────────────────────────────────────────────────── // ─── transcriptView (pure) ───────────────────────────────────────────────────
@@ -427,9 +409,11 @@ describe("the extraction pipeline (extractor → auditor → merger)", () => {
}); });
await pipeline.runExtraction(chatId); await pipeline.runExtraction(chatId);
// Live: mid-audit the extractor's usage was fully accumulated // Live: mid-audit the extractor's usage was fully accumulated (faux
// (the mock bills 10 input / 5 output per step, flat per the SDK). // estimates tokens from the prompt/response text — deterministic,
assert.deepEqual(usageSeen, [10]); // positive, no cache).
assert.equal(usageSeen.length, 1);
assert.ok(typeof usageSeen[0] === "number" && usageSeen[0]! > 0, "the extractor's step billed input tokens");
// Persisted: the run appended ONE usage entry to the chat's // Persisted: the run appended ONE usage entry to the chat's
// extractions log — per-stage model + steps, with merge stats. // extractions log — per-stage model + steps, with merge stats.
@@ -438,16 +422,14 @@ describe("the extraction pipeline (extractor → auditor → merger)", () => {
const run = record.extractions![0]!; const run = record.extractions![0]!;
assert.equal(run.stats?.nodes, 0); // linked, not created assert.equal(run.stats?.nodes, 0); // linked, not created
assert.equal(run.stats?.events, 3); assert.equal(run.stats?.events, 3);
assert.equal(run.stages.extractor.model, "mock-model-id"); // the mock's model id assert.equal(run.stages.extractor.model, "faux-1"); // the scripted engine's model id
assert.equal(run.stages.extractor.steps.length, 2); // tool-call step + text step assert.equal(run.stages.extractor.steps.length, 1); // the submit call — the run STOPS after a clean submit
assert.deepEqual(run.stages.extractor.steps[0], { const step = run.stages.extractor.steps[0]!;
inputTokens: 10, assert.ok(step.inputTokens! > 0 && step.outputTokens! > 0, "estimated usage captured");
outputTokens: 5, assert.equal(step.cacheReadTokens, 0);
cacheReadTokens: 0, assert.equal(step.cacheWriteTokens, 0);
cacheWriteTokens: 0, assert.equal(step.reasoningTokens, 0);
reasoningTokens: 0, assert.equal(run.stages.auditor.steps.length, 1);
});
assert.equal(run.stages.auditor.steps.length, 2);
assert.ok(!Number.isNaN(Date.parse(run.at))); assert.ok(!Number.isNaN(Date.parse(run.at)));
}); });
@@ -497,18 +479,13 @@ describe("the extraction pipeline (extractor → auditor → merger)", () => {
release = resolve; release = resolve;
}); });
let calls = 0; let calls = 0;
const gatedExtractor = new MockLanguageModelV4({ const gatedExtractor = fauxEngine([
doGenerate: async () => { async () => {
calls++; calls++;
if (calls === 1) await gate; if (calls === 1) await gate;
return { return fauxAssistantMessage("…");
content: [{ type: "text" as const, text: "…" }],
finishReason: { unified: "stop" as const, raw: "stop" },
usage: usage(),
warnings: [],
};
}, },
}); ]);
const pipeline = createExtractionPipeline({ const pipeline = createExtractionPipeline({
chatStore, chatStore,
+76 -135
View File
@@ -28,8 +28,7 @@ import * as fs from "node:fs";
import * as path from "node:path"; import * as path from "node:path";
import { after, describe, it } from "node:test"; import { after, describe, it } from "node:test";
import type { UIMessage } from "ai"; import type { UIMessage } from "ai";
import { MockLanguageModelV4, simulateReadableStream } from "ai/test"; import { createFauxCore, fauxAssistantMessage, fauxText, fauxToolCall, type AssistantMessage, type FauxResponseStep } from "@earendil-works/pi-ai";
import type { LanguageModelV4StreamPart } from "@ai-sdk/provider";
import { import {
accumulateJobUsage, accumulateJobUsage,
buildAbortSteer, buildAbortSteer,
@@ -45,16 +44,18 @@ import {
buildProbeHandoffSteer, buildProbeHandoffSteer,
type ProbeHandoffSteerMessage, type ProbeHandoffSteerMessage,
} from "../src/core/contracts.ts"; } from "../src/core/contracts.ts";
import { composeInstructions, createTeacherAgent } from "../src/agent/teacher-agent.ts"; import { composeInstructions, createTeacherAgent, type TeachingAgent } from "../src/agent/teacher-agent.ts";
import { createTeacherTools } from "../src/agent/tools.ts"; import { createTeacherTools } from "../src/agent/tools.ts";
import { webTools } from "../src/agent/tools/web.ts";
import { handleChatRequest } from "../src/server/chat-handler.ts"; import { handleChatRequest } from "../src/server/chat-handler.ts";
import { recordToPiMessages } from "../src/server/pi/record-adapter.ts";
import { abortJob, claimSteers, deliverJobResult, setTurnRunner } from "../src/server/jobs.ts"; import { abortJob, claimSteers, deliverJobResult, setTurnRunner } from "../src/server/jobs.ts";
import { beginTurn, endTurn } from "../src/server/turn-liveness.ts"; import { beginTurn, endTurn } from "../src/server/turn-liveness.ts";
import { persistTurn, withMessageIds } from "../src/server/turn-service.ts"; import { persistTurn, withMessageIds } from "../src/server/turn-service.ts";
import { chatStore } from "../src/store/chats.ts"; import { chatStore } from "../src/store/chats.ts";
import { jobStore } from "../src/store/jobs.ts"; import { jobStore } from "../src/store/jobs.ts";
import { createJobRunner, researchModelOverride } from "../src/agent/subagents/runner.ts"; import { createJobRunner, researchModelOverride } from "../src/agent/subagents/runner.ts";
import { createResearcherAgent } from "../src/agent/subagents/researcher.ts"; import type { PiEngine } from "../src/agent/models.ts";
import { PATHS } from "../config.ts"; import { PATHS } from "../config.ts";
const JOBS_DIR = path.join(PATHS.data, "jobs"); const JOBS_DIR = path.join(PATHS.data, "jobs");
@@ -88,11 +89,30 @@ after(() => {
researchModelOverride.current = undefined; researchModelOverride.current = undefined;
}); });
function usage() { /** A scripted pi engine (faux provider): the response queue replaces the
return { * old doStream/doGenerate turn counters. tokenSize is set so every test
inputTokens: { total: 10, noCache: 10, cacheRead: 0, cacheWrite: 0 }, * string streams as ONE delta — the SSE-body assertions ("Back with the
outputTokens: { total: 5, text: 5, reasoning: 0 }, * brief.") read contiguous text. */
}; function fauxEngine(steps: FauxResponseStep[]): PiEngine {
const faux = createFauxCore({ tokenSize: { min: 64, max: 64 } });
faux.setResponses(steps);
return { model: faux.getModel(), streamFn: faux.streamSimple };
}
/** A scripted response factory whose model call never resolves until its
* abort signal fires (a cancelled provider request); an already-aborted
* signal resolves at once — the faux stream then ends the message aborted. */
function abortableFactory(): FauxResponseStep {
const factory = (_context: unknown, options: { signal?: AbortSignal } | undefined) =>
new Promise<AssistantMessage>((resolve) => {
const sig = options?.signal;
if (sig?.aborted) {
resolve(fauxAssistantMessage("late"));
return;
}
sig?.addEventListener("abort", () => resolve(fauxAssistantMessage("late")), { once: true });
});
return factory as unknown as FauxResponseStep;
} }
async function drain(res: Response): Promise<string> { async function drain(res: Response): Promise<string> {
@@ -295,34 +315,14 @@ const RESEARCH_INPUT = {
}; };
/** A teacher that calls research (text → call), then plain text thereafter. */ /** A teacher that calls research (text → call), then plain text thereafter. */
function researchCallingTeacher(): MockLanguageModelV4 { function researchCallingTeacher(): PiEngine {
let turn = 0; return fauxEngine([
return new MockLanguageModelV4({ fauxAssistantMessage(
doStream: async () => { [fauxText("Let me check the current state of the library."), fauxToolCall("research", RESEARCH_INPUT, { id: "call-r1" })],
turn++; { stopReason: "toolUse" },
const chunks: LanguageModelV4StreamPart[] = ),
turn === 1 fauxAssistantMessage([fauxText("Back with the brief.")]),
? [ ]);
{ type: "stream-start", warnings: [] },
{ type: "text-start", id: "t1" },
{ type: "text-delta", id: "t1", delta: "Let me check the current state of the library." },
{ type: "text-end", id: "t1" },
{ type: "tool-input-start", id: "call-r1", toolName: "research" },
{ type: "tool-input-delta", id: "call-r1", delta: JSON.stringify(RESEARCH_INPUT) },
{ type: "tool-input-end", id: "call-r1" },
{ type: "tool-call", toolCallId: "call-r1", toolName: "research", input: JSON.stringify(RESEARCH_INPUT) },
{ type: "finish", finishReason: { unified: "tool-calls", raw: "tool_calls" }, usage: usage() },
]
: [
{ type: "stream-start", warnings: [] },
{ type: "text-start", id: "t2" },
{ type: "text-delta", id: "t2", delta: "Back with the brief." },
{ type: "text-end", id: "t2" },
{ type: "finish", finishReason: { unified: "stop", raw: "stop" }, usage: usage() },
];
return { stream: simulateReadableStream({ chunks, initialDelayInMs: 0, chunkDelayInMs: 0 }) };
},
});
} }
describe("dispatch (async protocol §2)", () => { describe("dispatch (async protocol §2)", () => {
@@ -332,17 +332,15 @@ describe("dispatch (async protocol §2)", () => {
const delivery = fakeDeliver(); const delivery = fakeDeliver();
const testRunner = createJobRunner({ chatStore, jobStore, deliver: delivery.deliver }); const testRunner = createJobRunner({ chatStore, jobStore, deliver: delivery.deliver });
const agent = createTeacherAgent( const engine = researchCallingTeacher();
undefined, const tools = createTeacherTools({
researchCallingTeacher(),
undefined,
createTeacherTools({
chatId: chat.id, chatId: chat.id,
chatStore, chatStore,
jobStore, jobStore,
spawn: (job, input) => void testRunner.runJob(job.id, input), spawn: (job, input) => void testRunner.runJob(job.id, input),
}), });
); const agent = (outgoing: UIMessage[]): TeachingAgent =>
createTeacherAgent(undefined, { engine, tools, messages: recordToPiMessages(outgoing) });
const restoreFetch = stubFetch(); const restoreFetch = stubFetch();
researchModelOverride.current = researcherModel(); researchModelOverride.current = researcherModel();
@@ -394,38 +392,15 @@ sigma.js v3 is current.
## Sources ## Sources
- sigma.js docs (https://www.sigmajs.org/docs)`; - sigma.js docs (https://www.sigmajs.org/docs)`;
/** /** The researcher loop's scripted engine. Turn 1: a web_search tool call;
* The researcher loop runs via agent.generate (non-streaming) — the mock * turn 2: the final text brief. */
* therefore scripts doGenerate. Turn 1: a web_search tool call; turn 2: the function researcherModel(): PiEngine {
* final text brief. return fauxEngine([
*/ fauxAssistantMessage([fauxToolCall("web_search", { query: "sigma.js latest version" }, { id: "ws1" })], {
function researcherModel(): MockLanguageModelV4 { stopReason: "toolUse",
let turn = 0; }),
return new MockLanguageModelV4({ fauxAssistantMessage([fauxText(BRIEF)]),
doGenerate: async () => { ]);
turn++;
return turn === 1
? {
content: [
{
type: "tool-call" as const,
toolCallId: "ws1",
toolName: "web_search",
input: JSON.stringify({ query: "sigma.js latest version" }),
},
],
finishReason: { unified: "tool-calls" as const, raw: "tool-calls" },
usage: usage(),
warnings: [],
}
: {
content: [{ type: "text" as const, text: BRIEF }],
finishReason: { unified: "stop" as const, raw: "stop" },
usage: usage(),
warnings: [],
};
},
});
} }
function stubFetch() { function stubFetch() {
@@ -516,11 +491,11 @@ describe("research runner (async protocol §3)", () => {
const delivery = fakeDeliver(); const delivery = fakeDeliver();
const runner = createJobRunner({ chatStore, jobStore, deliver: delivery.deliver }); const runner = createJobRunner({ chatStore, jobStore, deliver: delivery.deliver });
researchModelOverride.current = new MockLanguageModelV4({ researchModelOverride.current = fauxEngine([
doGenerate: async () => { () => {
throw new Error("model exploded"); throw new Error("model exploded");
}, },
}); ]);
try { try {
await runner.runJob(job.id); await runner.runJob(job.id);
} finally { } finally {
@@ -578,11 +553,12 @@ describe("research runner (async protocol §3)", () => {
// ─── the researcher agent construction (model resolution seam) ────────────── // ─── the researcher agent construction (model resolution seam) ──────────────
describe("researcher agent", () => { describe("researcher agent", () => {
it("builds with the web tool set (no network touched)", () => { it("runs with the web tool set (no network touched)", () => {
const agent = createResearcherAgent(researcherModel()); // The researcher is a runSubagent call over the webTools registry —
assert.ok(agent.tools.web_search, "web_search available"); // the tool set is what the loop sees.
assert.ok(agent.tools.web_fetch, "web_fetch available"); assert.ok(webTools.web_search, "web_search available");
assert.equal(Object.keys(agent.tools).length, 2, "ONLY the web tools — no question/plan tools"); assert.ok(webTools.web_fetch, "web_fetch available");
assert.equal(Object.keys(webTools).length, 2, "ONLY the web tools — no question/plan tools");
// reset any accidental global override // reset any accidental global override
researchModelOverride.current = undefined; researchModelOverride.current = undefined;
}); });
@@ -950,57 +926,20 @@ describe("prompt composition", () => {
// ─── abort (2026-09 semantics: real kill, cancelled status, teacher told) ─── // ─── abort (2026-09 semantics: real kill, cancelled status, teacher told) ───
describe("abort (src/server/jobs.ts abortJob)", () => { describe("abort (src/server/jobs.ts abortJob)", () => {
/** A hanging model that only settles when its abort signal fires — the /** A hanging model whose call never resolves until its abort signal
* real-kill path: the in-flight call itself is killed. */ * fires — the real-kill path: the in-flight call itself is killed. */
function hangingModel(): MockLanguageModelV4 { function hangingModel(): PiEngine {
return new MockLanguageModelV4({ return fauxEngine([abortableFactory()]);
doGenerate: ({ abortSignal }) =>
new Promise((_resolve, reject) => {
// an ALREADY-aborted signal never fires the event — real
// providers reject an aborted fetch at once, the mock must too
if (abortSignal?.aborted) {
reject(new DOMException("The operation was aborted", "AbortError"));
return;
}
abortSignal?.addEventListener("abort", () => {
reject(new DOMException("The operation was aborted", "AbortError"));
});
}),
});
} }
/** One step (with usage) then a hang — proves usage-so-far survives the kill. */ /** One step (billed) then a hang — proves usage-so-far survives the kill. */
function oneStepThenHangModel(): MockLanguageModelV4 { function oneStepThenHangModel(): PiEngine {
let turn = 0; return fauxEngine([
return new MockLanguageModelV4({ fauxAssistantMessage([fauxToolCall("web_search", { query: "anything" }, { id: "ws1" })], {
doGenerate: ({ abortSignal }) => { stopReason: "toolUse",
turn++; }),
if (turn === 1) { abortableFactory(),
return Promise.resolve({ ]);
content: [
{
type: "tool-call" as const,
toolCallId: "ws1",
toolName: "web_search",
input: JSON.stringify({ query: "anything" }),
},
],
finishReason: { unified: "tool-calls" as const, raw: "tool-calls" },
usage: usage(),
warnings: [],
});
}
return new Promise((_resolve, reject) => {
if (abortSignal?.aborted) {
reject(new DOMException("The operation was aborted", "AbortError"));
return;
}
abortSignal?.addEventListener("abort", () => {
reject(new DOMException("The operation was aborted", "AbortError"));
});
});
},
});
} }
function newChatWithResearchJob(label: string, toolCallId: string): { chatId: string; jobId: string } { function newChatWithResearchJob(label: string, toolCallId: string): { chatId: string; jobId: string } {
@@ -1082,8 +1021,10 @@ describe("abort (src/server/jobs.ts abortJob)", () => {
assert.equal(job.error, "Aborted by learner"); assert.equal(job.error, "Aborted by learner");
assert.equal(job.usage!.steps, 1, "usage-so-far persisted by the runner's finally"); assert.equal(job.usage!.steps, 1, "usage-so-far persisted by the runner's finally");
assert.equal(job.usage!.model, "z-ai/glm-5.3-flash", "the config role's OpenRouter id"); assert.equal(job.usage!.model, "z-ai/glm-5.3-flash", "the config role's OpenRouter id");
assert.equal(job.usage!.inputTokens, 10); // faux estimates tokens from the prompt/response text — the
assert.equal(job.usage!.outputTokens, 5); // pinned fact is that the ONE completed step's spend survived.
assert.ok(job.usage!.inputTokens! > 0, "input tokens billed");
assert.ok(job.usage!.outputTokens! > 0, "output tokens billed");
} finally { } finally {
researchModelOverride.current = undefined; researchModelOverride.current = undefined;
} }
+2 -2
View File
@@ -50,10 +50,10 @@ const validQuiz = {
}; };
interface QuizToolSpec { interface QuizToolSpec {
inputSchema: { safeParse: (v: unknown) => { success: boolean; error?: { issues: Array<{ message: string }> } } }; schema: { safeParse: (v: unknown) => { success: boolean; error?: { issues: Array<{ message: string }> } } };
} }
/** teachQuizToolFor returns the single quiz tool; createProberTools a registry. */ /** teachQuizToolFor returns the single quiz tool; createProberTools a registry. */
const schemaOf = (toolLike: unknown) => (toolLike as QuizToolSpec).inputSchema; const schemaOf = (toolLike: unknown) => (toolLike as QuizToolSpec).schema;
const quizSchemaOf = (tools: Record<string, unknown>) => schemaOf(tools.quiz); const quizSchemaOf = (tools: Record<string, unknown>) => schemaOf(tools.quiz);
describe("teachQuizToolFor: the dynamic plan-nodeId validation (position signal)", () => { describe("teachQuizToolFor: the dynamic plan-nodeId validation (position signal)", () => {