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
+84
-51
@@ -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 });
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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[];
|
||||
}
|
||||
@@ -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",
|
||||
},
|
||||
});
|
||||
}
|
||||
Reference in New Issue
Block a user