Allow users to add middlewares on toollist and toolcall

Signed-off-by: Marcos Candeia <[email protected]>

Author: mcandeia

deno.json

@@ -1,6 +1,6 @@
 {
   "name": "@deco/mcp",
-  "version": "0.2.6",
+  "version": "0.2.7",
   "exports": "./mod.ts",
   "tasks": {
     "check": "deno fmt && deno lint && deno check mod.ts"
@@ -9,6 +9,7 @@
     "@deco/deco": "jsr:@deco/deco@^1.110.4",
     "@modelcontextprotocol/sdk": "npm:@modelcontextprotocol/sdk@^1.6.1",
     "@hono/hono": "jsr:@hono/hono@^4.5.4",
-    "@std/http": "jsr:@std/http@^1.0.13"
+    "@std/http": "jsr:@std/http@^1.0.13",
+    "zod": "npm:zod@^3.24.2"
   }
 }

deno.lock

@@ -0,0 +1,32 @@
+// deno-lint-ignore-file no-explicit-any
+
+export interface RequestMiddlewareContext<T = any> {
+  next?(): Promise<T>;
+}
+export type RequestMiddleware<
+  TRequest = any,
+  TResponse = any,
+> = (request: TRequest, next?: () => Promise<TResponse>) => Promise<TResponse>;
+
+export const compose = <
+  TRequest,
+  TResponse,
+>(
+  ...middlewares: RequestMiddleware<TRequest, TResponse>[]
+): RequestMiddleware<TRequest, TResponse> => {
+  const last = middlewares[middlewares.length - 1];
+  return function composedResolver(request: TRequest) {
+    const dispatch = (
+      i: number,
+    ): Promise<TResponse> => {
+      const middleware = middlewares[i];
+      if (!middleware) {
+        return last(request);
+      }
+      const next = () => dispatch(i + 1);
+      return middleware(request, next);
+    };
+
+    return dispatch(0);
+  };
+};

mcp/middleware.ts

@@ -15,7 +15,8 @@ import { SSEServerTransport } from "./sse.ts";
 import { HttpServerTransport } from "./http.ts";
 import { dereferenceSchema } from "./utils.ts";
 import { WebSocketServerTransport } from "./websocket.ts";
-
+import { compose, type RequestMiddleware } from "./middleware.ts";
+import type { z } from "zod";
 const idFromDefinition = (definition: string) => {
   const [_, __, id] = definition.split("/");
   return id;
@@ -50,6 +51,10 @@ export interface Options<TManifest extends AppManifest> {
   exclude?: Array<keyof (TManifest["actions"] & TManifest["loaders"])>;
   blocks?: Array<keyof TManifest>;
   basePath?: string;
+  middlewares?: {
+    listTools?: ListToolsMiddleware[];
+    callTool?: CallToolMiddleware[];
+  };
 }
 
 interface RootSchema extends JSONSchema7 {
@@ -183,26 +188,45 @@ export const getTools = <TManifest extends AppManifest>(
   return tools.filter((tool) => tool !== undefined);
 };
 
+export interface ListToolsResult {
+  tools: Tool[];
+}
+
+export type ListToolsMiddleware = RequestMiddleware<
+  z.infer<typeof ListToolsRequestSchema>,
+  ListToolsResult
+>;
+
+export type CallToolMiddleware = RequestMiddleware<
+  z.infer<typeof CallToolRequestSchema>,
+  { content: { type: "text"; text: string }[] }
+>;
+
 function registerTools<TManifest extends AppManifest>(
   mcp: McpServer,
   deco: Deco<TManifest>,
   options?: Options<TManifest>,
 ) {
   // Add map to store slugified names to original names
   let toolNames: null | Map<string, string> = null;
-  const loadTools = async () => {
+  const loadTools: ListToolsMiddleware = async (): Promise<ListToolsResult> => {
     toolNames ??= new Map<string, string>();
     const meta = await deco.meta().then((v) => v?.value);
     if (!meta) return { tools: [] };
     const schemas = meta.schema;
     return { tools: getTools(toolNames, schemas, options) };
   };
 
-  mcp.server.setRequestHandler(ListToolsRequestSchema, async () => {
-    return await loadTools();
+  const listTools: ListToolsMiddleware = compose(
+    ...(options?.middlewares?.listTools ?? []),
+    loadTools,
+  );
+
+  mcp.server.setRequestHandler(ListToolsRequestSchema, async (request) => {
+    return await listTools(request);
   });
 
-  mcp.server.setRequestHandler(CallToolRequestSchema, async (req) => {
+  const invokeTool = async (req: z.infer<typeof CallToolRequestSchema>) => {
     IS_DEBUG && console.log(req);
     try {
       const state = await deco.prepareState({
@@ -213,7 +237,7 @@ function registerTools<TManifest extends AppManifest>(
       });
       // Use the original name from the map when invoking
       if (!toolNames) {
-        await loadTools();
+        await loadTools({ request: {} });
       }
       const originalName = toolNames!.get(req.params.name);
       if (!originalName) {
@@ -227,12 +251,21 @@ function registerTools<TManifest extends AppManifest>(
         state,
       );
       return {
-        content: [{ type: "text", text: JSON.stringify(result) }],
+        content: [{ type: "text" as const, text: JSON.stringify(result) }],
       };
     } catch (err) {
       console.error(err);
       throw err;
     }
+  };
+
+  const callToolMiddleware: CallToolMiddleware = compose(
+    ...(options?.middlewares?.callTool ?? []),
+    invokeTool,
+  );
+
+  mcp.server.setRequestHandler(CallToolRequestSchema, async (req) => {
+    return await callToolMiddleware(req);
   });
 }
 

mcp/server.ts

@@ -94,6 +94,10 @@ export function dereferenceSchema(
     );
   }
 
+  if ("allOf" in result && !Array.isArray(result.allOf)) {
+    delete result.allOf;
+  }
+
   return result;
 }
 

mcp/utils.ts

@@ -1,2 +1,9 @@
-export { getTools, mcpServer, type Tool } from "./mcp/server.ts";
+export {
+  type CallToolMiddleware,
+  getTools,
+  type ListToolsMiddleware,
+  mcpServer,
+  type Options,
+  type Tool,
+} from "./mcp/server.ts";
 export { HttpClientTransport } from "./mcp/http-client.ts";