My agentic slop goes here. Not intended for anyone else!
anil.recoil.org
1(** High-level MCP server session API *)
2
3(** {1 Types} *)
4
5type config = {
6 server_info : Capabilities.Implementation.t;
7 server_capabilities : Capabilities.Server.t;
8 instructions : string option;
9}
10
11type handlers = {
12 list_resources : (cursor:string option -> Messages.Resources.list_result) option;
13 list_resource_templates : (cursor:string option -> Messages.Resources.list_result) option;
14 read_resource : (uri:string -> Messages.Resources.read_result) option;
15 subscribe_resource : (uri:string -> unit) option;
16 unsubscribe_resource : (uri:string -> unit) option;
17 list_tools : (cursor:string option -> Messages.Tools.list_result) option;
18 call_tool : (name:string -> arguments:Jsont.json option -> Messages.Tools.call_result) option;
19 list_prompts : (cursor:string option -> Messages.Prompts.list_result) option;
20 get_prompt : (name:string -> arguments:(string * string) list option -> Messages.Prompts.get_result) option;
21 complete : (ref_:Messages.Completions.completion_ref -> argument:string -> Messages.Completions.result) option;
22 ping : (unit -> unit) option;
23}
24
25type t = {
26 session : Session.t;
27 config : config;
28 handlers : handlers;
29 mutable client_capabilities : Capabilities.Client.t option;
30 mutable client_info : Capabilities.Implementation.t option;
31 mutable protocol_version : string option;
32 mutable initialized : bool;
33}
34
35(** {1 Helper Functions} *)
36
37let encode_json jsont value =
38 match Jsont.Json.encode jsont value with
39 | Ok json -> json
40 | Error e -> failwith ("Failed to encode JSON: " ^ e)
41
42let decode_json jsont json =
43 match Jsont.Json.decode jsont json with
44 | Ok value -> value
45 | Error e -> failwith ("Failed to decode JSON: " ^ e)
46
47let method_not_found method_ =
48 let error = Jsonrpc.Error_data.make
49 ~code:Method_not_found
50 ~message:(Printf.sprintf "Method not found: %s" method_)
51 ()
52 in
53 raise (Session.Remote_error error)
54
55let invalid_params method_ msg =
56 let error = Jsonrpc.Error_data.make
57 ~code:Invalid_params
58 ~message:(Printf.sprintf "Invalid params for %s: %s" method_ msg)
59 ()
60 in
61 raise (Session.Remote_error error)
62
63(** {1 Request Handler} *)
64
65let handle_request t ~method_ ~params =
66 (* Ensure initialization has completed for non-init requests *)
67 if method_ <> Messages.Initialize.method_ && not t.initialized then begin
68 let error = Jsonrpc.Error_data.make
69 ~code:Internal_error
70 ~message:"Server not initialized"
71 ()
72 in
73 raise (Session.Remote_error error)
74 end;
75
76 (* Route to appropriate handler *)
77 match method_ with
78 | m when m = Messages.Initialize.method_ ->
79 (* Handle initialization *)
80 let req_params = match params with
81 | Some p -> decode_json Messages.Initialize.request_params_jsont p
82 | None -> invalid_params method_ "missing params"
83 in
84
85 (* Store client info *)
86 t.client_capabilities <- Some req_params.capabilities;
87 t.client_info <- Some req_params.client_info;
88 t.protocol_version <- Some req_params.protocol_version;
89
90 (* Build response *)
91 let result = Messages.Initialize.make_result
92 ~protocol_version:req_params.protocol_version
93 ~capabilities:t.config.server_capabilities
94 ~server_info:t.config.server_info
95 ?instructions:t.config.instructions
96 ()
97 in
98 encode_json Messages.Initialize.result_jsont result
99
100 | m when m = Messages.Ping.method_ ->
101 let handler = t.handlers.ping in
102 (match handler with
103 | None -> method_not_found method_
104 | Some h ->
105 h ();
106 let result = Messages.Ping.make_result () in
107 encode_json Messages.Ping.result_jsont result)
108
109 | m when m = Messages.Resources.list_method ->
110 let handler = t.handlers.list_resources in
111 (match handler with
112 | None -> method_not_found method_
113 | Some h ->
114 let req = match params with
115 | Some p -> decode_json Messages.Resources.list_request_jsont p
116 | None -> Messages.Resources.make_list_request ()
117 in
118 let result = h ~cursor:req.cursor in
119 encode_json Messages.Resources.list_result_jsont result)
120
121 | m when m = Messages.Resources.read_method ->
122 let handler = t.handlers.read_resource in
123 (match handler with
124 | None -> method_not_found method_
125 | Some h ->
126 let req = match params with
127 | Some p -> decode_json Messages.Resources.read_request_jsont p
128 | None -> invalid_params method_ "missing params"
129 in
130 let result = h ~uri:req.uri in
131 encode_json Messages.Resources.read_result_jsont result)
132
133 | m when m = Messages.Resources.subscribe_method ->
134 let handler = t.handlers.subscribe_resource in
135 (match handler with
136 | None -> method_not_found method_
137 | Some h ->
138 let req = match params with
139 | Some p -> decode_json Messages.Resources.subscribe_request_jsont p
140 | None -> invalid_params method_ "missing params"
141 in
142 h ~uri:req.uri;
143 Jsont.Object ([], Jsont.Meta.none)) (* Empty response *)
144
145 | m when m = Messages.Resources.unsubscribe_method ->
146 let handler = t.handlers.unsubscribe_resource in
147 (match handler with
148 | None -> method_not_found method_
149 | Some h ->
150 let req = match params with
151 | Some p -> decode_json Messages.Resources.unsubscribe_request_jsont p
152 | None -> invalid_params method_ "missing params"
153 in
154 h ~uri:req.uri;
155 Jsont.Object ([], Jsont.Meta.none)) (* Empty response *)
156
157 | m when m = Messages.Tools.list_method ->
158 let handler = t.handlers.list_tools in
159 (match handler with
160 | None -> method_not_found method_
161 | Some h ->
162 let req = match params with
163 | Some p -> decode_json Messages.Tools.list_request_jsont p
164 | None -> Messages.Tools.make_list_request ()
165 in
166 let result = h ~cursor:req.cursor in
167 encode_json Messages.Tools.list_result_jsont result)
168
169 | m when m = Messages.Tools.call_method ->
170 let handler = t.handlers.call_tool in
171 (match handler with
172 | None -> method_not_found method_
173 | Some h ->
174 let req = match params with
175 | Some p -> decode_json Messages.Tools.call_request_jsont p
176 | None -> invalid_params method_ "missing params"
177 in
178 let result = h ~name:req.name ~arguments:req.arguments in
179 encode_json Messages.Tools.call_result_jsont result)
180
181 | m when m = Messages.Prompts.list_method ->
182 let handler = t.handlers.list_prompts in
183 (match handler with
184 | None -> method_not_found method_
185 | Some h ->
186 let req = match params with
187 | Some p -> decode_json Messages.Prompts.list_request_jsont p
188 | None -> Messages.Prompts.make_list_request ()
189 in
190 let result = h ~cursor:req.cursor in
191 encode_json Messages.Prompts.list_result_jsont result)
192
193 | m when m = Messages.Prompts.get_method ->
194 let handler = t.handlers.get_prompt in
195 (match handler with
196 | None -> method_not_found method_
197 | Some h ->
198 let req = match params with
199 | Some p -> decode_json Messages.Prompts.get_request_jsont p
200 | None -> invalid_params method_ "missing params"
201 in
202 let result = h ~name:req.name ~arguments:req.arguments in
203 encode_json Messages.Prompts.get_result_jsont result)
204
205 | m when m = Messages.Completions.method_ ->
206 let handler = t.handlers.complete in
207 (match handler with
208 | None -> method_not_found method_
209 | Some h ->
210 let req = match params with
211 | Some p -> decode_json Messages.Completions.request_jsont p
212 | None -> invalid_params method_ "missing params"
213 in
214 let argument = match req.argument with
215 | Some a -> a
216 | None -> ""
217 in
218 let result = h ~ref_:req.ref_ ~argument in
219 encode_json Messages.Completions.result_jsont result)
220
221 | _ ->
222 method_not_found method_
223
224(** {1 Notification Handler} *)
225
226let handle_notification t ~method_ ~params =
227 match method_ with
228 | m when m = Messages.Initialized.method_ ->
229 (* Client has confirmed initialization *)
230 let _notif = match params with
231 | Some p -> decode_json Messages.Initialized.notification_jsont p
232 | None -> Messages.Initialized.make_notification ()
233 in
234 t.initialized <- true
235
236 | _ ->
237 (* Ignore unknown notifications *)
238 ()
239
240(** {1 Public API} *)
241
242let create ~sw ~transport ?timeout ?clock config handlers =
243 (* Create session with handlers *)
244 let t_ref = ref None in
245
246 let request_handler ~method_ ~params =
247 match !t_ref with
248 | None -> failwith "Server session not initialized"
249 | Some t -> handle_request t ~method_ ~params
250 in
251
252 let notification_handler ~method_ ~params =
253 match !t_ref with
254 | None -> ()
255 | Some t -> handle_notification t ~method_ ~params
256 in
257
258 let session_config = {
259 Session.transport;
260 request_handler;
261 notification_handler;
262 timeout;
263 clock;
264 } in
265
266 let session = Session.create ~sw session_config in
267
268 let t = {
269 session;
270 config;
271 handlers;
272 client_capabilities = None;
273 client_info = None;
274 protocol_version = None;
275 initialized = false;
276 } in
277
278 t_ref := Some t;
279
280 t
281
282let client_capabilities t =
283 match t.client_capabilities with
284 | Some c -> c
285 | None -> invalid_arg "Server_session.client_capabilities: not initialized"
286
287let client_info t =
288 match t.client_info with
289 | Some i -> i
290 | None -> invalid_arg "Server_session.client_info: not initialized"
291
292let protocol_version t =
293 match t.protocol_version with
294 | Some v -> v
295 | None -> invalid_arg "Server_session.protocol_version: not initialized"
296
297(** {1 Sending Notifications} *)
298
299let send_notification t method_ params_jsont params =
300 let params_json = encode_json params_jsont params in
301 Session.send_notification t.session ~method_ ~params:params_json ()
302
303let send_resource_updated t ~uri =
304 let notif = Messages.Resources.make_updated_notification ~uri in
305 send_notification t
306 Messages.Resources.updated_notification_method
307 Messages.Resources.updated_notification_jsont
308 notif
309
310let send_resource_list_changed t =
311 let notif = Messages.Resources.make_list_changed_notification () in
312 send_notification t
313 Messages.Resources.list_changed_notification_method
314 Messages.Resources.list_changed_notification_jsont
315 notif
316
317let send_tool_list_changed t =
318 let notif = Messages.Tools.make_list_changed_notification () in
319 send_notification t
320 Messages.Tools.list_changed_notification_method
321 Messages.Tools.list_changed_notification_jsont
322 notif
323
324let send_prompt_list_changed t =
325 let notif = Messages.Prompts.make_list_changed_notification () in
326 send_notification t
327 Messages.Prompts.list_changed_notification_method
328 Messages.Prompts.list_changed_notification_jsont
329 notif
330
331let send_roots_list_changed t =
332 let notif = Messages.Roots.make_list_changed_notification () in
333 send_notification t
334 Messages.Roots.list_changed_notification_method
335 Messages.Roots.list_changed_notification_jsont
336 notif
337
338let send_log_message t ~level ?logger ~data () =
339 let notif = Messages.Logging.make_notification ~level ?logger ~data () in
340 send_notification t
341 Messages.Logging.method_
342 Messages.Logging.notification_jsont
343 notif
344
345let send_progress t ~progress_token ~progress ?total () =
346 let notif = Messages.Progress.make_notification ~progress_token ~progress ?total () in
347 send_notification t
348 Messages.Progress.method_
349 Messages.Progress.notification_jsont
350 notif
351
352(** {1 Requesting from Client} *)
353
354let request_roots_list t =
355 match client_capabilities t with
356 | { roots = None; _ } -> None
357 | { roots = Some _; _ } ->
358 let req = Messages.Roots.make_list_request () in
359 let params = encode_json Messages.Roots.list_request_jsont req in
360 let result_json = Session.send_request t.session
361 ~method_:Messages.Roots.list_method
362 ~params
363 ()
364 in
365 let result = decode_json Messages.Roots.list_result_jsont result_json in
366 Some result
367
368(** {1 Session Management} *)
369
370let close t =
371 Session.close t.session