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
+84 -51
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 { createJobRunner } from "../agent/subagents/runner.ts";
import { createExtractionPipeline } from "../agent/subagents/extraction.ts";
import { MODELS } from "../../config.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 { beginTurn, endTurn } from "./turn-liveness.ts";
import { persistTurn, triggerExtraction, withMessageIds } from "./turn-service.ts";
import type { CallUsage } from "../core/contracts.ts";
/**
* One chat turn: run the teacher agent over the client's message list and
* 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)
* 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.
*
* 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
* wake pipeline's server-side steers (src/server/jobs.ts) — one code path
* 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({
@@ -48,7 +56,27 @@ setTurnRunner(async (chatId, messages) => {
});
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 };
try {
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
// 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
// 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 steps: CallUsage[] = [];
const collectUsage = (step: { usage: LanguageModelUsage }) => {
const collectUsage = (usage: Usage) => {
steps.push({
inputTokens: step.usage.inputTokens,
outputTokens: step.usage.outputTokens,
cacheReadTokens: step.usage.inputTokenDetails?.cacheReadTokens,
cacheWriteTokens: step.usage.inputTokenDetails?.cacheWriteTokens,
reasoningTokens: step.usage.outputTokenDetails?.reasoningTokens,
inputTokens: usage.input,
outputTokens: usage.output,
cacheReadTokens: usage.cacheRead,
cacheWriteTokens: usage.cacheWrite,
// Optional in pi's Usage — persisted as 0, never dropped.
reasoningTokens: usage.reasoning ?? 0,
});
};
@@ -119,49 +148,53 @@ export async function handleChatRequest(req: Request, agent?: TeachingAgent): Pr
// steers, usage, extraction) — only the agent differs.
const teaching = outgoing.some((m) => m.parts.some((p) => p.type === "tool-probe_done"));
const modelId = teaching ? MODELS.teacher.model : MODELS.prober.model;
const constructed: TeachingAgent = teaching
? createTeacherAgent(
{
chatId,
chatStore: stores.chatStore,
jobStore: stores.jobStore,
spawn: (job, input) => void jobRunner.runJob(job.id, input),
},
undefined,
collectUsage,
)
: createProberAgent(
{
chatId,
chatStore: stores.chatStore,
deliver: (steer) => void deliverJobResult(steer),
},
undefined,
collectUsage,
);
// 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(
{
chatId,
chatStore: stores.chatStore,
jobStore: stores.jobStore,
spawn: (job, input) => void jobRunner.runJob(job.id, input),
},
{ messages: recordToPiMessages(outgoing), onUsage: collectUsage },
)
: createProberAgent(
{
chatId,
chatStore: stores.chatStore,
deliver: (steer) => void deliverJobResult(steer),
},
{ messages: recordToPiMessages(outgoing), onUsage: 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 {
return await createAgentUIStreamResponse({
agent: agent ?? constructed,
uiMessages: outgoing,
// Tool-call validation failures are messages OUR schemas authored (the
// quiz authoring contract, unknown tools) — safe and useful in the
// transcript: they are exactly what the teacher saw and retried on. The
// tool loop delivers them as plain strings (what the model received) or
// as the original error object; expose both, and keep everything else
// (provider/transport internals) masked, as the SDK defaults.
onError: (error) => {
if (typeof error === "string") return error;
if (InvalidToolInputError.isInstance(error) || NoSuchToolError.isInstance(error)) return error.message;
return "An error occurred.";
},
onFinish: ({ messages, isAborted }) => {
assertValidUiMessages(outgoing);
return createPiTurnResponse({
agent: constructed,
uiMessages: outgoing,
signal: req.signal,
// Stream-level failures are provider/transport internals — masked,
// like the SDK defaults. Tool-input validation messages never
// route through here: they ride the error-result path and reach the
// transcript directly (the retry context the model saw).
onError: (error) => {
console.error(`[chat ${chatId}] turn error:`, error);
return "An error occurred.";
},
onFinish: (messages, isAborted) => {
persistTurn(stores.chatStore, chatId, messages as UIMessage[], isAborted);
// usage log append (idempotent per turn id) — the stats bar
// refolds this on the client after the turn settles. messageId
// comes from the id-FILLED list: onFinish messages carry no ids
// (withMessageIds fills copies in persistTurn, not these); the
// re-fill here is idempotent, so the key is stable.
// comes from the id-FILLED list: the bridge's response message
// already carries its id; the re-fill here is idempotent.
const filled = withMessageIds(messages as UIMessage[]);
const response = [...filled].reverse().find((m) => m.role === "assistant");
if (response && steps.length > 0) {
@@ -186,4 +219,4 @@ export async function handleChatRequest(req: Request, agent?: TeachingAgent): Pr
console.error(`[chat ${chatId}] turn failed to start:`, err);
return Response.json({ error: "Failed to start the turn" }, { status: 500 });
}
}
}