Add ToolAnnotations support to McpServer.tool() method

bhosmer-ant · bhosmer-ant · commit 0096be432f68 · 2025-04-22T15:24:32.000-04:00
diff --git a/src/examples/server/simpleStreamableHttp.ts b/src/examples/server/simpleStreamableHttp.ts
@@ -13,17 +13,88 @@ const getServer = () => {
     version: '1.0.0',
   }, { capabilities: { logging: {} } });
 
-  // Register a simple tool that returns a greeting
-  server.tool(
-    'greet',
-    'A simple greeting tool',
-    {
-      name: z.string().describe('Name to greet'),
-    },
-    async ({ name }): Promise<CallToolResult> => {
-      return {
-        content: [
-          {
+// Register a simple tool that returns a greeting
+server.tool(
+  'greet',
+  'A simple greeting tool',
+  {
+    name: z.string().describe('Name to greet'),
+  },
+  async ({ name }): Promise<CallToolResult> => {
+    return {
+      content: [
+        {
+          type: 'text',
+          text: `Hello, ${name}!`,
+        },
+      ],
+    };
+  },
+  {
+    title: 'Greeting Tool',
+    readOnlyHint: true,
+    openWorldHint: false
+  }
+);
+
+// Register a tool that sends multiple greetings with notifications
+server.tool(
+  'multi-greet',
+  'A tool that sends different greetings with delays between them',
+  {
+    name: z.string().describe('Name to greet'),
+  },
+  async ({ name }, { sendNotification }): Promise<CallToolResult> => {
+    const sleep = (ms: number) => new Promise(resolve => setTimeout(resolve, ms));
+
+    await sendNotification({
+      method: "notifications/message",
+      params: { level: "debug", data: `Starting multi-greet for ${name}` }
+    });
+
+    await sleep(1000); // Wait 1 second before first greeting
+
+    await sendNotification({
+      method: "notifications/message",
+      params: { level: "info", data: `Sending first greeting to ${name}` }
+    });
+
+    await sleep(1000); // Wait another second before second greeting
+
+    await sendNotification({
+      method: "notifications/message",
+      params: { level: "info", data: `Sending second greeting to ${name}` }
+    });
+
+    return {
+      content: [
+        {
+          type: 'text',
+          text: `Good morning, ${name}!`,
+        }
+      ],
+    };
+  },
+  {
+    title: 'Multiple Greeting Tool', 
+    readOnlyHint: true,
+    openWorldHint: false
+  }
+);
+
+// Register a simple prompt
+server.prompt(
+  'greeting-template',
+  'A simple greeting prompt template',
+  {
+    name: z.string().describe('Name to include in greeting'),
+  },
+  async ({ name }): Promise<GetPromptResult> => {
+    return {
+      messages: [
+        {
+          role: 'user',
+          content: {
             type: 'text',
             text: `Hello, ${name}!`,
           },
diff --git a/src/server/mcp.test.ts b/src/server/mcp.test.ts
@@ -494,6 +494,133 @@ describe("tool()", () => {
     expect(result.tools[0].name).toBe("test");
     expect(result.tools[0].description).toBe("Test description");
   });
+  
+  test("should register tool with annotations", async () => {
+    const mcpServer = new McpServer({
+      name: "test server",
+      version: "1.0",
+    });
+    const client = new Client({
+      name: "test client",
+      version: "1.0",
+    });
+
+    mcpServer.tool("test", { title: "Test Tool", readOnlyHint: true }, async () => ({
+      content: [
+        {
+          type: "text",
+          text: "Test response",
+        },
+      ],
+    }));
+
+    const [clientTransport, serverTransport] =
+      InMemoryTransport.createLinkedPair();
+
+    await Promise.all([
+      client.connect(clientTransport),
+      mcpServer.server.connect(serverTransport),
+    ]);
+
+    const result = await client.request(
+      {
+        method: "tools/list",
+      },
+      ListToolsResultSchema,
+    );
+
+    expect(result.tools).toHaveLength(1);
+    expect(result.tools[0].name).toBe("test");
+    expect(result.tools[0].annotations).toEqual({ title: "Test Tool", readOnlyHint: true });
+  });
+  
+  test("should register tool with params and annotations", async () => {
+    const mcpServer = new McpServer({
+      name: "test server",
+      version: "1.0",
+    });
+    const client = new Client({
+      name: "test client",
+      version: "1.0",
+    });
+
+    mcpServer.tool(
+      "test", 
+      { name: z.string() },
+      { title: "Test Tool", readOnlyHint: true },
+      async ({ name }) => ({
+        content: [{ type: "text", text: `Hello, ${name}!` }]
+      })
+    );
+
+    const [clientTransport, serverTransport] =
+      InMemoryTransport.createLinkedPair();
+
+    await Promise.all([
+      client.connect(clientTransport),
+      mcpServer.server.connect(serverTransport),
+    ]);
+
+    const result = await client.request(
+      { method: "tools/list" },
+      ListToolsResultSchema,
+    );
+
+    expect(result.tools).toHaveLength(1);
+    expect(result.tools[0].name).toBe("test");
+    expect(result.tools[0].inputSchema).toMatchObject({
+      type: "object",
+      properties: { name: { type: "string" } }
+    });
+    expect(result.tools[0].annotations).toEqual({ title: "Test Tool", readOnlyHint: true });
+  });
+  
+  test("should register tool with description, params, and annotations", async () => {
+    const mcpServer = new McpServer({
+      name: "test server",
+      version: "1.0",
+    });
+    const client = new Client({
+      name: "test client",
+      version: "1.0",
+    });
+
+    mcpServer.tool(
+      "test", 
+      "A tool with everything",
+      { name: z.string() },
+      { title: "Complete Test Tool", readOnlyHint: true, openWorldHint: false },
+      async ({ name }) => ({
+        content: [{ type: "text", text: `Hello, ${name}!` }]
+      })
+    );
+
+    const [clientTransport, serverTransport] =
+      InMemoryTransport.createLinkedPair();
+
+    await Promise.all([
+      client.connect(clientTransport),
+      mcpServer.server.connect(serverTransport),
+    ]);
+
+    const result = await client.request(
+      { method: "tools/list" },
+      ListToolsResultSchema,
+    );
+
+    expect(result.tools).toHaveLength(1);
+    expect(result.tools[0].name).toBe("test");
+    expect(result.tools[0].description).toBe("A tool with everything");
+    expect(result.tools[0].inputSchema).toMatchObject({
+      type: "object",
+      properties: { name: { type: "string" } }
+    });
+    expect(result.tools[0].annotations).toEqual({ 
+      title: "Complete Test Tool", 
+      readOnlyHint: true,
+      openWorldHint: false
+    });
+  });
 
   test("should validate tool args", async () => {
     const mcpServer = new McpServer({
diff --git a/src/server/mcp.ts b/src/server/mcp.ts
@@ -39,6 +39,7 @@ import {
   ReadResourceResult,
   ServerRequest,
   ServerNotification,
+  ToolAnnotations,
 } from "../types.js";
 import { Completable, CompletableDef } from "./completable.js";
 import { UriTemplate, Variables } from "../shared/uriTemplate.js";
@@ -118,6 +119,7 @@ export class McpServer {
                     strictUnions: true,
                   }) as Tool["inputSchema"])
                 : EMPTY_OBJECT_JSON_SCHEMA,
+              annotations: tool.annotations,
             };
           },
         ),
@@ -605,44 +607,103 @@ export class McpServer {
   tool(name: string, description: string, cb: ToolCallback): RegisteredTool;
 
   /**
-   * Registers a tool `name` accepting the given arguments, which must be an object containing named properties associated with Zod schemas. When the client calls it, the function will be run with the parsed and validated arguments.
+   * Registers a tool taking either a parameter schema for validation or annotations for additional metadata.
+   * This unified overload handles both `tool(name, paramsSchema, cb)` and `tool(name, annotations, cb)` cases.
+   * 
+   * Note: We use a union type for the second parameter because TypeScript cannot reliably disambiguate
+   * between ToolAnnotations and ZodRawShape during overload resolution, as both are plain object types.
    */
   tool<Args extends ZodRawShape>(
     name: string,
-    paramsSchema: Args,
+    paramsSchemaOrAnnotations: Args | ToolAnnotations,
     cb: ToolCallback<Args>,
   ): RegisteredTool;
 
   /**
-   * Registers a tool `name` (with a description) accepting the given arguments, which must be an object containing named properties associated with Zod schemas. When the client calls it, the function will be run with the parsed and validated arguments.
+   * Registers a tool `name` (with a description) taking either parameter schema or annotations.
+   * This unified overload handles both `tool(name, description, paramsSchema, cb)` and 
+   * `tool(name, description, annotations, cb)` cases.
+   * 
+   * Note: We use a union type for the third parameter because TypeScript cannot reliably disambiguate
+   * between ToolAnnotations and ZodRawShape during overload resolution, as both are plain object types.
+   */
+  tool<Args extends ZodRawShape>(
+    name: string,
+    description: string,
+    paramsSchemaOrAnnotations: Args | ToolAnnotations,
+    cb: ToolCallback<Args>,
+  ): RegisteredTool;
+  
+  /**
+   * Registers a tool with both parameter schema and annotations.
+   */
+  tool<Args extends ZodRawShape>(
+    name: string,
+    paramsSchema: Args,
+    annotations: ToolAnnotations,
+    cb: ToolCallback<Args>,
+  ): RegisteredTool;
+  
+  /**
+   * Registers a tool with description, parameter schema, and annotations.
    */
   tool<Args extends ZodRawShape>(
     name: string,
     description: string,
     paramsSchema: Args,
+    annotations: ToolAnnotations,
     cb: ToolCallback<Args>,
   ): RegisteredTool;
 
   tool(name: string, ...rest: unknown[]): RegisteredTool {
     if (this._registeredTools[name]) {
       throw new Error(`Tool ${name} is already registered`);
     }
+    
+    // Helper to check if an object is a Zod schema (ZodRawShape)
+    const isZodRawShape = (obj: unknown): obj is ZodRawShape => {
+      if (typeof obj !== "object" || obj === null) return false;
+      // Check that at least one property is a ZodType instance
+      return Object.values(obj as object).some(v => v instanceof ZodType);
+    };
 
     let description: string | undefined;
     if (typeof rest[0] === "string") {
       description = rest.shift() as string;
     }
 
     let paramsSchema: ZodRawShape | undefined;
+    let annotations: ToolAnnotations | undefined;
+    
+    // Handle the different overload combinations
     if (rest.length > 1) {
-      paramsSchema = rest.shift() as ZodRawShape;
+      // We have at least two more args before the callback
+      const firstArg = rest[0];
+      
+      if (isZodRawShape(firstArg)) {
+        // We have a params schema as the first arg
+        paramsSchema = rest.shift() as ZodRawShape;
+        
+        // Check if the next arg is potentially annotations  
+        if (rest.length > 1 && typeof rest[0] === "object" && rest[0] !== null && !(isZodRawShape(rest[0]))) {
+          // Case: tool(name, paramsSchema, annotations, cb)
+          // Or: tool(name, description, paramsSchema, annotations, cb)
+          annotations = rest.shift() as ToolAnnotations;
+        }
+      } else if (typeof firstArg === "object" && firstArg !== null) {
+        // Not a ZodRawShape, so must be annotations in this position
+        // Case: tool(name, annotations, cb)
+        // Or: tool(name, description, annotations, cb)
+        annotations = rest.shift() as ToolAnnotations;
+      }
     }
 
     const cb = rest[0] as ToolCallback<ZodRawShape | undefined>;
     const registeredTool: RegisteredTool = {
       description,
       inputSchema:
         paramsSchema === undefined ? undefined : z.object(paramsSchema),
+      annotations,
       callback: cb,
       enabled: true,
       disable: () => registeredTool.update({ enabled: false }),
@@ -656,6 +717,7 @@ export class McpServer {
         if (typeof updates.description !== "undefined") registeredTool.description = updates.description
         if (typeof updates.paramsSchema !== "undefined") registeredTool.inputSchema = z.object(updates.paramsSchema)
         if (typeof updates.callback !== "undefined") registeredTool.callback = updates.callback
+        if (typeof updates.annotations !== "undefined") registeredTool.annotations = updates.annotations
         if (typeof updates.enabled !== "undefined") registeredTool.enabled = updates.enabled
         this.sendToolListChanged()
       },
@@ -853,11 +915,12 @@ export type ToolCallback<Args extends undefined | ZodRawShape = undefined> =
 export type RegisteredTool = {
   description?: string;
   inputSchema?: AnyZodObject;
+  annotations?: ToolAnnotations;
   callback: ToolCallback<undefined | ZodRawShape>;
   enabled: boolean;
   enable(): void;
   disable(): void;
-  update<Args extends ZodRawShape>(updates: { name?: string | null, description?: string, paramsSchema?: Args, callback?: ToolCallback<Args>, enabled?: boolean }): void
+  update<Args extends ZodRawShape>(updates: { name?: string | null, description?: string, paramsSchema?: Args, callback?: ToolCallback<Args>, annotations?: ToolAnnotations, enabled?: boolean }): void
   remove(): void
 };
 
diff --git a/src/types.ts b/src/types.ts
