From a0bc6f30edb726e538b7d143a7c1af191ff5a488 Mon Sep 17 00:00:00 2001 From: AP3X Date: Sat, 11 Apr 2026 14:40:03 -0700 Subject: [PATCH] feat: add DynamicToolRegistry for hot-loadable tool extensions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduces a mutable handler map that sits behind a single compiled StateGraph node. Extensions register/unregister at runtime without recompiling the graph — the topology never changes. - DynamicToolRegistry class with register, unregister, list, has, asNode (StateGraph-compatible node fn), asSchema (OpenAI format) - ToolHandler / ToolResult types - Subpath export at @oni.bot/core/registry - 18 tests across 4 files covering execution, unknown-tool error handling, lifecycle, and schema output --- package.json | 4 + .../registry/register-and-execute.test.ts | 95 +++++++++++ .../registry/register-unregister.test.ts | 126 ++++++++++++++ src/__tests__/registry/schema-output.test.ts | 114 +++++++++++++ src/__tests__/registry/unknown-tool.test.ts | 74 ++++++++ src/registry/README.md | 125 ++++++++++++++ src/registry/dynamic-tool-registry.ts | 161 ++++++++++++++++++ src/registry/index.ts | 3 + src/registry/types.ts | 16 ++ 9 files changed, 718 insertions(+) create mode 100644 src/__tests__/registry/register-and-execute.test.ts create mode 100644 src/__tests__/registry/register-unregister.test.ts create mode 100644 src/__tests__/registry/schema-output.test.ts create mode 100644 src/__tests__/registry/unknown-tool.test.ts create mode 100644 src/registry/README.md create mode 100644 src/registry/dynamic-tool-registry.ts create mode 100644 src/registry/index.ts create mode 100644 src/registry/types.ts diff --git a/package.json b/package.json index bdda689..982021e 100644 --- a/package.json +++ b/package.json @@ -90,6 +90,10 @@ "./config": { "import": "./dist/config/index.js", "types": "./dist/config/index.d.ts" + }, + "./registry": { + "import": "./dist/registry/index.js", + "types": "./dist/registry/index.d.ts" } }, "workspaces": [ diff --git a/src/__tests__/registry/register-and-execute.test.ts b/src/__tests__/registry/register-and-execute.test.ts new file mode 100644 index 0000000..2851fcc --- /dev/null +++ b/src/__tests__/registry/register-and-execute.test.ts @@ -0,0 +1,95 @@ +import { describe, it, expect } from "vitest"; +import { DynamicToolRegistry } from "../../registry/index.js"; +import type { DynamicToolState } from "../../registry/index.js"; +import type { ToolResult } from "../../registry/index.js"; + +describe("DynamicToolRegistry — register and execute", () => { + it("registers a tool and executes it via asNode()", async () => { + const registry = new DynamicToolRegistry(); + + registry.register("greet", async (args) => ({ + tool_name: "greet", + success: true, + output: `Hello, ${args.name}!`, + })); + + const node = registry.asNode(); + + const state: DynamicToolState = { + messages: [], + pendingTools: [{ name: "greet", args: { name: "World" }, id: "tc-1" }], + }; + + const result = await node(state); + expect(result).toBeDefined(); + + const partial = result as Partial; + expect(partial.messages).toHaveLength(1); + + const msg = partial.messages![0]; + expect(msg.role).toBe("tool"); + expect(msg.name).toBe("greet"); + expect(msg.tool_call_id).toBe("tc-1"); + + const body: ToolResult = JSON.parse(msg.content); + expect(body.success).toBe(true); + expect(body.output).toBe("Hello, World!"); + expect(body.tool_name).toBe("greet"); + }); + + it("returns empty update when pendingTools is empty", async () => { + const registry = new DynamicToolRegistry(); + registry.register("noop", async () => ({ + tool_name: "noop", + success: true, + output: "ok", + })); + + const node = registry.asNode(); + const result = await node({ messages: [], pendingTools: [] }); + expect(result).toEqual({}); + }); + + it("passes state to handler as second argument", async () => { + const registry = new DynamicToolRegistry(); + let capturedState: Record | undefined; + + registry.register("spy", async (_args, state) => { + capturedState = state; + return { tool_name: "spy", success: true, output: "ok" }; + }); + + const node = registry.asNode(); + const state: DynamicToolState = { + messages: [{ role: "user", content: "hi" }], + pendingTools: [{ name: "spy", args: {}, id: "tc-2" }], + }; + + await node(state); + expect(capturedState).toBeDefined(); + expect((capturedState as DynamicToolState).messages).toHaveLength(1); + }); + + it("catches handler exceptions and returns structured error", async () => { + const registry = new DynamicToolRegistry(); + + registry.register("boom", async () => { + throw new Error("something broke"); + }); + + const node = registry.asNode(); + const state: DynamicToolState = { + messages: [], + pendingTools: [{ name: "boom", args: {}, id: "tc-3" }], + }; + + const result = await node(state); + const partial = result as Partial; + expect(partial.messages).toHaveLength(1); + + const body: ToolResult = JSON.parse(partial.messages![0].content); + expect(body.success).toBe(false); + expect(body.error).toBe("something broke"); + expect(body.tool_name).toBe("boom"); + }); +}); diff --git a/src/__tests__/registry/register-unregister.test.ts b/src/__tests__/registry/register-unregister.test.ts new file mode 100644 index 0000000..7141f88 --- /dev/null +++ b/src/__tests__/registry/register-unregister.test.ts @@ -0,0 +1,126 @@ +import { describe, it, expect } from "vitest"; +import { DynamicToolRegistry } from "../../registry/index.js"; +import type { DynamicToolState } from "../../registry/index.js"; +import type { ToolResult } from "../../registry/index.js"; + +describe("DynamicToolRegistry — register / unregister lifecycle", () => { + it("list() reflects registrations", () => { + const registry = new DynamicToolRegistry(); + expect(registry.list()).toEqual([]); + + registry.register("a", async () => ({ + tool_name: "a", success: true, output: "ok", + })); + expect(registry.list()).toEqual(["a"]); + + registry.register("b", async () => ({ + tool_name: "b", success: true, output: "ok", + })); + expect(registry.list()).toEqual(["a", "b"]); + }); + + it("unregister removes the tool from list()", () => { + const registry = new DynamicToolRegistry(); + + registry.register("x", async () => ({ + tool_name: "x", success: true, output: "ok", + })); + registry.register("y", async () => ({ + tool_name: "y", success: true, output: "ok", + })); + + expect(registry.list()).toContain("x"); + expect(registry.unregister("x")).toBe(true); + expect(registry.list()).not.toContain("x"); + expect(registry.list()).toContain("y"); + }); + + it("unregister returns false for unknown tool", () => { + const registry = new DynamicToolRegistry(); + expect(registry.unregister("nonexistent")).toBe(false); + }); + + it("has() returns correct value before and after unregister", () => { + const registry = new DynamicToolRegistry(); + + registry.register("tool", async () => ({ + tool_name: "tool", success: true, output: "ok", + })); + + expect(registry.has("tool")).toBe(true); + expect(registry.has("other")).toBe(false); + + registry.unregister("tool"); + expect(registry.has("tool")).toBe(false); + }); + + it("calling unregistered tool via asNode returns error result", async () => { + const registry = new DynamicToolRegistry(); + + registry.register("temp", async () => ({ + tool_name: "temp", success: true, output: "ok", + })); + + // Verify it works first + const node = registry.asNode(); + const state: DynamicToolState = { + messages: [], + pendingTools: [{ name: "temp", args: {}, id: "tc-1" }], + }; + + const good = await node(state); + const goodBody: ToolResult = JSON.parse( + (good as Partial).messages![0].content + ); + expect(goodBody.success).toBe(true); + + // Now unregister and try again + registry.unregister("temp"); + + const bad = await node(state); + const badBody: ToolResult = JSON.parse( + (bad as Partial).messages![0].content + ); + expect(badBody.success).toBe(false); + expect(badBody.error).toContain("not registered"); + }); + + it("register() replaces existing handler (hot-swap)", async () => { + const registry = new DynamicToolRegistry(); + + registry.register("swap", async () => ({ + tool_name: "swap", success: true, output: "v1", + })); + + const node = registry.asNode(); + const state: DynamicToolState = { + messages: [], + pendingTools: [{ name: "swap", args: {}, id: "tc-1" }], + }; + + const r1 = await node(state); + const b1: ToolResult = JSON.parse( + (r1 as Partial).messages![0].content + ); + expect(b1.output).toBe("v1"); + + // Replace with v2 + registry.register("swap", async () => ({ + tool_name: "swap", success: true, output: "v2", + })); + + const r2 = await node(state); + const b2: ToolResult = JSON.parse( + (r2 as Partial).messages![0].content + ); + expect(b2.output).toBe("v2"); + }); + + it("register() returns this for chaining", () => { + const registry = new DynamicToolRegistry(); + const returned = registry.register("a", async () => ({ + tool_name: "a", success: true, output: "ok", + })); + expect(returned).toBe(registry); + }); +}); diff --git a/src/__tests__/registry/schema-output.test.ts b/src/__tests__/registry/schema-output.test.ts new file mode 100644 index 0000000..af01f75 --- /dev/null +++ b/src/__tests__/registry/schema-output.test.ts @@ -0,0 +1,114 @@ +import { describe, it, expect } from "vitest"; +import { DynamicToolRegistry } from "../../registry/index.js"; + +describe("DynamicToolRegistry — asSchema()", () => { + it("returns empty array when no tools registered", () => { + const registry = new DynamicToolRegistry(); + expect(registry.asSchema()).toEqual([]); + }); + + it("returns valid function-calling schema for registered tools", () => { + const registry = new DynamicToolRegistry(); + + registry.register("read_file", async () => ({ + tool_name: "read_file", success: true, output: "ok", + }), { + description: "Read a file", + parameters: { + type: "object", + properties: { path: { type: "string" } }, + required: ["path"], + }, + }); + + registry.register("write_file", async () => ({ + tool_name: "write_file", success: true, output: "ok", + }), { + description: "Write a file", + parameters: { + type: "object", + properties: { + path: { type: "string" }, + content: { type: "string" }, + }, + required: ["path", "content"], + }, + }); + + registry.register("search", async () => ({ + tool_name: "search", success: true, output: "ok", + }), { + description: "Search the web", + parameters: { + type: "object", + properties: { query: { type: "string" } }, + required: ["query"], + }, + }); + + const schemas = registry.asSchema(); + expect(schemas).toHaveLength(3); + + // All entries have the correct shape + for (const entry of schemas) { + expect(entry.type).toBe("function"); + expect(entry.function).toHaveProperty("name"); + expect(entry.function).toHaveProperty("description"); + expect(entry.function).toHaveProperty("parameters"); + expect(typeof entry.function.name).toBe("string"); + expect(typeof entry.function.description).toBe("string"); + } + + // Check specific tools + const names = schemas.map((s) => s.function.name); + expect(names).toEqual(["read_file", "write_file", "search"]); + + // Check parameter schemas preserved + const readFile = schemas.find((s) => s.function.name === "read_file")!; + expect(readFile.function.description).toBe("Read a file"); + expect(readFile.function.parameters).toEqual({ + type: "object", + properties: { path: { type: "string" } }, + required: ["path"], + }); + + const writeFile = schemas.find((s) => s.function.name === "write_file")!; + expect(writeFile.function.parameters).toHaveProperty("properties.path"); + expect(writeFile.function.parameters).toHaveProperty("properties.content"); + }); + + it("reflects live changes — unregistered tools disappear from schema", () => { + const registry = new DynamicToolRegistry(); + + registry.register("a", async () => ({ + tool_name: "a", success: true, output: "ok", + }), { description: "Tool A" }); + + registry.register("b", async () => ({ + tool_name: "b", success: true, output: "ok", + }), { description: "Tool B" }); + + expect(registry.asSchema()).toHaveLength(2); + + registry.unregister("a"); + const schemas = registry.asSchema(); + expect(schemas).toHaveLength(1); + expect(schemas[0].function.name).toBe("b"); + }); + + it("uses default description and parameters when opts omitted", () => { + const registry = new DynamicToolRegistry(); + + registry.register("bare", async () => ({ + tool_name: "bare", success: true, output: "ok", + })); + + const schemas = registry.asSchema(); + expect(schemas).toHaveLength(1); + expect(schemas[0].function.description).toBe(""); + expect(schemas[0].function.parameters).toEqual({ + type: "object", + properties: {}, + }); + }); +}); diff --git a/src/__tests__/registry/unknown-tool.test.ts b/src/__tests__/registry/unknown-tool.test.ts new file mode 100644 index 0000000..b5905c4 --- /dev/null +++ b/src/__tests__/registry/unknown-tool.test.ts @@ -0,0 +1,74 @@ +import { describe, it, expect } from "vitest"; +import { DynamicToolRegistry } from "../../registry/index.js"; +import type { DynamicToolState } from "../../registry/index.js"; +import type { ToolResult } from "../../registry/index.js"; + +describe("DynamicToolRegistry — unknown tool", () => { + it("returns structured error for unregistered tool name", async () => { + const registry = new DynamicToolRegistry(); + + registry.register("known_tool", async () => ({ + tool_name: "known_tool", + success: true, + output: "ok", + })); + + const node = registry.asNode(); + const state: DynamicToolState = { + messages: [], + pendingTools: [{ name: "nonexistent", args: {}, id: "tc-1" }], + }; + + const result = await node(state); + const partial = result as Partial; + + expect(partial.messages).toHaveLength(1); + + const msg = partial.messages![0]; + expect(msg.role).toBe("tool"); + expect(msg.name).toBe("nonexistent"); + expect(msg.tool_call_id).toBe("tc-1"); + + const body: ToolResult = JSON.parse(msg.content); + expect(body.success).toBe(false); + expect(body.error).toContain("nonexistent"); + expect(body.error).toContain("not registered"); + }); + + it("lists available tools in the error message", async () => { + const registry = new DynamicToolRegistry(); + + registry.register("alpha", async () => ({ + tool_name: "alpha", success: true, output: "ok", + })); + registry.register("beta", async () => ({ + tool_name: "beta", success: true, output: "ok", + })); + + const node = registry.asNode(); + const state: DynamicToolState = { + messages: [], + pendingTools: [{ name: "gamma", args: {}, id: "tc-2" }], + }; + + const result = await node(state); + const partial = result as Partial; + const body: ToolResult = JSON.parse(partial.messages![0].content); + + expect(body.error).toContain("alpha"); + expect(body.error).toContain("beta"); + }); + + it("does not throw — error stays in the message for LLM self-correction", async () => { + const registry = new DynamicToolRegistry(); + const node = registry.asNode(); + + const state: DynamicToolState = { + messages: [], + pendingTools: [{ name: "missing", args: {}, id: "tc-3" }], + }; + + // Should NOT throw — the error is returned as a structured result + await expect(node(state)).resolves.toBeDefined(); + }); +}); diff --git a/src/registry/README.md b/src/registry/README.md new file mode 100644 index 0000000..9c1b465 --- /dev/null +++ b/src/registry/README.md @@ -0,0 +1,125 @@ +# @oni.bot/core/registry — DynamicToolRegistry + +A hot-loadable extension point for `StateGraph`. The registry acts as a single +compiled node that internally dispatches to a mutable handler map. Extensions +can be added and removed at runtime without recompiling the graph. + +## Usage + +```typescript +import { StateGraph, START, END, appendList, lastValue } from "@oni.bot/core"; +import { DynamicToolRegistry } from "@oni.bot/core/registry"; +import type { DynamicToolState } from "@oni.bot/core/registry"; + +// 1. Create the registry +const registry = new DynamicToolRegistry(); + +// 2. Register built-in tools at startup +registry.register("read_file", async (args) => ({ + tool_name: "read_file", + success: true, + output: `Contents of ${args.path}`, +}), { + description: "Read a file from disk", + parameters: { + type: "object", + properties: { path: { type: "string" } }, + required: ["path"], + }, +}); + +registry.register("write_file", async (args) => ({ + tool_name: "write_file", + success: true, + output: `Wrote ${args.path}`, +}), { + description: "Write content to a file", + parameters: { + type: "object", + properties: { + path: { type: "string" }, + content: { type: "string" }, + }, + required: ["path", "content"], + }, +}); + +// 3. Wire into a StateGraph +type State = DynamicToolState; + +const graph = new StateGraph({ + channels: { + messages: appendList(() => []), + pendingTools: lastValue(() => []), + }, +}) + .addNode("tool_executor", registry.asNode()) + .addConditionalEdges(START, (state) => { + return state.pendingTools.length > 0 ? "tool_executor" : "__end__"; + }, { tool_executor: "tool_executor", [END]: END }) + .addEdge("tool_executor", END); + +const app = graph.compile(); + +// 4. Register a new extension at runtime — no recompile +registry.register("search_web", async (args) => ({ + tool_name: "search_web", + success: true, + output: `Results for "${args.query}"`, +}), { + description: "Search the web", + parameters: { + type: "object", + properties: { query: { type: "string" } }, + required: ["query"], + }, +}); + +// 5. Unregister — also live +registry.unregister("search_web"); + +// 6. Pass registered tool schemas to the LLM +const toolSchemas = registry.asSchema(); +// toolSchemas is an array of { type: "function", function: { name, description, parameters } } +``` + +## API + +### `new DynamicToolRegistry()` + +Creates an empty registry. + +### `.register(name, handler, opts?)` + +Register a tool handler. Returns `this` for chaining. If a tool with the same +name is already registered, the handler is replaced (hot-swap). + +- `name` — unique tool name +- `handler` — `(args, state) => Promise` +- `opts.description` — tool description for schema generation +- `opts.parameters` — JSON Schema for the tool's arguments + +### `.unregister(name)` + +Remove a tool. Returns `true` if it existed. + +### `.list()` + +Returns an array of currently registered tool names. + +### `.has(name)` + +Returns `true` if the named tool is registered. + +### `.asNode()` + +Returns a `NodeFn` suitable for `graph.addNode()`. The node: +- Reads `state.pendingTools[0]` +- Dispatches to the matching handler +- Returns the result as a tool message in `messages` +- Returns a structured error (never throws) if the tool is not found + +### `.asSchema()` + +Returns an array of tool definitions in OpenAI function-calling format, suitable +for passing to an LLM node as `tools`. diff --git a/src/registry/dynamic-tool-registry.ts b/src/registry/dynamic-tool-registry.ts new file mode 100644 index 0000000..655e42f --- /dev/null +++ b/src/registry/dynamic-tool-registry.ts @@ -0,0 +1,161 @@ +// ============================================================ +// @oni.bot/core/registry — DynamicToolRegistry +// ============================================================ +// A mutable handler map that sits behind a single compiled +// StateGraph node. Extensions register/unregister at runtime; +// the graph topology never changes. +// ============================================================ + +import type { NodeFn } from "../types.js"; +import type { ONIMessage } from "../graph.js"; +import type { ToolHandler, ToolResult } from "./types.js"; + +// ---------------------------------------------------------------- +// State contract — the node reads pendingTools from state +// ---------------------------------------------------------------- + +export interface DynamicToolState { + messages: ONIMessage[]; + pendingTools: Array<{ name: string; args: Record; id: string }>; +} + +// ---------------------------------------------------------------- +// Tool metadata for schema generation +// ---------------------------------------------------------------- + +export interface ToolRegistration { + name: string; + description: string; + parameters: Record; + handler: ToolHandler; +} + +// ---------------------------------------------------------------- +// DynamicToolRegistry +// ---------------------------------------------------------------- + +export class DynamicToolRegistry { + private handlers = new Map(); + + /** + * Register a tool handler. If a handler with the same name already exists, + * it is replaced silently (hot-swap). + */ + register( + name: string, + handler: ToolHandler, + opts: { description?: string; parameters?: Record } = {} + ): this { + this.handlers.set(name, { + name, + description: opts.description ?? "", + parameters: opts.parameters ?? { type: "object", properties: {} }, + handler, + }); + return this; + } + + /** Remove a tool by name. Returns true if it existed. */ + unregister(name: string): boolean { + return this.handlers.delete(name); + } + + /** List currently registered tool names. */ + list(): string[] { + return [...this.handlers.keys()]; + } + + /** Check if a tool is registered. */ + has(name: string): boolean { + return this.handlers.has(name); + } + + /** + * Returns a StateGraph-compatible node function. + * + * The node reads `state.pendingTools[0]`, dispatches to the matching + * handler, and returns the result merged into messages. If the tool + * is not found it returns a structured error (never throws). + */ + asNode(): NodeFn { + return async (state) => { + const pending = state.pendingTools?.[0]; + + if (!pending) { + return {}; + } + + const registration = this.handlers.get(pending.name); + + if (!registration) { + const errorResult: ToolResult = { + tool_name: pending.name, + success: false, + output: "", + error: `Tool "${pending.name}" is not registered. Available tools: ${this.list().join(", ") || "(none)"}`, + }; + + return { + messages: [ + { + role: "tool" as const, + content: JSON.stringify(errorResult), + name: pending.name, + tool_call_id: pending.id, + }, + ], + }; + } + + try { + const result = await registration.handler(pending.args, state); + return { + messages: [ + { + role: "tool" as const, + content: JSON.stringify(result), + name: pending.name, + tool_call_id: pending.id, + }, + ], + }; + } catch (err) { + const errorResult: ToolResult = { + tool_name: pending.name, + success: false, + output: "", + error: err instanceof Error ? err.message : String(err), + }; + + return { + messages: [ + { + role: "tool" as const, + content: JSON.stringify(errorResult), + name: pending.name, + tool_call_id: pending.id, + }, + ], + }; + } + }; + } + + /** + * Returns a JSON schema array describing all registered tools, + * suitable for passing to an LLM node as available tool definitions. + */ + asSchema(): Array<{ + type: "function"; + function: { name: string; description: string; parameters: Record }; + }> { + return [...this.handlers.values()].map((reg) => ({ + type: "function" as const, + function: { + name: reg.name, + description: reg.description, + parameters: reg.parameters, + }, + })); + } +} diff --git a/src/registry/index.ts b/src/registry/index.ts new file mode 100644 index 0000000..e08f151 --- /dev/null +++ b/src/registry/index.ts @@ -0,0 +1,3 @@ +export { DynamicToolRegistry } from "./dynamic-tool-registry.js"; +export type { DynamicToolState, ToolRegistration } from "./dynamic-tool-registry.js"; +export type { ToolHandler, ToolResult } from "./types.js"; diff --git a/src/registry/types.ts b/src/registry/types.ts new file mode 100644 index 0000000..685796f --- /dev/null +++ b/src/registry/types.ts @@ -0,0 +1,16 @@ +// ============================================================ +// @oni.bot/core/registry — Types +// ============================================================ + +// eslint-disable-next-line @typescript-eslint/no-explicit-any +export type ToolHandler = ( + args: Record, + state: any // SAFE: state shape varies by graph — callers narrow via DynamicToolState +) => Promise; + +export type ToolResult = { + tool_name: string; + success: boolean; + output: string; + error?: string; +};