Skip to content

Add ToolAnnotations support to McpServer.tool() method #309

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Apr 22, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion src/examples/server/simpleStreamableHttp.ts
Original file line number Diff line number Diff line change
Expand Up @@ -32,13 +32,18 @@ const getServer = () => {
}
);

// Register a tool that sends multiple greetings with notifications
// Register a tool that sends multiple greetings with notifications (with annotations)
server.tool(
'multi-greet',
'A tool that sends different greetings with delays between them',
{
name: z.string().describe('Name to greet'),
},
{
title: 'Multiple Greeting Tool',
readOnlyHint: true,
openWorldHint: false
},
async ({ name }, { sendNotification }): Promise<CallToolResult> => {
const sleep = (ms: number) => new Promise(resolve => setTimeout(resolve, ms));

Expand Down
129 changes: 128 additions & 1 deletion src/server/mcp.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -404,7 +404,7 @@ describe("tool()", () => {
])
});

test("should register tool with args schema", async () => {
test("should register tool with params", async () => {
const mcpServer = new McpServer({
name: "test server",
version: "1.0",
Expand Down Expand Up @@ -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({
Expand Down
73 changes: 68 additions & 5 deletions src/server/mcp.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -118,6 +119,7 @@ export class McpServer {
strictUnions: true,
}) as Tool["inputSchema"])
: EMPTY_OBJECT_JSON_SCHEMA,
annotations: tool.annotations,
};
},
),
Expand Down Expand Up @@ -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 }),
Expand All @@ -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()
},
Expand Down Expand Up @@ -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
};

Expand Down
61 changes: 61 additions & 0 deletions src/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -752,6 +752,62 @@ export const PromptListChangedNotificationSchema = NotificationSchema.extend({
});

/* Tools */
/**
* Additional properties describing a Tool to clients.
*
* NOTE: all properties in ToolAnnotations are **hints**.
* They are not guaranteed to provide a faithful description of
* tool behavior (including descriptive properties like `title`).
*
* Clients should never make tool use decisions based on ToolAnnotations
* received from untrusted servers.
*/
export const ToolAnnotationsSchema = z
.object({
/**
* A human-readable title for the tool.
*/
title: z.optional(z.string()),

/**
* If true, the tool does not modify its environment.
*
* Default: false
*/
readOnlyHint: z.optional(z.boolean()),

/**
* If true, the tool may perform destructive updates to its environment.
* If false, the tool performs only additive updates.
*
* (This property is meaningful only when `readOnlyHint == false`)
*
* Default: true
*/
destructiveHint: z.optional(z.boolean()),

/**
* If true, calling the tool repeatedly with the same arguments
* will have no additional effect on the its environment.
*
* (This property is meaningful only when `readOnlyHint == false`)
*
* Default: false
*/
idempotentHint: z.optional(z.boolean()),

/**
* If true, this tool may interact with an "open world" of external
* entities. If false, the tool's domain of interaction is closed.
* For example, the world of a web search tool is open, whereas that
* of a memory tool is not.
*
* Default: true
*/
openWorldHint: z.optional(z.boolean()),
})
.passthrough();

/**
* Definition for a tool the client can call.
*/
Expand All @@ -774,6 +830,10 @@ export const ToolSchema = z
properties: z.optional(z.object({}).passthrough()),
})
.passthrough(),
/**
* Optional additional tool information.
*/
annotations: z.optional(ToolAnnotationsSchema),
})
.passthrough();

Expand Down Expand Up @@ -1246,6 +1306,7 @@ export type GetPromptResult = Infer<typeof GetPromptResultSchema>;
export type PromptListChangedNotification = Infer<typeof PromptListChangedNotificationSchema>;

/* Tools */
export type ToolAnnotations = Infer<typeof ToolAnnotationsSchema>;
export type Tool = Infer<typeof ToolSchema>;
export type ListToolsRequest = Infer<typeof ListToolsRequestSchema>;
export type ListToolsResult = Infer<typeof ListToolsResultSchema>;
Expand Down