From df2698fa0eacd1e5283d4e43154b3738eb02bc40 Mon Sep 17 00:00:00 2001 From: Huang Xin Date: Tue, 26 May 2026 10:14:53 +0800 Subject: [PATCH] =?UTF-8?q?feat(reedy):=20Appendix=20A=20=C2=B7=20Phase=20?= =?UTF-8?q?2.6=20=E2=80=94=20AgentRuntime=20+=20abort=20helper=20(#4301)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../src/__tests__/reedy/AgentRuntime.test.ts | 362 ++++++++++++++++++ .../services/reedy/runtime/AgentRuntime.ts | 301 +++++++++++++++ .../src/services/reedy/runtime/abort.ts | 46 +++ 3 files changed, 709 insertions(+) create mode 100644 apps/readest-app/src/__tests__/reedy/AgentRuntime.test.ts create mode 100644 apps/readest-app/src/services/reedy/runtime/AgentRuntime.ts create mode 100644 apps/readest-app/src/services/reedy/runtime/abort.ts diff --git a/apps/readest-app/src/__tests__/reedy/AgentRuntime.test.ts b/apps/readest-app/src/__tests__/reedy/AgentRuntime.test.ts new file mode 100644 index 00000000..dce7d416 --- /dev/null +++ b/apps/readest-app/src/__tests__/reedy/AgentRuntime.test.ts @@ -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(); + 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, + }; +} + +async function* asyncIter(items: T[]): AsyncGenerator { + for (const item of items) yield item; +} + +async function drain(stream: AsyncIterable): Promise { + 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 { + 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 { + 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; + 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(); + }); +}); diff --git a/apps/readest-app/src/services/reedy/runtime/AgentRuntime.ts b/apps/readest-app/src/services/reedy/runtime/AgentRuntime.ts new file mode 100644 index 00000000..c9c86a16 --- /dev/null +++ b/apps/readest-app/src/services/reedy/runtime/AgentRuntime.ts @@ -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 { + return this.runTurnGenerator(input); + } + + private async *runTurnGenerator(input: RunTurnInput): AsyncIterable { + 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 => 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)}`; +} diff --git a/apps/readest-app/src/services/reedy/runtime/abort.ts b/apps/readest-app/src/services/reedy/runtime/abort.ts new file mode 100644 index 00000000..2e8f8f53 --- /dev/null +++ b/apps/readest-app/src/services/reedy/runtime/abort.ts @@ -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): { + 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'; +}