Skip to content

Commit bef48bd

Browse files
omercnetomercnet
committed
fix
1 parent fc01943 commit bef48bd

File tree

2 files changed

+153
-61
lines changed

2 files changed

+153
-61
lines changed

‎src/utils.test.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -74,7 +74,7 @@ describe("utils", () => {
7474
"test_tool",
7575
expect.objectContaining({
7676
description: "Test tool",
77-
inputSchema: expect.objectContaining({ input: expect.any(Object) }),
77+
inputSchema: expect.any(Object), // Now it's a ZodObject
7878
}),
7979
expect.any(Function),
8080
);

‎src/utils.ts‎

Lines changed: 152 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,16 @@
1-
import { AuthInfo } from "./schemas/auth.js";
2-
import { z, ZodRawShape } from "zod";
31
import {
42
McpServer,
5-
RegisteredTool,
6-
ToolCallback,
3+
RegisteredTool
74
} from "@modelcontextprotocol/sdk/server/mcp.js";
5+
import { RequestHandlerExtra } from "@modelcontextprotocol/sdk/shared/protocol.js";
86
import {
97
CallToolResult,
10-
ServerRequest,
118
ServerNotification,
9+
ServerRequest,
1210
ToolAnnotations,
1311
} from "@modelcontextprotocol/sdk/types.js";
14-
import { RequestHandlerExtra } from "@modelcontextprotocol/sdk/shared/protocol.js";
12+
import { z, ZodRawShape } from "zod";
13+
import { AuthInfo } from "./schemas/auth.js";
1514
import { getOutboundToken } from "./utils/outboundToken.js";
1615
import {
1716
getRequestContext,
@@ -157,62 +156,127 @@ export function registerAuthenticatedTool(
157156
cb: (...args: any[]) => CallToolResult | Promise<CallToolResult>,
158157
requiredScopes: string[] = [],
159158
) {
160-
return (server: McpServer) => {
161-
// eslint-disable-next-line @typescript-eslint/no-explicit-any
162-
const wrapped: ToolCallback<any> = async (
163-
args: unknown,
164-
extra: RequestHandlerExtra<ServerRequest, ServerNotification>,
165-
) => {
166-
// Get auth context from the server
167-
const context = getRequestContext(server as ServerWithContext);
168-
169-
if (!context?.authInfo) {
170-
throw new Error(
171-
`Authentication required for tool "${name}". Ensure a valid bearer token is provided.`,
172-
);
173-
}
174-
175-
const { authInfo, descopeConfig } = context;
176-
177-
// Scope validation
178-
if (requiredScopes.length) {
179-
const missing = requiredScopes.filter(
180-
(s) => !authInfo.scopes?.includes(s),
181-
);
182-
if (missing.length) {
183-
const userScopes = authInfo.scopes?.join(", ") || "none";
159+
return (server: McpServer): RegisteredTool => {
160+
// Convert ZodRawShape to ZodObject and register with MCP server
161+
// We need to handle the two cases separately to maintain correct types
162+
if (config.inputSchema) {
163+
// Tool WITH input schema
164+
const wrapped = async (
165+
args: unknown,
166+
extra: RequestHandlerExtra<ServerRequest, ServerNotification>,
167+
) => {
168+
// Get auth context from the server
169+
const context = getRequestContext(server as ServerWithContext);
170+
171+
if (!context?.authInfo) {
184172
throw new Error(
185-
`Tool "${name}" requires scopes: ${requiredScopes.join(", ")}. ` +
186-
`User has scopes: ${userScopes}. ` +
187-
`Missing: ${missing.join(", ")}. ` +
188-
`Request these scopes during authentication.`,
173+
`Authentication required for tool "${name}". Ensure a valid bearer token is provided.`,
174+
);
175+
}
176+
177+
const { authInfo, descopeConfig } = context;
178+
179+
// Scope validation
180+
if (requiredScopes.length) {
181+
const missing = requiredScopes.filter(
182+
(s) => !authInfo.scopes?.includes(s),
189183
);
184+
if (missing.length) {
185+
const userScopes = authInfo.scopes?.join(", ") || "none";
186+
throw new Error(
187+
`Tool "${name}" requires scopes: ${requiredScopes.join(", ")}. ` +
188+
`User has scopes: ${userScopes}. ` +
189+
`Missing: ${missing.join(", ")}. ` +
190+
`Request these scopes during authentication.`,
191+
);
192+
}
190193
}
191-
}
192194

193-
// getOutboundToken bound to this request
194-
const getOutboundTokenFn = (appId: string, scopes?: string[]) =>
195-
descopeConfig
196-
? getOutboundToken(appId, authInfo, descopeConfig, scopes)
197-
: Promise.resolve(null);
195+
// getOutboundToken bound to this request
196+
const getOutboundTokenFn = (appId: string, scopes?: string[]) =>
197+
descopeConfig
198+
? getOutboundToken(appId, authInfo, descopeConfig, scopes)
199+
: Promise.resolve(null);
198200

199-
const authExtra: AuthenticatedExtra = {
201+
const authExtra: AuthenticatedExtra = {
202+
...extra,
203+
authInfo,
204+
getOutboundToken: getOutboundTokenFn,
205+
};
206+
207+
// Call user-supplied handler with args
200208
// eslint-disable-next-line @typescript-eslint/no-explicit-any
201-
...(extra as any),
202-
authInfo,
203-
getOutboundToken: getOutboundTokenFn,
209+
return (cb as any)(args, authExtra);
210+
};
211+
212+
// Use explicit any to prevent deep type instantiation errors with ZodObject
213+
const mcpConfigWithInput: any = {
214+
title: config.title,
215+
description: config.description,
216+
inputSchema: z.object(config.inputSchema),
217+
outputSchema: config.outputSchema ? z.object(config.outputSchema) : undefined,
218+
annotations: config.annotations,
204219
};
220+
// eslint-disable-next-line @typescript-eslint/no-explicit-any
221+
return server.registerTool(name, mcpConfigWithInput, wrapped as any);
222+
} else {
223+
// Tool WITHOUT input schema
224+
const wrapped = async (
225+
extra: RequestHandlerExtra<ServerRequest, ServerNotification>,
226+
) => {
227+
// Get auth context from the server
228+
const context = getRequestContext(server as ServerWithContext);
205229

206-
// Call user-supplied handler
207-
return config.inputSchema
208-
? // eslint-disable-next-line @typescript-eslint/no-explicit-any
209-
(cb as any)(args, authExtra)
210-
: // eslint-disable-next-line @typescript-eslint/no-explicit-any
211-
(cb as any)(authExtra);
212-
};
230+
if (!context?.authInfo) {
231+
throw new Error(
232+
`Authentication required for tool "${name}". Ensure a valid bearer token is provided.`,
233+
);
234+
}
235+
236+
const { authInfo, descopeConfig } = context;
237+
238+
// Scope validation
239+
if (requiredScopes.length) {
240+
const missing = requiredScopes.filter(
241+
(s) => !authInfo.scopes?.includes(s),
242+
);
243+
if (missing.length) {
244+
const userScopes = authInfo.scopes?.join(", ") || "none";
245+
throw new Error(
246+
`Tool "${name}" requires scopes: ${requiredScopes.join(", ")}. ` +
247+
`User has scopes: ${userScopes}. ` +
248+
`Missing: ${missing.join(", ")}. ` +
249+
`Request these scopes during authentication.`,
250+
);
251+
}
252+
}
253+
254+
// getOutboundToken bound to this request
255+
const getOutboundTokenFn = (appId: string, scopes?: string[]) =>
256+
descopeConfig
257+
? getOutboundToken(appId, authInfo, descopeConfig, scopes)
258+
: Promise.resolve(null);
259+
260+
const authExtra: AuthenticatedExtra = {
261+
...extra,
262+
authInfo,
263+
getOutboundToken: getOutboundTokenFn,
264+
};
213265

214-
// eslint-disable-next-line @typescript-eslint/no-explicit-any
215-
return server.registerTool(name, config, wrapped as any);
266+
// Call user-supplied handler without args
267+
// eslint-disable-next-line @typescript-eslint/no-explicit-any
268+
return (cb as any)(authExtra);
269+
};
270+
271+
// Use explicit any to prevent deep type instantiation errors with ZodObject
272+
const mcpConfigWithoutInput: any = {
273+
title: config.title,
274+
description: config.description,
275+
outputSchema: config.outputSchema ? z.object(config.outputSchema) : undefined,
276+
annotations: config.annotations,
277+
};
278+
return server.registerTool(name, mcpConfigWithoutInput, wrapped);
279+
}
216280
};
217281
}
218282

@@ -240,6 +304,38 @@ export function registerAuthenticatedTool(
240304
* });
241305
* ```
242306
*/
307+
308+
// Overload with input schema
309+
export function defineTool<
310+
I extends ZodRawShape,
311+
O extends ZodRawShape | undefined = undefined,
312+
>(cfg: {
313+
name: string;
314+
title?: string;
315+
description?: string;
316+
input: I;
317+
output?: O;
318+
scopes?: string[];
319+
annotations?: ToolAnnotations;
320+
handler: (
321+
args: z.infer<z.ZodObject<I>>,
322+
extra: AuthenticatedExtra,
323+
) => CallToolResult | Promise<CallToolResult>;
324+
}): (server: McpServer) => RegisteredTool;
325+
326+
// Overload without input schema
327+
export function defineTool<O extends ZodRawShape | undefined = undefined>(cfg: {
328+
name: string;
329+
title?: string;
330+
description?: string;
331+
input?: undefined;
332+
output?: O;
333+
scopes?: string[];
334+
annotations?: ToolAnnotations;
335+
handler: (extra: AuthenticatedExtra) => CallToolResult | Promise<CallToolResult>;
336+
}): (server: McpServer) => RegisteredTool;
337+
338+
// Implementation
243339
export function defineTool<
244340
I extends ZodRawShape | undefined = undefined,
245341
O extends ZodRawShape | undefined = undefined,
@@ -251,13 +347,9 @@ export function defineTool<
251347
output?: O;
252348
scopes?: string[];
253349
annotations?: ToolAnnotations;
254-
handler: I extends ZodRawShape
255-
? (
256-
args: z.infer<z.ZodObject<I>>,
257-
extra: AuthenticatedExtra,
258-
) => CallToolResult | Promise<CallToolResult>
259-
: (extra: AuthenticatedExtra) => CallToolResult | Promise<CallToolResult>;
260-
}) {
350+
// eslint-disable-next-line @typescript-eslint/no-explicit-any
351+
handler: any;
352+
}): (server: McpServer) => RegisteredTool {
261353
if (cfg.input) {
262354
// With input schema
263355
return registerAuthenticatedTool(

0 commit comments

Comments
 (0)