feat(reedy): wire Phase 5 skills + Phase 3 memory tools into the agent runtime (#4309)
This commit is contained in:
@@ -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<string, unknown> };
|
||||
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();
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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<string, unknown>;
|
||||
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', () => {
|
||||
|
||||
@@ -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<ToolRegistry['toVercelToolSet']> | 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[] {
|
||||
|
||||
@@ -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<string, unknown> = {};
|
||||
for (const t of this.tools.values()) {
|
||||
if (allowed && !allowed.has(t.name)) continue;
|
||||
set[t.name] = vercelTool({
|
||||
description: t.description,
|
||||
inputSchema: t.inputSchema,
|
||||
|
||||
@@ -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<Skill[]>([]);
|
||||
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<string | null>(null);
|
||||
const activeSkill = useMemo<Skill | null>(
|
||||
() => 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<SkillInstructions | null>(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<typeof createSkillLayer> => {
|
||||
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({
|
||||
}
|
||||
/>
|
||||
</div>
|
||||
<Composer isRunning={isRunning} onSend={handleSend} onAbort={abort} />
|
||||
<Composer
|
||||
isRunning={isRunning}
|
||||
onSend={handleSend}
|
||||
onAbort={abort}
|
||||
skills={availableSkills.map((s) => ({
|
||||
id: s.id,
|
||||
name: s.name,
|
||||
description: s.description,
|
||||
}))}
|
||||
activeSkillId={activeSkillId}
|
||||
onSkillSelect={setActiveSkillId}
|
||||
/>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
@@ -17,7 +17,13 @@ export function useReedyTurn(runtime: AgentRuntime | null) {
|
||||
const cancelRef = useRef<AbortController | null>(null);
|
||||
|
||||
const send = useCallback(
|
||||
async (args: { sessionId: string; bookHash: string; userMessage: string }): Promise<void> => {
|
||||
async (args: {
|
||||
sessionId: string;
|
||||
bookHash: string;
|
||||
userMessage: string;
|
||||
/** Optional per-turn tool-name allowlist (e.g. active skill's). */
|
||||
toolAllowlist?: readonly string[] | null;
|
||||
}): Promise<void> => {
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user