Skip to content

Commit ccbc6d3

Browse files
committed
feat: introduce CustomMcpServer class that handles JSONSchema
1 parent 89db562 commit ccbc6d3

18 files changed

+865
-583
lines changed

package-lock.json

Lines changed: 88 additions & 17 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

package.json

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,10 +17,11 @@
1717
"description": "",
1818
"dependencies": {
1919
"@modelcontextprotocol/sdk": "^1.11.0",
20+
"ajv": "^8.17.1",
2021
"algoliasearch": "^5.23.4",
2122
"commander": "^13.1.0",
2223
"open": "^10.1.0",
23-
"zod": "^3.24.2"
24+
"zod": "^3.24.4"
2425
},
2526
"devDependencies": {
2627
"@eslint/js": "^9.24.0",

src/CustomMcpServer.ts

Lines changed: 159 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,159 @@
1+
import { Server } from "@modelcontextprotocol/sdk/server/index.js";
2+
import type { RequestHandlerExtra } from "@modelcontextprotocol/sdk/shared/protocol.js";
3+
import type { Transport } from "@modelcontextprotocol/sdk/shared/transport.js";
4+
import type {
5+
CallToolResult as BaseCallToolResult,
6+
ServerNotification,
7+
ServerRequest,
8+
Tool,
9+
} from "@modelcontextprotocol/sdk/types.js";
10+
import {
11+
CallToolRequestSchema,
12+
ErrorCode,
13+
ListToolsRequestSchema,
14+
McpError,
15+
type ToolAnnotations,
16+
} from "@modelcontextprotocol/sdk/types.js";
17+
import { Ajv2020 as Ajv } from "ajv/dist/2020.js";
18+
import type { SomeJSONSchema } from "ajv/dist/types/json-schema.js";
19+
20+
export type InputJsonSchema = Partial<SomeJSONSchema>;
21+
22+
type CallToolResult = string | BaseCallToolResult;
23+
type ToolCallback<Args extends undefined | InputJsonSchema = undefined> =
24+
Args extends InputJsonSchema
25+
? (
26+
args: object,
27+
extra: RequestHandlerExtra<ServerRequest, ServerNotification>,
28+
) => CallToolResult | Promise<CallToolResult>
29+
: (
30+
extra: RequestHandlerExtra<ServerRequest, ServerNotification>,
31+
) => CallToolResult | Promise<CallToolResult>;
32+
33+
export type ToolDefinition<T extends InputJsonSchema | undefined = undefined> = {
34+
name: string;
35+
inputSchema: T;
36+
description?: string;
37+
annotations?: ToolAnnotations;
38+
cb: ToolCallback<T>;
39+
};
40+
41+
function formatToolError(error: unknown): BaseCallToolResult {
42+
return {
43+
content: [
44+
{
45+
type: "text",
46+
text: error instanceof Error ? error.message : String(error),
47+
},
48+
],
49+
isError: true,
50+
};
51+
}
52+
53+
function formatToolResult(result: CallToolResult): BaseCallToolResult {
54+
if (typeof result === "string") {
55+
return {
56+
content: [{ type: "text", text: result }],
57+
};
58+
}
59+
60+
return result;
61+
}
62+
63+
export class CustomMcpServer {
64+
private readonly server: Server;
65+
private ajv: Ajv;
66+
// eslint-disable-next-line @typescript-eslint/no-explicit-any
67+
private tools: Record<string, ToolDefinition<any>> = {};
68+
69+
constructor(...args: ConstructorParameters<typeof Server>) {
70+
this.ajv = new Ajv({ removeAdditional: true, strict: false });
71+
this.server = new Server(...args);
72+
73+
this.server.assertCanSetRequestHandler(ListToolsRequestSchema.shape.method.value);
74+
this.server.assertCanSetRequestHandler(CallToolRequestSchema.shape.method.value);
75+
this.server.registerCapabilities({ tools: { listChanged: true } });
76+
77+
this.server.setRequestHandler(ListToolsRequestSchema, () => {
78+
return {
79+
tools: Object.values(this.tools).map<Tool>((tool) => {
80+
return {
81+
name: tool.name,
82+
description: tool.description,
83+
// eslint-disable-next-line @typescript-eslint/no-explicit-any
84+
inputSchema: (tool.inputSchema as any) ?? { type: "object" },
85+
annotations: tool.annotations,
86+
};
87+
}),
88+
};
89+
});
90+
91+
this.server.setRequestHandler(
92+
CallToolRequestSchema,
93+
async (request, extra): Promise<BaseCallToolResult> => {
94+
const tool = this.tools[request.params.name];
95+
if (!tool) {
96+
throw new McpError(ErrorCode.InvalidParams, `Tool ${request.params.name} not found`);
97+
}
98+
99+
if (tool.inputSchema) {
100+
const validate = this.ajv.compile(tool.inputSchema);
101+
const callbackArguments: Record<string, unknown> =
102+
structuredClone(request.params.arguments) ?? {};
103+
104+
if (
105+
"requestBody" in callbackArguments &&
106+
typeof callbackArguments.requestBody === "string"
107+
) {
108+
callbackArguments.requestBody = JSON.parse(callbackArguments.requestBody);
109+
}
110+
111+
const isValid = validate(callbackArguments);
112+
if (!isValid) {
113+
throw new McpError(
114+
ErrorCode.InvalidParams,
115+
`Invalid arguments for tool ${request.params.name}: ${this.ajv.errorsText(validate.errors)}`,
116+
);
117+
}
118+
119+
try {
120+
const cb = tool.cb as unknown as ToolCallback<InputJsonSchema>;
121+
return formatToolResult(await cb(callbackArguments, extra));
122+
} catch (error) {
123+
return formatToolError(error);
124+
}
125+
}
126+
127+
try {
128+
const cb = tool.cb as unknown as ToolCallback<undefined>;
129+
return formatToolResult(await cb(extra));
130+
} catch (error) {
131+
return formatToolError(error);
132+
}
133+
},
134+
);
135+
}
136+
137+
async connect(transport: Transport): Promise<void> {
138+
return await this.server.connect(transport);
139+
}
140+
141+
async close(): Promise<void> {
142+
await this.server.close();
143+
}
144+
145+
tool<T extends InputJsonSchema | undefined>(options: ToolDefinition<T>) {
146+
if (this.tools[options.name]) {
147+
throw new Error(`Tool with name ${options.name} already exists`);
148+
}
149+
150+
this.tools[options.name] = options;
151+
this.sendToolListChanged();
152+
}
153+
154+
private sendToolListChanged() {
155+
if (this.server.transport) {
156+
this.server.sendToolListChanged();
157+
}
158+
}
159+
}

src/commands/start-server.ts

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
#!/usr/bin/env -S node --experimental-strip-types
22

3-
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
43
import { authenticate } from "../authentication.ts";
54
import { AppStateManager } from "../appState.ts";
65
import { DashboardApi } from "../DashboardApi.ts";
@@ -36,6 +35,8 @@ import {
3635
operationId as SetCustomRankingOperationId,
3736
} from "../tools/registerSetCustomRanking.ts";
3837

38+
import { CustomMcpServer } from "../CustomMcpServer.ts";
39+
3940
export type StartServerOptions = CliFilteringOptions;
4041

4142
export async function startServer(opts: StartServerOptions) {
@@ -56,7 +57,7 @@ export async function startServer(opts: StartServerOptions) {
5657
appState,
5758
});
5859

59-
const server = new McpServer({
60+
const server = new CustomMcpServer({
6061
name: "algolia",
6162
version: CONFIG.version,
6263
capabilities: {

0 commit comments

Comments
 (0)