feat(reedy): Appendix A · Phase 2.6 — AgentRuntime + abort helper (#4301)

This commit is contained in:
Huang Xin
2026-05-26 10:14:53 +08:00
committed by GitHub
parent a86b09dba5
commit df2698fa0e
3 changed files with 709 additions and 0 deletions
@@ -0,0 +1,362 @@
import { describe, it, expect, vi, beforeEach } from 'vitest';
import { z } from 'zod';
import type { ChatModel } from '@/services/reedy/models/ChatModel';
import { ToolRegistry } from '@/services/reedy/tools/ToolRegistry';
import type { ReedyTool } from '@/services/reedy/tools/types';
import { createPolicyLayer } from '@/services/reedy/context';
import type { ReedyEvent } from '@/services/reedy/runtime/events';
// Mock streamText so we drive what the runtime sees without an LLM.
// vi.hoisted lets the factory below access these — vi.mock is hoisted
// above regular consts, so a plain `const streamTextMock = vi.fn()` would
// fail with "cannot access before initialization".
const { streamTextMock, stepCountIsMock } = vi.hoisted(() => ({
streamTextMock: vi.fn(),
stepCountIsMock: vi.fn((n: number) => ({ __stepCountIs: n })),
}));
vi.mock('ai', async (importOriginal) => {
const actual = await importOriginal<typeof import('ai')>();
return {
...actual,
streamText: streamTextMock,
stepCountIs: stepCountIsMock,
};
});
const { AgentRuntime } = await import('@/services/reedy/runtime/AgentRuntime');
function fakeModel(): ChatModel {
return {
id: 'fake',
contextWindow: 32_000,
reservedOutput: 1_000,
supportsTools: true,
getLanguageModel: () =>
({ __mock: 'lm' }) as unknown as ReturnType<ChatModel['getLanguageModel']>,
};
}
async function* asyncIter<T>(items: T[]): AsyncGenerator<T> {
for (const item of items) yield item;
}
async function drain(stream: AsyncIterable<ReedyEvent>): Promise<ReedyEvent[]> {
const out: ReedyEvent[] = [];
for await (const ev of stream) out.push(ev);
return out;
}
beforeEach(() => {
streamTextMock.mockReset();
stepCountIsMock.mockClear();
});
describe('AgentRuntime — happy path', () => {
it('emits turn_start → text_delta(s) → finish-step → usage → done for a text-only stream', async () => {
streamTextMock.mockReturnValue({
fullStream: asyncIter([
{ type: 'text-delta', id: 't1', text: 'Hello, ' },
{ type: 'text-delta', id: 't1', text: 'world.' },
{
type: 'finish-step',
finishReason: 'stop',
usage: { inputTokens: 50, outputTokens: 10 },
},
{ type: 'finish', finishReason: 'stop', totalUsage: { inputTokens: 50, outputTokens: 10 } },
]),
});
const runtime = new AgentRuntime({
model: fakeModel(),
tools: new ToolRegistry(),
layers: [createPolicyLayer('You are Reedy.')],
});
const events = await drain(
runtime.runTurn({ sessionId: 's1', bookHash: 'bk1', userMessage: 'hi' }),
);
const types = events.map((e) => e.type);
expect(types[0]).toBe('turn_start');
expect(types).toContain('text_delta');
expect(types).toContain('step_finish');
expect(types).toContain('usage');
expect(types.at(-1)).toBe('done');
const done = events.at(-1)!;
if (done.type === 'done') {
expect(done.output.finishReason).toBe('stop');
expect(done.output.usage).toEqual({ promptTokens: 50, completionTokens: 10 });
}
});
});
describe('AgentRuntime — tool calls', () => {
function readTool(name: string): ReedyTool<{ q: string }, { hit: string }> {
return {
name,
description: 'find',
permission: 'read',
parallelSafe: true,
inputSchema: z.object({ q: z.string() }),
async run(args) {
return { hit: args.q };
},
};
}
it('emits tool_call → tool_result(ok) when streamText reports a successful dispatch', async () => {
streamTextMock.mockReturnValue({
fullStream: asyncIter([
{
type: 'tool-call',
toolCallId: 'tc1',
toolName: 'find',
input: { q: 'alice' },
},
{
type: 'tool-result',
toolCallId: 'tc1',
toolName: 'find',
output: { hit: 'alice' },
},
{
type: 'finish-step',
finishReason: 'tool-calls',
usage: { inputTokens: 5, outputTokens: 0 },
},
{ type: 'finish', finishReason: 'stop', totalUsage: { inputTokens: 5, outputTokens: 0 } },
]),
});
const reg = new ToolRegistry();
reg.register(readTool('find'));
const runtime = new AgentRuntime({
model: fakeModel(),
tools: reg,
layers: [createPolicyLayer('p')],
});
const events = await drain(
runtime.runTurn({ sessionId: 's1', bookHash: 'bk1', userMessage: 'find alice' }),
);
const toolCall = events.find((e) => e.type === 'tool_call')!;
const toolResult = events.find((e) => e.type === 'tool_result')!;
expect(toolCall.type).toBe('tool_call');
if (toolCall.type === 'tool_call') {
expect(toolCall.name).toBe('find');
expect(toolCall.permission).toBe('read');
expect(toolCall.args).toEqual({ q: 'alice' });
}
expect(toolResult.type).toBe('tool_result');
if (toolResult.type === 'tool_result' && toolResult.ok) {
expect(toolResult.result).toEqual({ hit: 'alice' });
}
});
it('emits tool_result(ok=false) with the tool error when streamText reports tool-error', async () => {
streamTextMock.mockReturnValue({
fullStream: asyncIter([
{
type: 'tool-error',
toolCallId: 'tc1',
toolName: 'find',
error: new Error('boom'),
},
{
type: 'finish-step',
finishReason: 'tool-calls',
usage: { inputTokens: 1, outputTokens: 0 },
},
{ type: 'finish', finishReason: 'stop', totalUsage: { inputTokens: 1, outputTokens: 0 } },
]),
});
const runtime = new AgentRuntime({
model: fakeModel(),
tools: new ToolRegistry(),
layers: [createPolicyLayer('p')],
});
const events = await drain(
runtime.runTurn({ sessionId: 's1', bookHash: 'bk1', userMessage: 'do thing' }),
);
const toolResult = events.find((e) => e.type === 'tool_result');
expect(toolResult).toBeDefined();
if (toolResult?.type === 'tool_result' && !toolResult.ok) {
expect(toolResult.error.kind).toBe('tool_runtime_error');
expect(toolResult.error.toolName).toBe('find');
}
});
it('extracts citation events from lookupPassage tool results by default', async () => {
streamTextMock.mockReturnValue({
fullStream: asyncIter([
{
type: 'tool-result',
toolCallId: 'tc1',
toolName: 'lookupPassage',
output: {
passages: [
{ cfi: 'epubcfi(/6/2)', sectionIndex: 0, chapter: 'Ch1', text: 'snippet a' },
{ cfi: 'epubcfi(/6/4)', sectionIndex: 1, chapter: 'Ch2', text: 'snippet b' },
],
status: 'ok',
},
},
{ type: 'finish-step', finishReason: 'stop', usage: { inputTokens: 1, outputTokens: 1 } },
{ type: 'finish', finishReason: 'stop', totalUsage: { inputTokens: 1, outputTokens: 1 } },
]),
});
const runtime = new AgentRuntime({
model: fakeModel(),
tools: new ToolRegistry(),
layers: [createPolicyLayer('p')],
});
const events = await drain(
runtime.runTurn({ sessionId: 's1', bookHash: 'bk1', userMessage: 'alice?' }),
);
const citations = events.filter((e) => e.type === 'citation');
expect(citations).toHaveLength(2);
if (citations[0]!.type === 'citation') {
expect(citations[0]!.cfi).toBe('epubcfi(/6/2)');
expect(citations[0]!.snippet).toBe('snippet a');
}
});
});
describe('AgentRuntime — abort + error paths', () => {
it('yields events.abort + done(finishReason=abort) when caller signal fires mid-stream', async () => {
const controller = new AbortController();
streamTextMock.mockImplementation(() => ({
fullStream: (async function* (): AsyncGenerator<unknown> {
yield { type: 'text-delta', id: 't1', text: 'part 1' };
controller.abort();
// Simulate the SDK respecting the signal — yield an abort part.
yield { type: 'abort', reason: 'caller cancelled' };
})(),
}));
const runtime = new AgentRuntime({
model: fakeModel(),
tools: new ToolRegistry(),
layers: [createPolicyLayer('p')],
});
const events = await drain(
runtime.runTurn({
sessionId: 's1',
bookHash: 'bk1',
userMessage: 'long stream',
signal: controller.signal,
}),
);
expect(events.some((e) => e.type === 'abort')).toBe(true);
const done = events.at(-1)!;
expect(done.type).toBe('done');
if (done.type === 'done') expect(done.output.finishReason).toBe('abort');
});
it('emits events.error and done(finishReason=error) when streamText throws (non-abort)', async () => {
// A real "fullStream" that rejects mid-iteration. We return an
// AsyncIterable whose next() rejects — Biome doesn't flag this the
// way it does an empty-bodied generator with throw.
streamTextMock.mockImplementation(() => ({
fullStream: {
[Symbol.asyncIterator](): AsyncIterator<unknown> {
return {
next() {
return Promise.reject(new Error('upstream provider 500'));
},
};
},
},
}));
const runtime = new AgentRuntime({
model: fakeModel(),
tools: new ToolRegistry(),
layers: [createPolicyLayer('p')],
});
const events = await drain(
runtime.runTurn({ sessionId: 's1', bookHash: 'bk1', userMessage: 'hi' }),
);
const errorEv = events.find((e) => e.type === 'error');
expect(errorEv).toBeDefined();
if (errorEv?.type === 'error') {
expect(errorEv.kind).toBe('model_error');
expect(errorEv.message).toContain('500');
}
const done = events.at(-1)!;
if (done.type === 'done') expect(done.output.finishReason).toBe('error');
});
});
describe('AgentRuntime — streamText arg wiring', () => {
it('passes the composed system prompt + tools + stopWhen(maxSteps) to streamText', async () => {
streamTextMock.mockReturnValue({
fullStream: asyncIter([
{ type: 'finish', finishReason: 'stop', totalUsage: { inputTokens: 0, outputTokens: 0 } },
]),
});
const reg = new ToolRegistry();
reg.register({
name: 'a',
description: 'a desc',
permission: 'read',
parallelSafe: true,
inputSchema: z.object({}),
async run() {
return null;
},
});
const runtime = new AgentRuntime({
model: fakeModel(),
tools: reg,
layers: [createPolicyLayer('POL.')],
maxSteps: 4,
});
await drain(runtime.runTurn({ sessionId: 's', bookHash: 'b', userMessage: 'q' }));
expect(streamTextMock).toHaveBeenCalledTimes(1);
const call = streamTextMock.mock.calls[0]![0] as {
system: string;
tools?: Record<string, unknown>;
stopWhen: unknown;
messages: Array<{ role: string; content: string }>;
};
expect(call.system).toContain('POL.');
expect(call.tools).toBeDefined();
expect(call.tools!['a']).toBeDefined();
expect(call.stopWhen).toEqual({ __stepCountIs: 4 });
expect(call.messages.at(-1)).toEqual({ role: 'user', content: 'q' });
});
it('does not pass tools when the registry is empty', async () => {
streamTextMock.mockReturnValue({
fullStream: asyncIter([
{ type: 'finish', finishReason: 'stop', totalUsage: { inputTokens: 0, outputTokens: 0 } },
]),
});
const runtime = new AgentRuntime({
model: fakeModel(),
tools: new ToolRegistry(),
layers: [createPolicyLayer('p')],
});
await drain(runtime.runTurn({ sessionId: 's', bookHash: 'b', userMessage: 'q' }));
const call = streamTextMock.mock.calls[0]![0] as { tools?: unknown };
expect(call.tools).toBeUndefined();
});
});
@@ -0,0 +1,301 @@
import { streamText, stepCountIs, type ModelMessage } from 'ai';
import type { ChatModel } from '../models/ChatModel';
import type { PromptLayer, BuiltContext } from '../context';
import { buildPromptContext } from '../context';
import type { ToolRegistry } from '../tools/ToolRegistry';
import type { ToolContext } from '../tools/types';
import { events, type ReedyEvent, type ReedyTurnOutput } from './events';
import { isAbortError } from './abort';
import { ReedyToolError } from './errors';
const DEFAULT_MAX_STEPS = 8;
/**
* Per-turn input the runtime needs. `sessionId` + `assistantMessageId`
* thread the turn through downstream persistence (Phase 3+); the
* runtime itself stays stateless beyond the per-turn ToolContext.
*/
export interface RunTurnInput {
sessionId: string;
bookHash: string;
userMessage: string;
/**
* Prior message history to seed the conversation. The runtime adds the
* user message itself; pass everything before it. Defaults to [].
*/
history?: ModelMessage[];
/**
* Optional assistant message id. Generated if not provided so callers
* who don't need stable ids (tests) don't have to make one up.
*/
assistantMessageId?: string;
/** Caller signal — runtime composes with its own so tools see both. */
signal?: AbortSignal;
}
/**
* Synthesized citation event the runtime emits when a tool result
* implies one. `extractCitations` is the customization point.
*/
export interface CitationLike {
cfi: string;
sectionIndex: number;
chapterTitle?: string;
snippet: string;
}
export interface AgentRuntimeOptions {
model: ChatModel;
tools: ToolRegistry;
/**
* PromptLayers the runtime composes into the system message. Pass
* factories that resolve current state at runTurn() time (e.g.
* ReadingLayer wraps the latest snapshot). The runtime doesn't own
* layer construction — callers wire whatever applies.
*/
layers: PromptLayer[];
/** Cap on agent steps inside one turn. @default 8 */
maxSteps?: number;
/** Optional per-call permission prompt; defaults to read auto-approve, others deny. */
requestPermission?: ToolContext['requestPermission'];
/**
* Optional citation extractor — called for every successful tool
* result. Returning an array of citations causes the runtime to emit
* a `{ type: 'citation', ... }` event per item. Defaults to extracting
* from lookupPassage-shaped results.
*/
extractCitations?: (toolName: string, result: unknown) => CitationLike[];
}
/**
* Composes ChatModel + ToolRegistry + PromptContextBuilder into one
* runTurn() entrypoint that streams ReedyEvents (per plan §6 / §2.6).
*
* The runtime owns the per-turn glue (build system prompt, construct
* ToolContext, dispatch streamText, fan TextStreamParts out as
* ReedyEvents) but stays stateless — sessions, persistence, memory
* writes are wired in by callers (Phase 3+ services hook in via
* extractCitations / history / requestPermission).
*
* Vercel SDK does the multi-step tool loop internally via
* stopWhen=stepCountIs(maxSteps); we walk fullStream once and map each
* part to a ReedyEvent without re-implementing the dispatch loop. The
* plan reserves a per-step outer loop for memory-write checkpoints —
* defer that until Phase 3 actually surfaces a hook.
*/
export class AgentRuntime {
private readonly maxSteps: number;
constructor(private readonly opts: AgentRuntimeOptions) {
this.maxSteps = opts.maxSteps ?? DEFAULT_MAX_STEPS;
}
runTurn(input: RunTurnInput): AsyncIterable<ReedyEvent> {
return this.runTurnGenerator(input);
}
private async *runTurnGenerator(input: RunTurnInput): AsyncIterable<ReedyEvent> {
const assistantMessageId = input.assistantMessageId ?? randomId('msg');
yield events.turnStart(input.sessionId, assistantMessageId);
// Build the system prompt and the ToolContext the registry will hand
// to every tool dispatched in this turn.
let ctx: BuiltContext;
try {
ctx = buildPromptContext({ model: this.opts.model, layers: this.opts.layers });
} catch (err) {
yield events.error('unknown', `failed to build prompt context: ${stringifyErr(err)}`, false);
yield events.done({
sessionId: input.sessionId,
assistantMessageId,
finishReason: 'error',
});
return;
}
const toolCtx: ToolContext = {
bookHash: input.bookHash,
sessionId: input.sessionId,
assistantMessageId,
signal: input.signal ?? new AbortController().signal,
// ToolRegistry only invokes requestPermission for non-read tools, so
// the default here is reached only when the caller didn't supply a
// prompt — in which case we deny rather than silently writing /
// navigating without explicit user consent.
requestPermission: this.opts.requestPermission ?? (async (): Promise<boolean> => false),
};
const messages: ModelMessage[] = [
...(input.history ?? []),
{ role: 'user', content: input.userMessage },
];
let lastUsage: { promptTokens: number; completionTokens: number } | undefined;
let lastFinishReason: ReedyTurnOutput['finishReason'] = 'stop';
let aborted = false;
const extractCitations = this.opts.extractCitations ?? defaultExtractCitations;
try {
const result = streamText({
model: this.opts.model.getLanguageModel(),
system: ctx.system,
messages,
tools:
this.opts.tools.list().length > 0 ? this.opts.tools.toVercelToolSet(toolCtx) : undefined,
stopWhen: stepCountIs(this.maxSteps),
abortSignal: input.signal,
});
let stepIndex = 0;
for await (const part of result.fullStream) {
switch (part.type) {
case 'text-delta':
if (part.text.length > 0) yield events.textDelta(part.text);
break;
case 'tool-call':
yield events.toolCall({
id: part.toolCallId,
name: part.toolName,
args: part.input,
permission: this.opts.tools.get(part.toolName)?.permission ?? 'read',
});
break;
case 'tool-result': {
const toolName = part.toolName;
yield events.toolResultOk({
id: part.toolCallId,
name: toolName,
result: part.output,
durationMs: 0, // Vercel doesn't surface tool duration here; runtime
// can compute it from start-step + finish-step if needed.
});
// Synthesize citation events from the tool result.
for (const c of extractCitations(toolName, part.output)) {
yield events.citation(c);
}
break;
}
case 'tool-error': {
const err =
part.error instanceof ReedyToolError
? part.error
: new ReedyToolError(stringifyErr(part.error), {
kind: 'tool_runtime_error',
toolName: part.toolName,
cause: part.error,
});
yield events.toolResultErr({
id: part.toolCallId,
name: part.toolName,
error: err,
});
break;
}
case 'finish-step':
yield events.stepFinish(stepIndex, normaliseFinishReason(part.finishReason));
if (part.usage) {
lastUsage = {
promptTokens: part.usage.inputTokens ?? 0,
completionTokens: part.usage.outputTokens ?? 0,
};
yield events.usage(lastUsage.promptTokens, lastUsage.completionTokens);
}
stepIndex++;
break;
case 'finish':
lastFinishReason = normaliseTurnFinish(part.finishReason);
if (part.totalUsage) {
lastUsage = {
promptTokens: part.totalUsage.inputTokens ?? 0,
completionTokens: part.totalUsage.outputTokens ?? 0,
};
}
break;
case 'abort':
aborted = true;
break;
case 'error':
yield events.error('model_error', stringifyErr(part.error), true);
lastFinishReason = 'error';
break;
default:
// text-start / text-end / reasoning / source / file / start /
// start-step / tool-input-* / tool-output-denied are not
// load-bearing for the v1 event stream.
break;
}
}
} catch (err) {
if (isAbortError(err) || input.signal?.aborted) {
aborted = true;
} else {
yield events.error('model_error', stringifyErr(err), true);
lastFinishReason = 'error';
}
}
if (aborted) {
yield events.abort(/* partial */ true);
lastFinishReason = 'abort';
}
yield events.done({
sessionId: input.sessionId,
assistantMessageId,
finishReason: lastFinishReason,
usage: lastUsage,
});
}
}
function defaultExtractCitations(toolName: string, result: unknown): CitationLike[] {
if (toolName === 'lookupPassage' && result && typeof result === 'object') {
const r = result as {
passages?: Array<{
cfi?: string;
chapter?: string;
sectionIndex?: number;
text?: string;
}>;
};
if (!Array.isArray(r.passages)) return [];
return r.passages
.filter((p) => p.cfi && typeof p.text === 'string')
.map((p) => ({
cfi: p.cfi!,
sectionIndex: p.sectionIndex ?? 0,
chapterTitle: p.chapter,
snippet: (p.text ?? '').slice(0, 200),
}));
}
if (toolName === 'addCitation' && result && typeof result === 'object') {
// addCitation acks via { ok: true }; the citation was already pushed
// through its onCite callback. No additional event needed here.
return [];
}
return [];
}
function normaliseFinishReason(reason: string): 'stop' | 'tool-calls' | 'length' {
if (reason === 'tool-calls') return 'tool-calls';
if (reason === 'length') return 'length';
return 'stop';
}
function normaliseTurnFinish(reason: string): ReedyTurnOutput['finishReason'] {
if (reason === 'length') return 'length';
if (reason === 'tool-calls') return 'tool-error';
if (reason === 'error') return 'error';
return 'stop';
}
function stringifyErr(err: unknown): string {
if (err instanceof Error) return err.message;
return String(err);
}
function randomId(prefix: string): string {
if (typeof crypto !== 'undefined' && 'randomUUID' in crypto)
return `${prefix}-${crypto.randomUUID()}`;
return `${prefix}-${Date.now()}-${Math.random().toString(36).slice(2, 10)}`;
}
@@ -0,0 +1,46 @@
/**
* Tiny abort-signal helpers used by the runtime. Compose the
* caller's signal with a runtime-local one so internal cancellation
* (e.g. context_overflow retry) doesn't propagate out to the caller
* unless the caller themselves cancelled.
*/
/**
* Returns an AbortSignal that fires when EITHER input fires. The
* `cleanup` function unsubscribes listeners so we don't leak across
* long-lived consumers.
*/
export function anySignal(...signals: Array<AbortSignal | undefined>): {
signal: AbortSignal;
cleanup: () => void;
} {
const controller = new AbortController();
const listeners: Array<{ signal: AbortSignal; fn: () => void }> = [];
for (const sig of signals) {
if (!sig) continue;
if (sig.aborted) {
controller.abort(sig.reason);
break;
}
const fn = (): void => controller.abort(sig.reason);
sig.addEventListener('abort', fn, { once: true });
listeners.push({ signal: sig, fn });
}
const cleanup = (): void => {
for (const { signal, fn } of listeners) signal.removeEventListener('abort', fn);
};
return { signal: controller.signal, cleanup };
}
/**
* Type guard so the runtime can distinguish AbortError-shaped exceptions
* from real errors without doing instanceof against DOMException (which
* isn't always available in Node).
*/
export function isAbortError(err: unknown): boolean {
if (!err || typeof err !== 'object') return false;
const name = (err as { name?: unknown }).name;
return name === 'AbortError' || name === 'TimeoutError';
}