fix: update ToolCallback type to include output arguments for better type safety

This commit is contained in:
sushichan044 2025-06-20 16:47:53 +09:00
commit fc9ca5ff2e
No known key found for this signature in database
2 changed files with 45 additions and 26 deletions

View file

@ -1248,8 +1248,7 @@ describe("tool()", () => {
processedInput: input,
resultType: "structured",
// Missing required 'timestamp' field
someExtraField: "unexpected" // Extra field not in schema
},
} as unknown as { processedInput: string; resultType: string; timestamp: string }, // Type assertion to bypass TypeScript validation for testing purposes
})
);

View file

@ -169,7 +169,7 @@ export class McpServer {
}
const args = parseResult.data;
const cb = tool.callback as ToolCallback<ZodRawShape>;
const cb = tool.callback as ToolCallback<ZodRawShape, ZodRawShape>;
try {
result = await Promise.resolve(cb(args, extra));
} catch (error) {
@ -184,7 +184,7 @@ export class McpServer {
};
}
} else {
const cb = tool.callback as ToolCallback<undefined>;
const cb = tool.callback as ToolCallback<undefined, ZodRawShape>;
try {
result = await Promise.resolve(cb(extra));
} catch (error) {
@ -760,7 +760,7 @@ export class McpServer {
inputSchema: ZodRawShape | undefined,
outputSchema: ZodRawShape | undefined,
annotations: ToolAnnotations | undefined,
callback: ToolCallback<ZodRawShape | undefined>
callback: ToolCallback<ZodRawShape | undefined, ZodRawShape | undefined>
): RegisteredTool {
const registeredTool: RegisteredTool = {
title,
@ -917,7 +917,7 @@ export class McpServer {
outputSchema?: OutputArgs;
annotations?: ToolAnnotations;
},
cb: ToolCallback<InputArgs>
cb: ToolCallback<InputArgs, OutputArgs>
): RegisteredTool {
if (this._registeredTools[name]) {
throw new Error(`Tool ${name} is already registered`);
@ -932,7 +932,7 @@ export class McpServer {
inputSchema,
outputSchema,
annotations,
cb as ToolCallback<ZodRawShape | undefined>
cb as ToolCallback<ZodRawShape | undefined, ZodRawShape | undefined>
);
}
@ -1126,6 +1126,16 @@ export class ResourceTemplate {
}
}
/**
* Type helper to create a strongly-typed CallToolResult with structuredContent
*/
type TypedCallToolResult<OutputArgs extends undefined | ZodRawShape> =
OutputArgs extends ZodRawShape
? CallToolResult & {
structuredContent?: z.objectOutputType<OutputArgs, ZodTypeAny>;
}
: CallToolResult;
/**
* Callback for a tool handler registered with Server.tool().
*
@ -1136,13 +1146,21 @@ export class ResourceTemplate {
* - `content` if the tool does not have an outputSchema
* - Both fields are optional but typically one should be provided
*/
export type ToolCallback<Args extends undefined | ZodRawShape = undefined> =
Args extends ZodRawShape
export type ToolCallback<
InputArgs extends undefined | ZodRawShape = undefined,
OutputArgs extends undefined | ZodRawShape = undefined
> = InputArgs extends ZodRawShape
? (
args: z.objectOutputType<Args, ZodTypeAny>,
extra: RequestHandlerExtra<ServerRequest, ServerNotification>,
) => CallToolResult | Promise<CallToolResult>
: (extra: RequestHandlerExtra<ServerRequest, ServerNotification>) => CallToolResult | Promise<CallToolResult>;
args: z.objectOutputType<InputArgs, ZodTypeAny>,
extra: RequestHandlerExtra<ServerRequest, ServerNotification>
) =>
| TypedCallToolResult<OutputArgs>
| Promise<TypedCallToolResult<OutputArgs>>
: (
extra: RequestHandlerExtra<ServerRequest, ServerNotification>
) =>
| TypedCallToolResult<OutputArgs>
| Promise<TypedCallToolResult<OutputArgs>>;
export type RegisteredTool = {
title?: string;
@ -1150,22 +1168,24 @@ export type RegisteredTool = {
inputSchema?: AnyZodObject;
outputSchema?: AnyZodObject;
annotations?: ToolAnnotations;
callback: ToolCallback<undefined | ZodRawShape>;
callback: ToolCallback<ZodRawShape | undefined, ZodRawShape | undefined>;
enabled: boolean;
enable(): void;
disable(): void;
update<InputArgs extends ZodRawShape, OutputArgs extends ZodRawShape>(
updates: {
name?: string | null,
title?: string,
description?: string,
paramsSchema?: InputArgs,
outputSchema?: OutputArgs,
annotations?: ToolAnnotations,
callback?: ToolCallback<InputArgs>,
enabled?: boolean
}): void
remove(): void
update<
InputArgs extends ZodRawShape,
OutputArgs extends ZodRawShape
>(updates: {
name?: string | null;
title?: string;
description?: string;
paramsSchema?: InputArgs;
outputSchema?: OutputArgs;
annotations?: ToolAnnotations;
callback?: ToolCallback<InputArgs, OutputArgs>
enabled?: boolean
}): void;
remove(): void;
};
const EMPTY_OBJECT_JSON_SCHEMA = {