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
21 let debug = log Debug
22 let info = log Info
23 let warning = log Warning
24 let error = log Error
25end
26
27(* Context for tools and resources *)
28module Context = struct
29 type t = {
30 request_id: RequestId.t option;
31 lifespan_context: (string * Json.t) list;
32 mutable progress_token: ProgressToken.t option;
33 }
34
35 let create ?request_id ?(lifespan_context=[]) () =
36 { request_id; lifespan_context; progress_token = None }
37
38 let get_context_value ctx key =
39 List.assoc_opt key ctx.lifespan_context
40
41 let report_progress ctx value total =
42 match ctx.progress_token, ctx.request_id with
43 | Some token, Some _id ->
44 let params = `Assoc [
45 ("progress", `Float value);
46 ("total", `Float total);
47 ("progressToken", ProgressToken.yojson_of_t token)
48 ] in
49 Some (create_notification ~method_:"notifications/progress" ~params:(Some params) ())
50 | _ -> None
51end
52
53(* Tools for the MCP server *)
54module Tool = struct
55 type handler = Context.t -> Json.t -> (Json.t, string) result
56
57 type t = {
58 name: string;
59 description: string option;
60 input_schema: Json.t; (* JSON Schema *)
61 handler: handler;
62 }
63
64 let create ~name ?description ~input_schema ~handler () =
65 { name; description; input_schema; handler }
66
67 let to_json tool =
68 let assoc = [
69 ("name", `String tool.name);
70 ("inputSchema", tool.input_schema);
71 ] in
72 let assoc = match tool.description with
73 | Some desc -> ("description", `String desc) :: assoc
74 | None -> assoc
75 in
76 `Assoc assoc
77end
78
79(* Resources for the MCP server *)
80module Resource = struct
81 type handler = Context.t -> string list -> (string, string) result
82
83 type t = {
84 uri_template: string;
85 description: string option;
86 mime_type: string option;
87 handler: handler;
88 }
89
90 let create ~uri_template ?description ?mime_type ~handler () =
91 { uri_template; description; mime_type; handler }
92
93 let to_json resource =
94 let assoc = [
95 ("uriTemplate", `String resource.uri_template);
96 ] in
97 let assoc = match resource.description with
98 | Some desc -> ("description", `String desc) :: assoc
99 | None -> assoc
100 in
101 let assoc = match resource.mime_type with
102 | Some mime -> ("mimeType", `String mime) :: assoc
103 | None -> assoc
104 in
105 `Assoc assoc
106end
107
108(* Prompts for the MCP server *)
109module Prompt = struct
110 type argument = {
111 name: string;
112 description: string option;
113 required: bool;
114 }
115
116 type message = {
117 role: Role.t;
118 content: content;
119 }
120
121 type handler = Context.t -> (string * string) list -> (message list, string) result
122
123 type t = {
124 name: string;
125 description: string option;
126 arguments: argument list;
127 handler: handler;
128 }
129
130 let create ~name ?description ?(arguments=[]) ~handler () =
131 { name; description; arguments; handler }
132
133 let create_argument ~name ?description ?(required=false) () =
134 { name; description; required }
135
136 let yojson_of_message { role; content } =
137 `Assoc [
138 ("role", Role.yojson_of_t role);
139 ("content", yojson_of_content content);
140 ]
141
142 (* This function must match the structure expected by the PromptMessage module in mcp.ml *)
143 let message_of_yojson json =
144 match json with
145 | `Assoc fields -> begin
146 let role = match List.assoc_opt "role" fields with
147 | Some json -> begin
148 Role.t_of_yojson json
149 end
150 | None -> begin
151 raise (Json.Of_json ("Missing role field", `Assoc fields))
152 end
153 in
154 let content = match List.assoc_opt "content" fields with
155 | Some json -> begin
156 content_of_yojson json
157 end
158 | None -> begin
159 raise (Json.Of_json ("Missing content field", `Assoc fields))
160 end
161 in
162 { role; content }
163 end
164 | j -> begin
165 raise (Json.Of_json ("Expected object for PromptMessage", j))
166 end
167
168 let to_json prompt =
169 let assoc = [
170 ("name", `String prompt.name);
171 ] in
172 let assoc = match prompt.description with
173 | Some desc -> ("description", `String desc) :: assoc
174 | None -> assoc
175 in
176 let assoc = if prompt.arguments <> [] then
177 let args = List.map (fun (arg: argument) ->
178 let arg_assoc = [
179 ("name", `String arg.name);
180 ] in
181 let arg_assoc = match arg.description with
182 | Some desc -> ("description", `String desc) :: arg_assoc
183 | None -> arg_assoc
184 in
185 let arg_assoc =
186 if arg.required then
187 ("required", `Bool true) :: arg_assoc
188 else
189 arg_assoc
190 in
191 `Assoc arg_assoc
192 ) prompt.arguments in
193 ("arguments", `List args) :: assoc
194 else
195 assoc
196 in
197 `Assoc assoc
198end
199
200(* Helper functions for creating common objects *)
201let make_text_content text =
202 Text (TextContent.{ text; annotations = None })
203
204let make_text_content_with_annotations text annotations =
205 Text (TextContent.{ text; annotations = Some annotations })
206
207let make_image_content data mime_type =
208 Image (ImageContent.{ data; mime_type; annotations = None })
209
210let make_image_content_with_annotations data mime_type annotations =
211 Image (ImageContent.{ data; mime_type; annotations = Some annotations })
212
213let make_audio_content data mime_type =
214 Audio (AudioContent.{ data; mime_type; annotations = None })
215
216let make_audio_content_with_annotations data mime_type annotations =
217 Audio (AudioContent.{ data; mime_type; annotations = Some annotations })
218
219let make_text_resource_content uri text ?mime_type () =
220 Resource (EmbeddedResource.{
221 resource = `Text TextResourceContents.{ uri; text; mime_type };
222 annotations = None
223 })
224
225let make_blob_resource_content uri blob ?mime_type () =
226 Resource (EmbeddedResource.{
227 resource = `Blob BlobResourceContents.{ uri; blob; mime_type };
228 annotations = None
229 })
230
231let make_tool_schema properties required =
232 let props = List.map (fun (name, schema_type, description) ->
233 (name, `Assoc [
234 ("type", `String schema_type);
235 ("description", `String description)
236 ])
237 ) properties in
238 let required_json = `List (List.map (fun name -> `String name) required) in
239 `Assoc [
240 ("type", `String "object");
241 ("properties", `Assoc props);
242 ("required", required_json)
243 ]
244
245(* Main server type *)
246type server = {
247 name: string;
248 version: string;
249 protocol_version: string;
250 mutable capabilities: Json.t;
251 mutable tools: Tool.t list;
252 mutable resources: Resource.t list;
253 mutable prompts: Prompt.t list;
254 mutable lifespan_context: (string * Json.t) list;
255 mutable startup_hook: (unit -> unit) option;
256 mutable shutdown_hook: (unit -> unit) option;
257}
258
259(* Create a new server *)
260let create_server ~name ?(version="0.1.0") ?(protocol_version="2024-11-05") () =
261 {
262 name;
263 version;
264 protocol_version;
265 capabilities = `Assoc [];
266 tools = [];
267 resources = [];
268 prompts = [];
269 lifespan_context = [];
270 startup_hook = None;
271 shutdown_hook = None;
272 }
273
274(* Default capabilities for the server *)
275let default_capabilities ?(with_tools=true) ?(with_resources=false) ?(with_prompts=false) () =
276 let caps = [] in
277 let caps =
278 if with_tools then
279 ("tools", `Assoc [
280 ("listChanged", `Bool true)
281 ]) :: caps
282 else
283 caps
284 in
285 let caps =
286 if with_resources then
287 ("resources", `Assoc [
288 ("listChanged", `Bool true);
289 ("subscribe", `Bool false)
290 ]) :: caps
291 else if not with_resources then
292 ("resources", `Assoc [
293 ("listChanged", `Bool false);
294 ("subscribe", `Bool false)
295 ]) :: caps
296 else
297 caps
298 in
299 let caps =
300 if with_prompts then
301 ("prompts", `Assoc [
302 ("listChanged", `Bool true)
303 ]) :: caps
304 else if not with_prompts then
305 ("prompts", `Assoc [
306 ("listChanged", `Bool false)
307 ]) :: caps
308 else
309 caps
310 in
311 `Assoc caps
312
313(* Register a tool *)
314let register_tool server tool =
315 server.tools <- tool :: server.tools;
316 tool
317
318(* Create and register a tool in one step *)
319let add_tool server ~name ?description ?(schema_properties=[]) ?(schema_required=[]) handler =
320 let input_schema = make_tool_schema schema_properties schema_required in
321 let handler' ctx args =
322 try
323 Ok (handler args)
324 with exn ->
325 Error (Printexc.to_string exn)
326 in
327 let tool = Tool.create
328 ~name
329 ?description
330 ~input_schema
331 ~handler:handler'
332 ()
333 in
334 register_tool server tool
335
336(* Register a resource *)
337let register_resource server resource =
338 server.resources <- resource :: server.resources;
339 resource
340
341(* Create and register a resource in one step *)
342let add_resource server ~uri_template ?description ?mime_type handler =
343 let handler' _ctx params =
344 try
345 Ok (handler params)
346 with exn ->
347 Error (Printexc.to_string exn)
348 in
349 let resource = Resource.create
350 ~uri_template
351 ?description
352 ?mime_type
353 ~handler:handler'
354 ()
355 in
356 register_resource server resource
357
358(* Register a prompt *)
359let register_prompt server prompt =
360 server.prompts <- prompt :: server.prompts;
361 prompt
362
363(* Create and register a prompt in one step *)
364let add_prompt server ~name ?description ?(arguments=[]) handler =
365 let prompt_args = List.map (fun (name, desc, required) ->
366 Prompt.create_argument ~name ?description:desc ~required ()
367 ) arguments in
368 let handler' _ctx args =
369 try
370 Ok (handler args)
371 with exn ->
372 Error (Printexc.to_string exn)
373 in
374 let prompt = Prompt.create
375 ~name
376 ?description
377 ~arguments:prompt_args
378 ~handler:handler'
379 ()
380 in
381 register_prompt server prompt
382
383(* Set server capabilities *)
384let set_capabilities server capabilities =
385 server.capabilities <- capabilities
386
387(* Configure server with default capabilities based on registered components *)
388let configure_server server ?with_tools ?with_resources ?with_prompts () =
389 let with_tools = match with_tools with
390 | Some b -> b
391 | None -> server.tools <> []
392 in
393 let with_resources = match with_resources with
394 | Some b -> b
395 | None -> server.resources <> []
396 in
397 let with_prompts = match with_prompts with
398 | Some b -> b
399 | None -> server.prompts <> []
400 in
401 let capabilities = default_capabilities ~with_tools ~with_resources ~with_prompts () in
402 set_capabilities server capabilities;
403 server
404
405(* Set startup and shutdown hooks *)
406let set_startup_hook server hook =
407 server.startup_hook <- Some hook
408
409let set_shutdown_hook server hook =
410 server.shutdown_hook <- Some hook
411
412(* Transport type for server *)
413type transport_type =
414 | Stdio (* Read/write to stdin/stdout *)
415 | Http (* HTTP server - to be implemented *)
416
417(* Run server with stdio transport *)
418let run_server server =
419 (* Setup *)
420 Printexc.record_backtrace true;
421
422 Log.info (Printf.sprintf "%s server starting" server.name);
423 Log.debug (Printf.sprintf "Protocol version: %s" server.protocol_version);
424 Log.debug (Printf.sprintf "Server info: %s v%s" server.name server.version);
425
426 (* Initialize capabilities if not already set *)
427 if server.capabilities = `Assoc [] then
428 ignore (configure_server server ());
429
430 (* Run startup hook if provided *)
431 (match server.startup_hook with
432 | Some hook -> hook ()
433 | None -> ());
434
435 (* This function will be replaced by a full implementation in the mcp_server module *)
436 Log.info "Server initialized and ready."
437
438(* Placeholder for running server with different transports *)
439let run_server_with_transport server transport =
440 match transport with
441 | Http ->
442 Log.info "HTTP server not implemented in this version, using stdio instead";
443 run_server server
444 | Stdio ->
445 run_server server