feat(tools): add execution wrapper support to tool registration

src/tools/common.ts

@@ -65,20 +65,54 @@ export interface ToolDefinition<
   execute: (args: TArgs) => Promise<McpToolResponse>;
 }
 
+/**
+ * Options for registering an individual tool on an McpServer instance.
+ */
+export interface RegisterToolOptions<
+  TShape extends ZodRawShape = ZodRawShape,
+  TArgs = Record<string, unknown>,
+> {
+  /**
+   * Optional wrapper function to intercept tool execution, enabling
+   * context injection (runWithContext) or execution guards (e.g. read-only checks).
+   */
+  wrapExecute?: (
+    fn: () => Promise<McpToolResponse>,
+    tool: ToolDefinition<TShape, TArgs>,
+    args: TArgs
+  ) => Promise<McpToolResponse>;
+}
+
 /**
  * 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<string, unknown>,
->(server: McpServer, tool: ToolDefinition<TShape, TArgs>): void {
+>(
+  server: McpServer,
+  tool: ToolDefinition<TShape, TArgs>,
+  options?: RegisterToolOptions<TShape, TArgs>
+): 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);
     });
   }
 }
+

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<McpToolResponse>,
+    tool: AnyToolDefinition,
+    args: Record<string, unknown>
+  ) => Promise<McpToolResponse>;
 }
 
 /**
  * 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<string, unknown>
+            ),
+        }
+      : undefined;
+
+    registerTool(server, tool as unknown as ToolDefinition, toolOptions);
   }
 }
+

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<string, unknown>) => 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<string, unknown> = {};
+
+    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<string, unknown>;
+        return fn();
+      },
+    });
+
+    const registered = (
+      server as unknown as {
+        _registeredTools: Record<
+          string,
+          {
+            handler: (args: Record<string, unknown>) => 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<string, unknown>) => 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 || "";