Model Context Protocol in OCaml
anil.recoil.org
1open Mcp
2open Jsonrpc
3
4(* SDK version *)
5let version = "0.1.0"
6
7(* Logging utilities *)
8module Log = struct
9 type level = Debug | Info | Warning | Error
10
11 let string_of_level = function
12 | Debug -> "DEBUG"
13 | Info -> "INFO"
14 | Warning -> "WARNING"
15 | Error -> "ERROR"
16
17 let log level msg =
18 Printf.eprintf "[%s] %s\n" (string_of_level level) msg;
19 flush stderr;
20 Printf.printf "[%s] %s\n" (string_of_level level) msg;
21 flush stdout
22
23 let debug = log Debug
24 let info = log Info
25 let warning = log Warning
26 let error = log Error
27end
28
29(* Context for tools and resources *)
30module Context = struct
31 type t = {
32 request_id: RequestId.t option;
33 lifespan_context: (string * Json.t) list;
34 mutable progress_token: ProgressToken.t option;
35 }
36
37 let create ?request_id ?(lifespan_context=[]) () =
38 { request_id; lifespan_context; progress_token = None }
39
40 let get_context_value ctx key =
41 List.assoc_opt key ctx.lifespan_context
42
43 let report_progress ctx value total =
44 match ctx.progress_token, ctx.request_id with
45 | Some token, Some _id ->
46 let params = `Assoc [
47 ("progress", `Float value);
48 ("total", `Float total);
49 ("progressToken", ProgressToken.yojson_of_t token)
50 ] in
51 Some (create_notification ~method_:"notifications/progress" ~params:(Some params) ())
52 | _ -> None
53end
54
55(* Tools for the MCP server *)
56module Tool = struct
57 type handler = Context.t -> Json.t -> (Json.t, string) result
58
59 type t = {
60 name: string;
61 description: string option;
62 input_schema: Json.t; (* JSON Schema *)
63 handler: handler;
64 }
65
66 let create ~name ?description ~input_schema ~handler () =
67 { name; description; input_schema; handler }
68
69 let to_json tool =
70 let assoc = [
71 ("name", `String tool.name);
72 ("inputSchema", tool.input_schema);
73 ] in
74 let assoc = match tool.description with
75 | Some desc -> ("description", `String desc) :: assoc
76 | None -> assoc
77 in
78 `Assoc assoc
79end
80
81(* Resources for the MCP server *)
82module Resource = struct
83 type handler = Context.t -> string list -> (string, string) result
84
85 type t = {
86 uri_template: string;
87 description: string option;
88 mime_type: string option;
89 handler: handler;
90 }
91
92 let create ~uri_template ?description ?mime_type ~handler () =
93 { uri_template; description; mime_type; handler }
94
95 let to_json resource =
96 let assoc = [
97 ("uriTemplate", `String resource.uri_template);
98 ] in
99 let assoc = match resource.description with
100 | Some desc -> ("description", `String desc) :: assoc
101 | None -> assoc
102 in
103 let assoc = match resource.mime_type with
104 | Some mime -> ("mimeType", `String mime) :: assoc
105 | None -> assoc
106 in
107 `Assoc assoc
108end
109
110(* Prompts for the MCP server *)
111module Prompt = struct
112 type argument = {
113 name: string;
114 description: string option;
115 required: bool;
116 }
117
118 type message = {
119 role: Role.t;
120 content: content;
121 }
122
123 type handler = Context.t -> (string * string) list -> (message list, string) result
124
125 type t = {
126 name: string;
127 description: string option;
128 arguments: argument list;
129 handler: handler;
130 }
131
132 let create ~name ?description ?(arguments=[]) ~handler () =
133 { name; description; arguments; handler }
134
135 let create_argument ~name ?description ?(required=false) () =
136 { name; description; required }
137
138 let to_json prompt =
139 let assoc = [
140 ("name", `String prompt.name);
141 ] in
142 let assoc = match prompt.description with
143 | Some desc -> ("description", `String desc) :: assoc
144 | None -> assoc
145 in
146 let assoc = if prompt.arguments <> [] then
147 let args = List.map (fun (arg: argument) ->
148 let arg_assoc = [
149 ("name", `String arg.name);
150 ] in
151 let arg_assoc = match arg.description with
152 | Some desc -> ("description", `String desc) :: arg_assoc
153 | None -> arg_assoc
154 in
155 let arg_assoc =
156 if arg.required then
157 ("required", `Bool true) :: arg_assoc
158 else
159 arg_assoc
160 in
161 `Assoc arg_assoc
162 ) prompt.arguments in
163 ("arguments", `List args) :: assoc
164 else
165 assoc
166 in
167 `Assoc assoc
168end
169
170(* Helper functions for creating common objects *)
171let make_text_content text =
172 Text (TextContent.{ text; annotations = None })
173
174let make_tool_schema properties required =
175 let props = List.map (fun (name, schema_type, description) ->
176 (name, `Assoc [
177 ("type", `String schema_type);
178 ("description", `String description)
179 ])
180 ) properties in
181 let required_json = `List (List.map (fun name -> `String name) required) in
182 `Assoc [
183 ("type", `String "object");
184 ("properties", `Assoc props);
185 ("required", required_json)
186 ]
187
188(* Main server type *)
189type server = {
190 name: string;
191 version: string;
192 protocol_version: string;
193 mutable capabilities: Json.t;
194 mutable tools: Tool.t list;
195 mutable resources: Resource.t list;
196 mutable prompts: Prompt.t list;
197 mutable lifespan_context: (string * Json.t) list;
198 mutable startup_hook: (unit -> unit) option;
199 mutable shutdown_hook: (unit -> unit) option;
200}
201
202(* Create a new server *)
203let create_server ~name ?(version="0.1.0") ?(protocol_version="2024-11-05") () =
204 {
205 name;
206 version;
207 protocol_version;
208 capabilities = `Assoc [];
209 tools = [];
210 resources = [];
211 prompts = [];
212 lifespan_context = [];
213 startup_hook = None;
214 shutdown_hook = None;
215 }
216
217(* Default capabilities for the server *)
218let default_capabilities ?(with_tools=true) ?(with_resources=false) ?(with_prompts=false) () =
219 let caps = [] in
220 let caps =
221 if with_tools then
222 ("tools", `Assoc [
223 ("listChanged", `Bool true)
224 ]) :: caps
225 else
226 caps
227 in
228 let caps =
229 if with_resources then
230 ("resources", `Assoc [
231 ("listChanged", `Bool true);
232 ("subscribe", `Bool false)
233 ]) :: caps
234 else if not with_resources then
235 ("resources", `Assoc [
236 ("listChanged", `Bool false);
237 ("subscribe", `Bool false)
238 ]) :: caps
239 else
240 caps
241 in
242 let caps =
243 if with_prompts then
244 ("prompts", `Assoc [
245 ("listChanged", `Bool true)
246 ]) :: caps
247 else if not with_prompts then
248 ("prompts", `Assoc [
249 ("listChanged", `Bool false)
250 ]) :: caps
251 else
252 caps
253 in
254 `Assoc caps
255
256(* Register a tool *)
257let register_tool server tool =
258 server.tools <- tool :: server.tools;
259 tool
260
261(* Create and register a tool in one step *)
262let add_tool server ~name ?description ?(schema_properties=[]) ?(schema_required=[]) handler =
263 let input_schema = make_tool_schema schema_properties schema_required in
264 let handler' ctx args =
265 try
266 Ok (handler args)
267 with exn ->
268 Error (Printexc.to_string exn)
269 in
270 let tool = Tool.create
271 ~name
272 ?description
273 ~input_schema
274 ~handler:handler'
275 ()
276 in
277 register_tool server tool
278
279(* Register a resource *)
280let register_resource server resource =
281 server.resources <- resource :: server.resources;
282 resource
283
284(* Create and register a resource in one step *)
285let add_resource server ~uri_template ?description ?mime_type handler =
286 let handler' _ctx params =
287 try
288 Ok (handler params)
289 with exn ->
290 Error (Printexc.to_string exn)
291 in
292 let resource = Resource.create
293 ~uri_template
294 ?description
295 ?mime_type
296 ~handler:handler'
297 ()
298 in
299 register_resource server resource
300
301(* Register a prompt *)
302let register_prompt server prompt =
303 server.prompts <- prompt :: server.prompts;
304 prompt
305
306(* Create and register a prompt in one step *)
307let add_prompt server ~name ?description ?(arguments=[]) handler =
308 let prompt_args = List.map (fun (name, desc, required) ->
309 Prompt.create_argument ~name ?description:desc ~required ()
310 ) arguments in
311 let handler' _ctx args =
312 try
313 Ok (handler args)
314 with exn ->
315 Error (Printexc.to_string exn)
316 in
317 let prompt = Prompt.create
318 ~name
319 ?description
320 ~arguments:prompt_args
321 ~handler:handler'
322 ()
323 in
324 register_prompt server prompt
325
326(* Set server capabilities *)
327let set_capabilities server capabilities =
328 server.capabilities <- capabilities
329
330(* Configure server with default capabilities based on registered components *)
331let configure_server server ?with_tools ?with_resources ?with_prompts () =
332 let with_tools = match with_tools with
333 | Some b -> b
334 | None -> server.tools <> []
335 in
336 let with_resources = match with_resources with
337 | Some b -> b
338 | None -> server.resources <> []
339 in
340 let with_prompts = match with_prompts with
341 | Some b -> b
342 | None -> server.prompts <> []
343 in
344 let capabilities = default_capabilities ~with_tools ~with_resources ~with_prompts () in
345 set_capabilities server capabilities;
346 server
347
348(* Set startup and shutdown hooks *)
349let set_startup_hook server hook =
350 server.startup_hook <- Some hook
351
352let set_shutdown_hook server hook =
353 server.shutdown_hook <- Some hook
354
355(* Run the server *)
356let run_server server =
357 (* Setup *)
358 Printexc.record_backtrace true;
359 set_binary_mode_out stdout false;
360
361 Log.info (Printf.sprintf "%s server started" server.name);
362 Log.debug (Printf.sprintf "Protocol version: %s" server.protocol_version);
363 Log.debug (Printf.sprintf "Server info: %s v%s" server.name server.version);
364
365 (* Initialize capabilities if not already set *)
366 if server.capabilities = `Assoc [] then
367 ignore (configure_server server ());
368
369 (* Run startup hook if provided *)
370 (match server.startup_hook with
371 | Some hook -> hook ()
372 | None -> ());
373
374 Log.info "Server initialized and ready."