Model Context Protocol in OCaml
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 logf level fmt = 18 Printf.fprintf stderr "[%s] " (string_of_level level); 19 Printf.kfprintf (fun oc -> Printf.fprintf oc "\n"; flush oc) stderr fmt 20 21 let debugf fmt = logf Debug fmt 22 let infof fmt = logf Info fmt 23 let warningf fmt = logf Warning fmt 24 let errorf fmt = logf Error fmt 25 26 (* Backward compatibility functions that take a simple string *) 27 let log level msg = logf level "%s" msg 28 let debug msg = debugf "%s" msg 29 let info msg = infof "%s" msg 30 let warning msg = warningf "%s" msg 31 let error msg = errorf "%s" msg 32end 33 34(* Context for tools and resources *) 35module Context = struct 36 type t = { 37 request_id: RequestId.t option; 38 lifespan_context: (string * Json.t) list; 39 progress_token: ProgressToken.t option; 40 } 41 42 let create ?request_id ?progress_token ?(lifespan_context=[]) () = 43 { request_id; lifespan_context; progress_token } 44 45 let get_context_value ctx key = 46 List.assoc_opt key ctx.lifespan_context 47 48 let report_progress ctx value total = 49 match ctx.progress_token, ctx.request_id with 50 | Some token, Some _id -> 51 let params = `Assoc [ 52 ("progress", `Float value); 53 ("total", `Float total); 54 ("progressToken", ProgressToken.yojson_of_t token) 55 ] in 56 Some (create_notification ~meth:Method.Progress ~params:(Some params) ()) 57 | _ -> None 58end 59 60(* Tools for the MCP server *) 61module Tool = struct 62 type handler = Context.t -> Json.t -> (Json.t, string) result 63 64 type t = { 65 name: string; 66 description: string option; 67 input_schema: Json.t; (* JSON Schema *) 68 handler: handler; 69 } 70 71 let create ~name ?description ~input_schema ~handler () = 72 { name; description; input_schema; handler } 73 74 let to_json tool = 75 let assoc = [ 76 ("name", `String tool.name); 77 ("inputSchema", tool.input_schema); 78 ] in 79 let assoc = match tool.description with 80 | Some desc -> ("description", `String desc) :: assoc 81 | None -> assoc 82 in 83 `Assoc assoc 84 85 (* Convert to Mcp_rpc.ToolsList.Tool.t *) 86 let to_rpc_tool_list_tool (tool:t) = 87 Mcp_rpc.ToolsList.Tool.{ 88 name = tool.name; 89 description = tool.description; 90 input_schema = tool.input_schema; 91 annotations = None; (* Could be extended to support annotations *) 92 } 93 94 (* Convert a list of Tool.t to the format needed for tools/list response *) 95 let to_rpc_tools_list tools = 96 List.map to_rpc_tool_list_tool tools 97 98 (* Convert Mcp_rpc.ToolsCall response content to Mcp.content list *) 99 let rpc_content_to_mcp_content content = 100 List.map (function 101 | Mcp_rpc.ToolsCall.ToolContent.Text t -> 102 Mcp.Text { TextContent.text = t.text; annotations = None } 103 | Mcp_rpc.ToolsCall.ToolContent.Image i -> 104 Mcp.Image { 105 ImageContent.mime_type = i.mime_type; 106 data = i.data; 107 annotations = None 108 } 109 | Mcp_rpc.ToolsCall.ToolContent.Audio a -> 110 Mcp.Audio { 111 AudioContent.mime_type = a.mime_type; 112 data = a.data; 113 annotations = None 114 } 115 | Mcp_rpc.ToolsCall.ToolContent.Resource r -> 116 (* Create a simple text resource from the embedded resource *) 117 let uri = match r with 118 | { EmbeddedResource.resource = `Text tr; _ } -> tr.uri 119 | { EmbeddedResource.resource = `Blob br; _ } -> br.uri 120 in 121 let text_content = match r with 122 | { EmbeddedResource.resource = `Text tr; _ } -> tr.text 123 | { EmbeddedResource.resource = `Blob br; _ } -> "Binary content" 124 in 125 let mime_type = match r with 126 | { EmbeddedResource.resource = `Text tr; _ } -> tr.mime_type 127 | { EmbeddedResource.resource = `Blob br; _ } -> br.mime_type 128 in 129 let text_resource = { 130 TextResourceContents.uri; 131 text = text_content; 132 mime_type 133 } in 134 Mcp.Resource { 135 EmbeddedResource.resource = `Text text_resource; 136 annotations = None 137 } 138 ) content 139 140 (* Convert Mcp.content list to Mcp_rpc.ToolsCall.ToolContent.t list *) 141 let mcp_content_to_rpc_content content = 142 List.map (function 143 | Mcp.Text t -> 144 Mcp_rpc.ToolsCall.ToolContent.Text t 145 | Mcp.Image img -> 146 Mcp_rpc.ToolsCall.ToolContent.Image img 147 | Mcp.Audio aud -> 148 Mcp_rpc.ToolsCall.ToolContent.Audio aud 149 | Mcp.Resource res -> 150 let resource_data = match res.resource with 151 | `Text txt -> `Text txt 152 | `Blob blob -> `Blob blob 153 in 154 let resource = { 155 EmbeddedResource.resource = resource_data; 156 annotations = res.annotations 157 } in 158 Mcp_rpc.ToolsCall.ToolContent.Resource resource 159 ) content 160 161 (* Create a tool result with content *) 162 let create_tool_result content ~is_error = 163 `Assoc [ 164 ("content", `List (List.map Mcp.yojson_of_content content)); 165 ("isError", `Bool is_error); 166 ] 167 168 (* Create a tool error result with structured content *) 169 let create_error_result error = 170 Log.errorf "Error result: %s" error; 171 create_tool_result [Mcp.make_text_content error] ~is_error:true 172 173 (* Handle tool execution errors *) 174 let handle_execution_error err = 175 create_error_result (Printf.sprintf "Error executing tool: %s" err) 176 177 (* Handle unknown tool error *) 178 let handle_unknown_tool_error name = 179 create_error_result (Printf.sprintf "Unknown tool: %s" name) 180 181 (* Handle general tool execution exception *) 182 let handle_execution_exception exn = 183 create_error_result (Printf.sprintf "Internal error: %s" (Printexc.to_string exn)) 184end 185 186(* Resources for the MCP server *) 187module Resource = struct 188 type handler = Context.t -> string list -> (string, string) result 189 190 type t = { 191 uri: string; (* For resources, this is the exact URI (no variables) *) 192 name: string; 193 description: string option; 194 mime_type: string option; 195 handler: handler; 196 } 197 198 let create ~uri ~name ?description ?mime_type ~handler () = 199 (* Validate that the URI doesn't contain template variables *) 200 if String.contains uri '{' || String.contains uri '}' then 201 Log.warningf "Resource '%s' contains template variables. Consider using add_resource_template instead." uri; 202 { uri; name; description; mime_type; handler } 203 204 let to_json resource = 205 let assoc = [ 206 ("uri", `String resource.uri); 207 ("name", `String resource.name); 208 ] in 209 let assoc = match resource.description with 210 | Some desc -> ("description", `String desc) :: assoc 211 | None -> assoc 212 in 213 let assoc = match resource.mime_type with 214 | Some mime -> ("mimeType", `String mime) :: assoc 215 | None -> assoc 216 in 217 `Assoc assoc 218 219 (* Convert to Mcp_rpc.ResourcesList.Resource.t *) 220 let to_rpc_resource_list_resource (resource:t) = 221 Mcp_rpc.ResourcesList.Resource.{ 222 uri = resource.uri; 223 name = resource.name; 224 description = resource.description; 225 mime_type = resource.mime_type; 226 size = None; (* Size can be added when we have actual resource content *) 227 } 228 229 (* Convert a list of Resource.t to the format needed for resources/list response *) 230 let to_rpc_resources_list resources = 231 List.map to_rpc_resource_list_resource resources 232end 233 234(* Prompts for the MCP server *) 235module Prompt = struct 236 type argument = { 237 name: string; 238 description: string option; 239 required: bool; 240 } 241 242 type message = { 243 role: Role.t; 244 content: content; 245 } 246 247 type handler = Context.t -> (string * string) list -> (message list, string) result 248 249 type t = { 250 name: string; 251 description: string option; 252 arguments: argument list; 253 handler: handler; 254 } 255 256 let create ~name ?description ?(arguments=[]) ~handler () = 257 { name; description; arguments; handler } 258 259 let create_argument ~name ?description ?(required=false) () = 260 { name; description; required } 261 262 let to_json prompt = 263 let assoc = [ 264 ("name", `String prompt.name); 265 ] in 266 let assoc = match prompt.description with 267 | Some desc -> ("description", `String desc) :: assoc 268 | None -> assoc 269 in 270 let assoc = if prompt.arguments <> [] then 271 let args = List.map (fun (arg: argument) -> 272 let arg_assoc = [ 273 ("name", `String arg.name); 274 ] in 275 let arg_assoc = match arg.description with 276 | Some desc -> ("description", `String desc) :: arg_assoc 277 | None -> arg_assoc 278 in 279 let arg_assoc = 280 if arg.required then 281 ("required", `Bool true) :: arg_assoc 282 else 283 arg_assoc 284 in 285 `Assoc arg_assoc 286 ) prompt.arguments in 287 ("arguments", `List args) :: assoc 288 else 289 assoc 290 in 291 `Assoc assoc 292 293 (* Convert argument to Mcp_rpc.PromptsList.PromptArgument.t *) 294 let argument_to_rpc_prompt_argument (arg:argument) = 295 Mcp_rpc.PromptsList.PromptArgument.{ 296 name = arg.name; 297 description = arg.description; 298 required = arg.required; 299 } 300 301 (* Convert to Mcp_rpc.PromptsList.Prompt.t *) 302 let to_rpc_prompt_list_prompt (prompt:t) = 303 Mcp_rpc.PromptsList.Prompt.{ 304 name = prompt.name; 305 description = prompt.description; 306 arguments = List.map argument_to_rpc_prompt_argument prompt.arguments; 307 } 308 309 (* Convert a list of Prompt.t to the format needed for prompts/list response *) 310 let to_rpc_prompts_list prompts = 311 List.map to_rpc_prompt_list_prompt prompts 312 313 (* Convert message to Mcp_rpc.PromptMessage.t *) 314 let message_to_rpc_prompt_message msg = 315 { 316 PromptMessage.role = msg.role; 317 PromptMessage.content = msg.content; 318 } 319 320 (* Convert a list of messages to the format needed for prompts/get response *) 321 let messages_to_rpc_prompt_messages messages = 322 List.map message_to_rpc_prompt_message messages 323end 324 325let make_tool_schema properties required = 326 let props = List.map (fun (name, schema_type, description) -> 327 (name, `Assoc [ 328 ("type", `String schema_type); 329 ("description", `String description) 330 ]) 331 ) properties in 332 let required_json = `List (List.map (fun name -> `String name) required) in 333 `Assoc [ 334 ("type", `String "object"); 335 ("properties", `Assoc props); 336 ("required", required_json) 337 ] 338 339(* Resource Templates for the MCP server *) 340module ResourceTemplate = struct 341 type handler = Context.t -> string list -> (string, string) result 342 343 type t = { 344 uri_template: string; 345 name: string; 346 description: string option; 347 mime_type: string option; 348 handler: handler; 349 } 350 351 let create ~uri_template ~name ?description ?mime_type ~handler () = 352 { uri_template; name; description; mime_type; handler } 353 354 let to_json resource_template = 355 let assoc = [ 356 ("uriTemplate", `String resource_template.uri_template); 357 ("name", `String resource_template.name); 358 ] in 359 let assoc = match resource_template.description with 360 | Some desc -> ("description", `String desc) :: assoc 361 | None -> assoc 362 in 363 let assoc = match resource_template.mime_type with 364 | Some mime -> ("mimeType", `String mime) :: assoc 365 | None -> assoc 366 in 367 `Assoc assoc 368 369 (* Convert to Mcp_rpc.ResourceTemplatesList.ResourceTemplate.t *) 370 let to_rpc_resource_template (template:t) = 371 Mcp_rpc.ListResourceTemplatesResult.ResourceTemplate.{ 372 uri_template = template.uri_template; 373 name = template.name; 374 description = template.description; 375 mime_type = template.mime_type; 376 } 377 378 (* Convert a list of ResourceTemplate.t to the format needed for resources/templates/list response *) 379 let to_rpc_resource_templates_list templates = 380 List.map to_rpc_resource_template templates 381end 382 383(* Main server type *) 384type server = { 385 name: string; 386 version: string; 387 protocol_version: string; 388 lifespan_context: (string * Json.t) list; 389 mutable capabilities: Json.t; 390 mutable tools: Tool.t list; 391 mutable resources: Resource.t list; 392 mutable resource_templates: ResourceTemplate.t list; 393 mutable prompts: Prompt.t list; 394} 395 396let name { name; _ } = name 397let version { version; _ } = version 398let capabilities { capabilities; _ } = capabilities 399let lifespan_context { lifespan_context; _ } = lifespan_context 400let protocol_version { protocol_version; _ } = protocol_version 401let tools { tools; _ } = tools 402let resources { resources; _ } = resources 403let resource_templates { resource_templates; _ } = resource_templates 404let prompts { prompts; _ } = prompts 405 406(* Create a new server *) 407let create_server ~name ?(version="0.1.0") ?(protocol_version="2024-11-05") () = 408 { 409 name; 410 version; 411 protocol_version; 412 capabilities = `Assoc []; 413 tools = []; 414 resources = []; 415 resource_templates = []; 416 prompts = []; 417 lifespan_context = []; 418 } 419 420(* Default capabilities for the server *) 421let default_capabilities ?(with_tools=true) ?(with_resources=false) ?(with_resource_templates=false) ?(with_prompts=false) () = 422 let caps = [] in 423 let caps = 424 if with_tools then 425 ("tools", `Assoc [ 426 ("listChanged", `Bool true) 427 ]) :: caps 428 else 429 caps 430 in 431 let caps = 432 if with_resources then 433 ("resources", `Assoc [ 434 ("listChanged", `Bool true); 435 ("subscribe", `Bool false) 436 ]) :: caps 437 else if not with_resources then 438 ("resources", `Assoc [ 439 ("listChanged", `Bool false); 440 ("subscribe", `Bool false) 441 ]) :: caps 442 else 443 caps 444 in 445 let caps = 446 if with_resource_templates then 447 ("resourceTemplates", `Assoc [ 448 ("listChanged", `Bool true) 449 ]) :: caps 450 else if not with_resource_templates then 451 ("resourceTemplates", `Assoc [ 452 ("listChanged", `Bool false) 453 ]) :: caps 454 else 455 caps 456 in 457 let caps = 458 if with_prompts then 459 ("prompts", `Assoc [ 460 ("listChanged", `Bool true) 461 ]) :: caps 462 else if not with_prompts then 463 ("prompts", `Assoc [ 464 ("listChanged", `Bool false) 465 ]) :: caps 466 else 467 caps 468 in 469 `Assoc caps 470 471(* Register a tool *) 472let register_tool server tool = 473 server.tools <- tool :: server.tools; 474 tool 475 476(* Create and register a tool in one step *) 477let add_tool server ~name ?description ?(schema_properties=[]) ?(schema_required=[]) handler = 478 let input_schema = make_tool_schema schema_properties schema_required in 479 let handler' ctx args = 480 try 481 Ok (handler args) 482 with exn -> 483 Error (Printexc.to_string exn) 484 in 485 let tool = Tool.create 486 ~name 487 ?description 488 ~input_schema 489 ~handler:handler' 490 () 491 in 492 register_tool server tool 493 494(* Register a resource *) 495let register_resource server resource = 496 server.resources <- resource :: server.resources; 497 resource 498 499(* Create and register a resource in one step *) 500let add_resource server ~uri ~name ?description ?mime_type handler = 501 let handler' _ctx params = 502 try 503 Ok (handler params) 504 with exn -> 505 Error (Printexc.to_string exn) 506 in 507 let resource = Resource.create 508 ~uri 509 ~name 510 ?description 511 ?mime_type 512 ~handler:handler' 513 () 514 in 515 register_resource server resource 516 517(* Register a resource template *) 518let register_resource_template server template = 519 server.resource_templates <- template :: server.resource_templates; 520 template 521 522(* Create and register a resource template in one step *) 523let add_resource_template server ~uri_template ~name ?description ?mime_type handler = 524 let handler' _ctx params = 525 try 526 Ok (handler params) 527 with exn -> 528 Error (Printexc.to_string exn) 529 in 530 let template = ResourceTemplate.create 531 ~uri_template 532 ~name 533 ?description 534 ?mime_type 535 ~handler:handler' 536 () 537 in 538 register_resource_template server template 539 540(* Register a prompt *) 541let register_prompt server prompt = 542 server.prompts <- prompt :: server.prompts; 543 prompt 544 545(* Create and register a prompt in one step *) 546let add_prompt server ~name ?description ?(arguments=[]) handler = 547 let prompt_args = List.map (fun (name, desc, required) -> 548 Prompt.create_argument ~name ?description:desc ~required () 549 ) arguments in 550 let handler' _ctx args = 551 try 552 Ok (handler args) 553 with exn -> 554 Error (Printexc.to_string exn) 555 in 556 let prompt = Prompt.create 557 ~name 558 ?description 559 ~arguments:prompt_args 560 ~handler:handler' 561 () 562 in 563 register_prompt server prompt 564 565(* Set server capabilities *) 566let set_capabilities server capabilities = 567 server.capabilities <- capabilities 568 569(* Configure server with default capabilities based on registered components *) 570let configure_server server ?with_tools ?with_resources ?with_resource_templates ?with_prompts () = 571 let with_tools = match with_tools with 572 | Some b -> b 573 | None -> server.tools <> [] 574 in 575 let with_resources = match with_resources with 576 | Some b -> b 577 | None -> server.resources <> [] 578 in 579 let with_resource_templates = match with_resource_templates with 580 | Some b -> b 581 | None -> server.resource_templates <> [] 582 in 583 let with_prompts = match with_prompts with 584 | Some b -> b 585 | None -> server.prompts <> [] 586 in 587 let capabilities = default_capabilities ~with_tools ~with_resources ~with_resource_templates ~with_prompts () in 588 set_capabilities server capabilities; 589 server