My agentic slop goes here. Not intended for anyone else!
anil.recoil.org
1(** High-level MCP client session implementation *)
2
3(** {1 Configuration} *)
4
5type config = {
6 client_info : Capabilities.Implementation.t;
7 client_capabilities : Capabilities.Client.t;
8}
9
10(** {1 Internal State} *)
11
12exception Initialization_error of string
13
14type notification_handlers = {
15 mutable on_resource_updated : (uri:string -> unit) option;
16 mutable on_resource_list_changed : (unit -> unit) option;
17 mutable on_tool_list_changed : (unit -> unit) option;
18 mutable on_prompt_list_changed : (unit -> unit) option;
19 mutable on_log_message : (level:Messages.Logging.level -> logger:string option -> data:Jsont.json -> unit) option;
20}
21
22type t = {
23 session : Session.t;
24 server_capabilities : Capabilities.Server.t;
25 server_info : Capabilities.Implementation.t;
26 server_instructions : string option;
27 handlers : notification_handlers;
28}
29
30(** {1 Helper Functions} *)
31
32(* Encode a value to JSON using jsont codec *)
33let encode codec value =
34 match Jsont.Json.encode codec value with
35 | Ok json -> json
36 | Error msg -> failwith ("Failed to encode: " ^ msg)
37
38(* Decode a JSON value using jsont codec *)
39let decode codec json =
40 match Jsont.Json.decode codec json with
41 | Ok value -> value
42 | Error msg -> failwith ("Failed to decode: " ^ msg)
43
44(* Parse notification parameters - returns None if params is None or null *)
45let parse_notification_params codec params_opt =
46 match params_opt with
47 | None -> None
48 | Some (Jsont.Null _) -> None
49 | Some json -> Some (decode codec json)
50
51(** {1 Notification Routing} *)
52
53let create_notification_handler handlers : Session.notification_handler =
54 fun ~method_ ~params ->
55 match method_ with
56 | "notifications/resources/updated" ->
57 (match handlers.on_resource_updated with
58 | None -> ()
59 | Some handler ->
60 let notif = parse_notification_params
61 Messages.Resources.updated_notification_jsont params in
62 (match notif with
63 | None -> ()
64 | Some n -> handler ~uri:n.Messages.Resources.uri))
65
66 | "notifications/resources/list_changed" ->
67 (match handlers.on_resource_list_changed with
68 | None -> ()
69 | Some handler -> handler ())
70
71 | "notifications/tools/list_changed" ->
72 (match handlers.on_tool_list_changed with
73 | None -> ()
74 | Some handler -> handler ())
75
76 | "notifications/prompts/list_changed" ->
77 (match handlers.on_prompt_list_changed with
78 | None -> ()
79 | Some handler -> handler ())
80
81 | "notifications/message" ->
82 (match handlers.on_log_message with
83 | None -> ()
84 | Some handler ->
85 let notif = parse_notification_params
86 Messages.Logging.notification_jsont params in
87 (match notif with
88 | None -> ()
89 | Some n ->
90 let data = match n.Messages.Logging.data with
91 | None -> Jsont.Null ((), Jsont.Meta.none)
92 | Some d -> d
93 in
94 handler
95 ~level:n.Messages.Logging.level
96 ~logger:n.Messages.Logging.logger
97 ~data))
98
99 | _ ->
100 (* Unknown notification - ignore *)
101 ()
102
103(** {1 Request Handler} *)
104
105(* Client doesn't expect to receive requests from server in most cases *)
106let create_request_handler () : Session.request_handler =
107 fun ~method_ ~params:_ ->
108 (* Default: return method not found error *)
109 let error = Jsonrpc.Error_data.make
110 ~code:Method_not_found
111 ~message:(Printf.sprintf "Client does not handle method: %s" method_)
112 ()
113 in
114 raise (Session.Remote_error error)
115
116(** {1 Initialization} *)
117
118let perform_initialization session config =
119 (* Send Initialize request *)
120 let init_params = Messages.Initialize.make_request_params
121 ~protocol_version:"2024-11-05"
122 ~capabilities:config.client_capabilities
123 ~client_info:config.client_info
124 ()
125 in
126 let params_json = encode Messages.Initialize.request_params_jsont init_params in
127
128 let response_json = Session.send_request session
129 ~method_:Messages.Initialize.method_
130 ~params:params_json
131 ()
132 in
133
134 (* Decode Initialize result *)
135 let init_result = decode Messages.Initialize.result_jsont response_json in
136
137 (* Send Initialized notification *)
138 let initialized_notif = Messages.Initialized.make_notification () in
139 let notif_json = encode Messages.Initialized.notification_jsont initialized_notif in
140 Session.send_notification session
141 ~method_:Messages.Initialized.method_
142 ~params:notif_json
143 ();
144
145 (* Return server info *)
146 (init_result.Messages.Initialize.capabilities,
147 init_result.Messages.Initialize.server_info,
148 init_result.Messages.Initialize.instructions)
149
150(** {1 Public API} *)
151
152let create ~sw ~transport ?timeout ?clock config =
153 (* Create notification handlers *)
154 let handlers = {
155 on_resource_updated = None;
156 on_resource_list_changed = None;
157 on_tool_list_changed = None;
158 on_prompt_list_changed = None;
159 on_log_message = None;
160 } in
161
162 (* Create session config *)
163 let session_config : Session.config = {
164 transport;
165 request_handler = create_request_handler ();
166 notification_handler = create_notification_handler handlers;
167 timeout;
168 clock;
169 } in
170
171 (* Create underlying session *)
172 let session = Session.create ~sw session_config in
173
174 try
175 (* Perform initialization handshake *)
176 let (server_capabilities, server_info, server_instructions) =
177 perform_initialization session config
178 in
179
180 (* Return client session *)
181 {
182 session;
183 server_capabilities;
184 server_info;
185 server_instructions;
186 handlers;
187 }
188 with
189 | Session.Remote_error err ->
190 Session.close session;
191 raise (Initialization_error
192 (Printf.sprintf "Server returned error: %s" err.Jsonrpc.Error_data.message))
193 | Session.Timeout msg ->
194 Session.close session;
195 raise (Initialization_error ("Initialization timeout: " ^ msg))
196 | exn ->
197 Session.close session;
198 raise (Initialization_error
199 (Printf.sprintf "Initialization failed: %s" (Printexc.to_string exn)))
200
201(** {1 Server Information} *)
202
203let server_capabilities t = t.server_capabilities
204let server_info t = t.server_info
205let server_instructions t = t.server_instructions
206
207(** {1 Basic Operations} *)
208
209let ping t =
210 let params = Messages.Ping.make_params () in
211 let params_json = encode Messages.Ping.params_jsont params in
212 let response_json = Session.send_request t.session
213 ~method_:Messages.Ping.method_
214 ~params:params_json
215 ()
216 in
217 let _result = decode Messages.Ping.result_jsont response_json in
218 ()
219
220(** {1 Resources} *)
221
222let list_resources t ?cursor () =
223 let request = Messages.Resources.make_list_request ?cursor () in
224 let params_json = encode Messages.Resources.list_request_jsont request in
225 let response_json = Session.send_request t.session
226 ~method_:Messages.Resources.list_method
227 ~params:params_json
228 ()
229 in
230 decode Messages.Resources.list_result_jsont response_json
231
232let read_resource t ~uri =
233 let request = Messages.Resources.make_read_request ~uri in
234 let params_json = encode Messages.Resources.read_request_jsont request in
235 let response_json = Session.send_request t.session
236 ~method_:Messages.Resources.read_method
237 ~params:params_json
238 ()
239 in
240 decode Messages.Resources.read_result_jsont response_json
241
242let subscribe_resource t ~uri =
243 let request = Messages.Resources.make_subscribe_request ~uri in
244 let params_json = encode Messages.Resources.subscribe_request_jsont request in
245 let _response_json = Session.send_request t.session
246 ~method_:Messages.Resources.subscribe_method
247 ~params:params_json
248 ()
249 in
250 ()
251
252let unsubscribe_resource t ~uri =
253 let request = Messages.Resources.make_unsubscribe_request ~uri in
254 let params_json = encode Messages.Resources.unsubscribe_request_jsont request in
255 let _response_json = Session.send_request t.session
256 ~method_:Messages.Resources.unsubscribe_method
257 ~params:params_json
258 ()
259 in
260 ()
261
262(** {1 Tools} *)
263
264let list_tools t ?cursor () =
265 let request = Messages.Tools.make_list_request ?cursor () in
266 let params_json = encode Messages.Tools.list_request_jsont request in
267 let response_json = Session.send_request t.session
268 ~method_:Messages.Tools.list_method
269 ~params:params_json
270 ()
271 in
272 decode Messages.Tools.list_result_jsont response_json
273
274let call_tool t ~name ?arguments () =
275 let request = Messages.Tools.make_call_request ~name ?arguments () in
276 let params_json = encode Messages.Tools.call_request_jsont request in
277 let response_json = Session.send_request t.session
278 ~method_:Messages.Tools.call_method
279 ~params:params_json
280 ()
281 in
282 decode Messages.Tools.call_result_jsont response_json
283
284(** {1 Prompts} *)
285
286let list_prompts t ?cursor () =
287 let request = Messages.Prompts.make_list_request ?cursor () in
288 let params_json = encode Messages.Prompts.list_request_jsont request in
289 let response_json = Session.send_request t.session
290 ~method_:Messages.Prompts.list_method
291 ~params:params_json
292 ()
293 in
294 decode Messages.Prompts.list_result_jsont response_json
295
296let get_prompt t ~name ?arguments () =
297 let request = Messages.Prompts.make_get_request ~name ?arguments () in
298 let params_json = encode Messages.Prompts.get_request_jsont request in
299 let response_json = Session.send_request t.session
300 ~method_:Messages.Prompts.get_method
301 ~params:params_json
302 ()
303 in
304 decode Messages.Prompts.get_result_jsont response_json
305
306(** {1 Completions} *)
307
308let complete t ~ref ~argument =
309 let request = Messages.Completions.make_request ~ref_:ref ~argument () in
310 let params_json = encode Messages.Completions.request_jsont request in
311 let response_json = Session.send_request t.session
312 ~method_:Messages.Completions.method_
313 ~params:params_json
314 ()
315 in
316 decode Messages.Completions.result_jsont response_json
317
318(** {1 Logging} *)
319
320let set_log_level t level =
321 (* Create a simple request with level parameter *)
322 let level_json = encode Messages.Logging.level_jsont level in
323 let params = Jsont.Object ([
324 (("level", Jsont.Meta.none), level_json)
325 ], Jsont.Meta.none) in
326 let _response_json = Session.send_request t.session
327 ~method_:"logging/setLevel"
328 ~params
329 ()
330 in
331 ()
332
333(** {1 Notification Handlers} *)
334
335let on_resource_updated t handler =
336 t.handlers.on_resource_updated <- Some handler
337
338let on_resource_list_changed t handler =
339 t.handlers.on_resource_list_changed <- Some handler
340
341let on_tool_list_changed t handler =
342 t.handlers.on_tool_list_changed <- Some handler
343
344let on_prompt_list_changed t handler =
345 t.handlers.on_prompt_list_changed <- Some handler
346
347let on_log_message t handler =
348 t.handlers.on_log_message <- Some handler
349
350(** {1 Session Control} *)
351
352let close t =
353 Session.close t.session
354
355let is_closed t =
356 Session.is_closed t.session