From 3f618b2cf3d9fef65f22b45587e173872d823cd3 Mon Sep 17 00:00:00 2001 From: Lucas Winther Date: Sun, 6 Sep 2026 18:52:22 +0200 Subject: [PATCH] feat(tools): add execution wrapper support to tool registration --- src/tools/common.ts | 40 ++++++++++++-- src/tools/index.ts | 28 +++++++++- tests/tools.test.ts | 127 ++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 190 insertions(+), 5 deletions(-) diff --git a/src/tools/common.ts b/src/tools/common.ts index f2e29b8..7054ac3 100644 --- a/src/tools/common.ts +++ b/src/tools/common.ts @@ -65,20 +65,54 @@ export interface ToolDefinition< execute: (args: TArgs) => Promise; } +/** + * Options for registering an individual tool on an McpServer instance. + */ +export interface RegisterToolOptions< + TShape extends ZodRawShape = ZodRawShape, + TArgs = Record, +> { + /** + * Optional wrapper function to intercept tool execution, enabling + * context injection (runWithContext) or execution guards (e.g. read-only checks). + */ + wrapExecute?: ( + fn: () => Promise, + tool: ToolDefinition, + args: TArgs + ) => Promise; +} + /** * Registers a ToolDefinition instance on an McpServer instance. + * + * @param server The McpServer instance to register the tool on. + * @param tool The ToolDefinition describing the tool metadata, parameters, and handler. + * @param options Optional registration options, such as an execution wrapper. */ export function registerTool< TShape extends ZodRawShape = ZodRawShape, TArgs = Record, ->(server: McpServer, tool: ToolDefinition): void { +>( + server: McpServer, + tool: ToolDefinition, + options?: RegisterToolOptions +): void { + const executeFn = (args: TArgs) => { + if (options?.wrapExecute) { + return options.wrapExecute(() => tool.execute(args), tool, args); + } + return tool.execute(args); + }; + if (tool.parameters) { server.tool(tool.name, tool.description, tool.parameters, async (args) => { - return tool.execute(args as unknown as TArgs); + return executeFn(args as unknown as TArgs); }); } else { server.tool(tool.name, tool.description, async () => { - return tool.execute({} as unknown as TArgs); + return executeFn({} as unknown as TArgs); }); } } + diff --git a/src/tools/index.ts b/src/tools/index.ts index ba75ff3..a86d84f 100644 --- a/src/tools/index.ts +++ b/src/tools/index.ts @@ -3,7 +3,12 @@ */ import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; -import { registerTool, type ToolDefinition } from "./common"; +import { + registerTool, + type ToolDefinition, + type RegisterToolOptions, + type McpToolResponse, +} from "./common"; import { projectTools } from "./projects"; import { workPackageTools } from "./work-packages"; import { queryTools } from "./queries"; @@ -31,11 +36,17 @@ export const allTools: AnyToolDefinition[] = [ export interface RegisterToolsOptions { readOnly?: boolean; tools?: AnyToolDefinition[]; + wrapExecute?: ( + fn: () => Promise, + tool: AnyToolDefinition, + args: Record + ) => Promise; } /** * Registers all OpenProject tools with the given McpServer instance. * If options.readOnly is true, tools not marked as readOnly are skipped. + * If options.wrapExecute is provided, tool execution is intercepted through the wrapper. */ export function registerAllTools( server: McpServer, @@ -49,6 +60,19 @@ export function registerAllTools( if (isReadOnly && !tool.readOnly) { continue; } - registerTool(server, tool as unknown as ToolDefinition); + + const toolOptions: RegisterToolOptions | undefined = options?.wrapExecute + ? { + wrapExecute: (fn, t, args) => + options.wrapExecute!( + fn, + t as unknown as AnyToolDefinition, + args as Record + ), + } + : undefined; + + registerTool(server, tool as unknown as ToolDefinition, toolOptions); } } + diff --git a/tests/tools.test.ts b/tests/tools.test.ts index 3ef605e..25ad1eb 100644 --- a/tests/tools.test.ts +++ b/tests/tools.test.ts @@ -10,6 +10,7 @@ import { formatToolError, registerTool, type ToolDefinition, + type RegisterToolOptions, } from "../src/tools/common"; import { listProjectsShape, @@ -1026,6 +1027,132 @@ describe("Tool Registry", () => { }); }); +describe("Tool Execution Wrapper", () => { + test("registerTool passes execution through wrapExecute when provided", async () => { + const server = new McpServer({ name: "test", version: "1" }); + let wrapperCalled = false; + const dummyTool: ToolDefinition = { + name: "wrapped_tool", + description: "A wrapped tool", + readOnly: true, + execute: async () => ({ content: [{ type: "text", text: "success" }] }), + }; + + registerTool(server, dummyTool, { + wrapExecute: async (fn) => { + wrapperCalled = true; + return fn(); + }, + }); + + const registered = ( + server as unknown as { + _registeredTools: Record< + string, + { + handler: (args: Record) => Promise<{ + content: Array<{ type: string; text: string }>; + isError?: boolean; + }>; + } + >; + } + )._registeredTools["wrapped_tool"]; + const result = await registered.handler({}); + expect(wrapperCalled).toBe(true); + expect(result.content[0]?.text).toBe("success"); + }); + + test("registerTool passes tool definition and arguments to wrapExecute", async () => { + const server = new McpServer({ name: "test", version: "1" }); + let capturedToolName = ""; + let capturedArgs: Record = {}; + + const dummyToolWithParams: ToolDefinition<{ id: z.ZodNumber }, { id: number }> = { + name: "param_wrapped_tool", + description: "Param wrapped tool", + parameters: { id: z.number() }, + readOnly: true, + execute: async (args: { id: number }) => ({ + content: [{ type: "text", text: `id: ${args.id}` }], + }), + }; + + registerTool(server, dummyToolWithParams, { + wrapExecute: async (fn, tool, args) => { + capturedToolName = tool.name; + capturedArgs = args as Record; + return fn(); + }, + }); + + const registered = ( + server as unknown as { + _registeredTools: Record< + string, + { + handler: (args: Record) => Promise<{ + content: Array<{ type: string; text: string }>; + isError?: boolean; + }>; + } + >; + } + )._registeredTools["param_wrapped_tool"]; + + const result = await registered.handler({ id: 123 }); + expect(capturedToolName).toBe("param_wrapped_tool"); + expect(capturedArgs).toEqual({ id: 123 }); + expect(result.content[0]?.text).toBe("id: 123"); + }); + + test("registerAllTools forwards wrapExecute to registered tools", async () => { + const server = new McpServer({ name: "test", version: "1" }); + const interceptedTools: string[] = []; + + registerAllTools(server, { + wrapExecute: async (fn, tool) => { + interceptedTools.push(tool.name); + return fn(); + }, + }); + + const registeredTools = ( + server as unknown as { + _registeredTools: Record< + string, + { + handler: (args: Record) => Promise<{ + content: Array<{ type: string; text: string }>; + isError?: boolean; + }>; + } + >; + } + )._registeredTools; + + const listProjects = registeredTools["openproject_list_projects"]; + expect(listProjects).toBeDefined(); + + const dummyClient = { + get: async () => ({ + _embedded: { elements: [] }, + total: 0, + count: 0, + pageSize: 20, + offset: 1, + }), + } as unknown as OpenProjectClient; + + await runWithContext({ client: dummyClient, isReadOnly: false }, () => + listProjects.handler({}) + ); + + expect(interceptedTools).toContain("openproject_list_projects"); + }); +}); + + describe("Live Container Integration (All 10 MCP Tools)", () => { const baseUrl = process.env.OPENPROJECT_BASE_URL || "http://localhost:8080"; const apiKey = process.env.OPENPROJECT_API_KEY || "";