From 6be7606da464c5703cee2be84577fe58636e662a Mon Sep 17 00:00:00 2001 From: Huang Xin Date: Tue, 26 May 2026 15:31:19 +0800 Subject: [PATCH] feat(reedy): wire Phase 5 skills + Phase 3 memory tools into the agent runtime (#4309) --- .../src/__tests__/reedy/AgentRuntime.test.ts | 70 ++++++++++++ .../src/__tests__/reedy/ToolRegistry.test.ts | 24 ++++ .../services/reedy/runtime/AgentRuntime.ts | 29 ++++- .../src/services/reedy/tools/ToolRegistry.ts | 12 +- .../src/services/reedy/ui/ReedyAssistant.tsx | 108 ++++++++++++++++-- .../src/services/reedy/ui/useReedyTurn.ts | 9 +- 6 files changed, 237 insertions(+), 15 deletions(-) diff --git a/apps/readest-app/src/__tests__/reedy/AgentRuntime.test.ts b/apps/readest-app/src/__tests__/reedy/AgentRuntime.test.ts index dce7d416..fe19068d 100644 --- a/apps/readest-app/src/__tests__/reedy/AgentRuntime.test.ts +++ b/apps/readest-app/src/__tests__/reedy/AgentRuntime.test.ts @@ -359,4 +359,74 @@ describe('AgentRuntime — streamText arg wiring', () => { const call = streamTextMock.mock.calls[0]![0] as { tools?: unknown }; expect(call.tools).toBeUndefined(); }); + + it('honors toolAllowlist on the input — only allowlisted tools reach streamText', async () => { + streamTextMock.mockReturnValue({ + fullStream: asyncIter([ + { type: 'finish', finishReason: 'stop', totalUsage: { inputTokens: 0, outputTokens: 0 } }, + ]), + }); + const reg = new ToolRegistry(); + for (const name of ['lookupPassage', 'getReadingContext', 'navigateToCfi']) { + reg.register({ + name, + description: name, + permission: 'read', + parallelSafe: true, + inputSchema: z.object({}), + async run() { + return null; + }, + }); + } + const runtime = new AgentRuntime({ + model: fakeModel(), + tools: reg, + layers: [createPolicyLayer('p')], + }); + await drain( + runtime.runTurn({ + sessionId: 's', + bookHash: 'b', + userMessage: 'q', + toolAllowlist: ['lookupPassage', 'getReadingContext'], + }), + ); + const call = streamTextMock.mock.calls[0]![0] as { tools?: Record }; + expect(Object.keys(call.tools!).sort()).toEqual(['getReadingContext', 'lookupPassage']); + }); + + it('omits tools entirely when the allowlist matches zero registered tools', async () => { + streamTextMock.mockReturnValue({ + fullStream: asyncIter([ + { type: 'finish', finishReason: 'stop', totalUsage: { inputTokens: 0, outputTokens: 0 } }, + ]), + }); + const reg = new ToolRegistry(); + reg.register({ + name: 'lookupPassage', + description: 'find', + permission: 'read', + parallelSafe: true, + inputSchema: z.object({}), + async run() { + return null; + }, + }); + const runtime = new AgentRuntime({ + model: fakeModel(), + tools: reg, + layers: [createPolicyLayer('p')], + }); + await drain( + runtime.runTurn({ + sessionId: 's', + bookHash: 'b', + userMessage: 'q', + toolAllowlist: ['nonExistentTool'], + }), + ); + const call = streamTextMock.mock.calls[0]![0] as { tools?: unknown }; + expect(call.tools).toBeUndefined(); + }); }); diff --git a/apps/readest-app/src/__tests__/reedy/ToolRegistry.test.ts b/apps/readest-app/src/__tests__/reedy/ToolRegistry.test.ts index 03b0db46..fa9ceca7 100644 --- a/apps/readest-app/src/__tests__/reedy/ToolRegistry.test.ts +++ b/apps/readest-app/src/__tests__/reedy/ToolRegistry.test.ts @@ -214,6 +214,30 @@ describe('ToolRegistry', () => { expect(Object.keys(set).sort()).toEqual(['a', 'b']); expect(typeof set['a']!.execute).toBe('function'); }); + + it('with an allowlist, only the named tools appear in the set', () => { + reg.register(readTool('a')); + reg.register(readTool('b')); + reg.register(readTool('c')); + const set = reg.toVercelToolSet(ctxFor(), { allowlist: ['a', 'c'] }) as Record< + string, + unknown + >; + expect(Object.keys(set).sort()).toEqual(['a', 'c']); + }); + + it('allowlist=null behaves like no allowlist (all tools)', () => { + reg.register(readTool('a')); + reg.register(readTool('b')); + const set = reg.toVercelToolSet(ctxFor(), { allowlist: null }) as Record; + expect(Object.keys(set).sort()).toEqual(['a', 'b']); + }); + + it('empty allowlist returns an empty set', () => { + reg.register(readTool('a')); + const set = reg.toVercelToolSet(ctxFor(), { allowlist: [] }); + expect(Object.keys(set)).toEqual([]); + }); }); describe('unknown / runtime errors', () => { diff --git a/apps/readest-app/src/services/reedy/runtime/AgentRuntime.ts b/apps/readest-app/src/services/reedy/runtime/AgentRuntime.ts index c9c86a16..74d8fe23 100644 --- a/apps/readest-app/src/services/reedy/runtime/AgentRuntime.ts +++ b/apps/readest-app/src/services/reedy/runtime/AgentRuntime.ts @@ -31,6 +31,13 @@ export interface RunTurnInput { assistantMessageId?: string; /** Caller signal — runtime composes with its own so tools see both. */ signal?: AbortSignal; + /** + * Optional per-turn tool-name allowlist. When provided, only tools + * whose name appears here are exposed to the model. Used by the + * notebook to honor the active skill's `tool_allowlist`. Null or + * undefined means every registered tool is available. + */ + toolAllowlist?: readonly string[] | null; } /** @@ -140,8 +147,7 @@ export class AgentRuntime { model: this.opts.model.getLanguageModel(), system: ctx.system, messages, - tools: - this.opts.tools.list().length > 0 ? this.opts.tools.toVercelToolSet(toolCtx) : undefined, + tools: this.buildToolSet(toolCtx, input.toolAllowlist), stopWhen: stepCountIs(this.maxSteps), abortSignal: input.signal, }); @@ -246,6 +252,25 @@ export class AgentRuntime { usage: lastUsage, }); } + + /** + * Build the Vercel ToolSet for this turn, applying the per-turn + * allowlist (active skill's `tool_allowlist`) if one was supplied. + * Returns undefined when no tools end up exposed, so streamText skips + * tool-call wiring entirely instead of advertising an empty catalog. + */ + private buildToolSet( + ctx: ToolContext, + allowlist: readonly string[] | null | undefined, + ): ReturnType | undefined { + const all = this.opts.tools.list(); + if (all.length === 0) return undefined; + if (allowlist == null) return this.opts.tools.toVercelToolSet(ctx); + const allowedNames = new Set(allowlist); + const overlap = all.filter((t) => allowedNames.has(t.name)); + if (overlap.length === 0) return undefined; + return this.opts.tools.toVercelToolSet(ctx, { allowlist }); + } } function defaultExtractCitations(toolName: string, result: unknown): CitationLike[] { diff --git a/apps/readest-app/src/services/reedy/tools/ToolRegistry.ts b/apps/readest-app/src/services/reedy/tools/ToolRegistry.ts index 7af3ca3f..2f32f56f 100644 --- a/apps/readest-app/src/services/reedy/tools/ToolRegistry.ts +++ b/apps/readest-app/src/services/reedy/tools/ToolRegistry.ts @@ -46,10 +46,20 @@ export class ToolRegistry { * Adapt every registered tool to the Vercel ToolSet shape. Each tool's * `execute` is wrapped per the policy above. The caller passes the * per-turn ToolContext once; every dispatch in that turn reuses it. + * + * When `options.allowlist` is provided, only tools whose name appears + * in the array are exposed to the model. The plan §5 has skills carry + * an optional `tool_allowlist`; the runtime computes the active + * skill's allowlist and forwards it here. */ - toVercelToolSet(ctx: ToolContext): ToolSet { + toVercelToolSet( + ctx: ToolContext, + options: { allowlist?: readonly string[] | null } = {}, + ): ToolSet { + const allowed = options.allowlist == null ? null : new Set(options.allowlist); const set: Record = {}; for (const t of this.tools.values()) { + if (allowed && !allowed.has(t.name)) continue; set[t.name] = vercelTool({ description: t.description, inputSchema: t.inputSchema, diff --git a/apps/readest-app/src/services/reedy/ui/ReedyAssistant.tsx b/apps/readest-app/src/services/reedy/ui/ReedyAssistant.tsx index 6ee973b9..84a68a52 100644 --- a/apps/readest-app/src/services/reedy/ui/ReedyAssistant.tsx +++ b/apps/readest-app/src/services/reedy/ui/ReedyAssistant.tsx @@ -15,15 +15,26 @@ import { createGetReadingContextTool, createGetSelectionTool, createLookupPassageTool, + createSearchBookMemoryTool, + createSearchSessionMemoryTool, + createSearchUserMemoryTool, + createWriteBookMemoryTool, + createWriteUserMemoryTool, type ReadingContextSnapshot, } from '../tools/builtins'; import { DEFAULT_POLICY, + createBookMemoryLayer, createPolicyLayer, createReadingLayer, createSkillLayer, createToolCatalogLayer, + createUserMemoryLayer, + type SkillInstructions, } from '../context'; +import { MemoryService } from '../memory/MemoryService'; +import { SkillRegistry } from '../skills/SkillRegistry'; +import type { Skill } from '../skills/types'; import { useReedyStore } from '../store/reedyStore'; import { useReedyTurn } from './useReedyTurn'; import { AgentThread } from './AgentThread'; @@ -37,6 +48,8 @@ export interface ReedyAssistantProps { bookKey: string; aiSettings: AISettings; readingContext: ReadingContextSnapshot; + /** Stable id for the current user (used as the scope_key for user memory). */ + userId?: string; /** Wired by the notebook to `getView(bookKey)?.goTo(cfi)` on click. */ onNavigateToCfi?: (cfi: string) => void; } @@ -56,26 +69,44 @@ export function ReedyAssistant({ bookHash, aiSettings, readingContext, + userId = 'local', onNavigateToCfi, }: ReedyAssistantProps) { const models = useMemo(() => createReedyModels(aiSettings), [aiSettings]); // Lazily open reedy.db on first mount. The promise resolves once and - // we share the same ReedyDb + Indexer + Retriever for the lifetime of - // this component instance. + // we share the same ReedyDb + Indexer + Retriever + MemoryService + + // SkillRegistry for the lifetime of this component instance. The + // skill registry is seeded on first init() with the 3 builtins; on + // subsequent mounts the call is a no-op. const [reedy, setReedy] = useState<{ db: ReedyDb; indexer: BookIndexer; retriever: BookRetriever; + memory: MemoryService; + skills: SkillRegistry; } | null>(null); + const [availableSkills, setAvailableSkills] = useState([]); useEffect(() => { let alive = true; void appService .openDatabase('reedy', 'reedy.db', 'Data', { experimental: ['index_method'] }) - .then((svc) => { + .then(async (svc) => { if (!alive) return; const db = new ReedyDb(svc); - setReedy({ db, indexer: new BookIndexer(db), retriever: new BookRetriever(db) }); + const skills = new SkillRegistry(svc); + await skills.init(); + const memory = new MemoryService(db, models.embedding); + const enabledSkills = await skills.listEnabled(); + if (!alive) return; + setReedy({ + db, + indexer: new BookIndexer(db), + retriever: new BookRetriever(db), + memory, + skills, + }); + setAvailableSkills(enabledSkills); }) .catch((err) => { console.error('[Reedy] failed to open reedy.db', err); @@ -83,7 +114,15 @@ export function ReedyAssistant({ return () => { alive = false; }; - }, [appService]); + }, [appService, models.embedding]); + + // Active skill selection — null means no skill (the model gets the + // default Policy + Reading + ToolCatalog system prompt). + const [activeSkillId, setActiveSkillId] = useState(null); + const activeSkill = useMemo( + () => availableSkills.find((s) => s.id === activeSkillId) ?? null, + [availableSkills, activeSkillId], + ); // Snapshot the reading context in a ref so the tool factories don't // re-render the registry on every page turn. @@ -92,7 +131,20 @@ export function ReedyAssistant({ readingRef.current = readingContext; }, [readingContext]); - // Build the tool registry + runtime once per (reedy ready, model) pair. + // Active-skill state is read via a ref inside the SkillLayer's + // resolution so the runtime memo doesn't have to rebuild on every + // chip click. The layer captures the closure once; the closure peeks + // at the ref at prompt-build time. + const activeSkillRef = useRef(null); + useEffect(() => { + activeSkillRef.current = activeSkill + ? { id: activeSkill.id, instructions: activeSkill.instructions } + : null; + }, [activeSkill]); + + // Build the tool registry + runtime once per (reedy ready, model) + // pair. The userId and bookHash flow through the memory tool factories + // so each tool dispatches under the right scope_key. const runtime = useMemo(() => { if (!reedy) return null; const reg = new ToolRegistry(); @@ -113,16 +165,34 @@ export function ReedyAssistant({ // the store via a closure here. }), ); + // Memory tools (Phase 3.1) — bookHash scopes book memory; userId + // scopes user memory; sessionId mirrors the runtime's per-turn + // sessionId which we set to bookHash today (sessions are + // book-scoped in this build). + reg.register(createSearchBookMemoryTool({ service: reedy.memory, scopeKey: () => bookHash })); + reg.register(createWriteBookMemoryTool({ service: reedy.memory, scopeKey: () => bookHash })); + reg.register(createSearchUserMemoryTool({ service: reedy.memory, scopeKey: () => userId })); + reg.register(createWriteUserMemoryTool({ service: reedy.memory, scopeKey: () => userId })); + reg.register( + createSearchSessionMemoryTool({ service: reedy.memory, scopeKey: () => bookHash }), + ); const layers = [ createPolicyLayer(DEFAULT_POLICY), - createSkillLayer(null), + // Use a tiny wrapper so the SkillLayer resolves the *current* + // active skill each time the runtime rebuilds the prompt. + ((): ReturnType => { + const snapshot = activeSkillRef.current; + return createSkillLayer(snapshot); + })(), createReadingLayer(readingRef.current), createToolCatalogLayer(reg.list()), + createBookMemoryLayer(() => ''), + createUserMemoryLayer(() => ''), ]; return new AgentRuntime({ model: models.chat, tools: reg, layers }); - }, [reedy, models.chat, models.embedding, bookHash]); + }, [reedy, models.chat, models.embedding, bookHash, userId]); const messages = useReedyStore((s) => s.messages); const isRunning = useReedyStore((s) => s.isRunning); @@ -192,9 +262,14 @@ export function ReedyAssistant({ const handleSend = useCallback( (text: string) => { if (!runtime) return; - void send({ sessionId: bookHash, bookHash, userMessage: text }); + void send({ + sessionId: bookHash, + bookHash, + userMessage: text, + toolAllowlist: activeSkill?.toolAllowlist ?? null, + }); }, - [runtime, send, bookHash], + [runtime, send, bookHash, activeSkill], ); if (!aiSettings.enabled) { @@ -241,7 +316,18 @@ export function ReedyAssistant({ } /> - + ({ + id: s.id, + name: s.name, + description: s.description, + }))} + activeSkillId={activeSkillId} + onSkillSelect={setActiveSkillId} + /> ); } diff --git a/apps/readest-app/src/services/reedy/ui/useReedyTurn.ts b/apps/readest-app/src/services/reedy/ui/useReedyTurn.ts index e0531f90..2f4f6d96 100644 --- a/apps/readest-app/src/services/reedy/ui/useReedyTurn.ts +++ b/apps/readest-app/src/services/reedy/ui/useReedyTurn.ts @@ -17,7 +17,13 @@ export function useReedyTurn(runtime: AgentRuntime | null) { const cancelRef = useRef(null); const send = useCallback( - async (args: { sessionId: string; bookHash: string; userMessage: string }): Promise => { + async (args: { + sessionId: string; + bookHash: string; + userMessage: string; + /** Optional per-turn tool-name allowlist (e.g. active skill's). */ + toolAllowlist?: readonly string[] | null; + }): Promise => { if (!runtime) return; // Cancel any in-flight turn before starting a new one so we never // have two AgentRuntime streams mutating the same active message. @@ -33,6 +39,7 @@ export function useReedyTurn(runtime: AgentRuntime | null) { bookHash: args.bookHash, userMessage: args.userMessage, signal: controller.signal, + toolAllowlist: args.toolAllowlist ?? null, }; // The first event the runtime yields is `turn_start` carrying the