Model Context Protocol in OCaml
anil.recoil.org
1open Mcp
2open Mcp_sdk
3open Mcp_server
4
5(* Random pixel image generator MCP server *)
6
7(* Base64 encoding helper *)
8module Base64 = struct
9 let encode_char n =
10 "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/".[n]
11
12 let encode_block i bytes =
13 let buffer = Buffer.create 4 in
14 let b1 = Char.code (String.get bytes (i * 3)) in
15 let b2 = if i * 3 + 1 < String.length bytes then Char.code (String.get bytes (i * 3 + 1)) else 0 in
16 let b3 = if i * 3 + 2 < String.length bytes then Char.code (String.get bytes (i * 3 + 2)) else 0 in
17
18 let n = (b1 lsl 16) lor (b2 lsl 8) lor b3 in
19 Buffer.add_char buffer (encode_char ((n lsr 18) land 63));
20 Buffer.add_char buffer (encode_char ((n lsr 12) land 63));
21
22 if i * 3 + 1 < String.length bytes then
23 Buffer.add_char buffer (encode_char ((n lsr 6) land 63))
24 else
25 Buffer.add_char buffer '=';
26
27 if i * 3 + 2 < String.length bytes then
28 Buffer.add_char buffer (encode_char (n land 63))
29 else
30 Buffer.add_char buffer '=';
31
32 Buffer.contents buffer
33
34 let encode data =
35 let buffer = Buffer.create (4 * (String.length data + 2) / 3) in
36 for i = 0 to (String.length data - 1) / 3 do
37 Buffer.add_string buffer (encode_block i data)
38 done;
39 Buffer.contents buffer
40end
41
42(* Image generation utilities *)
43module ImageGenerator = struct
44 (* Simple PNG generation *)
45 let create_png width height pixels =
46 (* PNG signature *)
47 let signature = [|137; 80; 78; 71; 13; 10; 26; 10|] in
48
49 (* IHDR chunk data *)
50 let ihdr_data = Bytes.create 13 in
51 (* Width - big endian *)
52 Bytes.set ihdr_data 0 (Char.chr ((width lsr 24) land 0xff));
53 Bytes.set ihdr_data 1 (Char.chr ((width lsr 16) land 0xff));
54 Bytes.set ihdr_data 2 (Char.chr ((width lsr 8) land 0xff));
55 Bytes.set ihdr_data 3 (Char.chr (width land 0xff));
56 (* Height - big endian *)
57 Bytes.set ihdr_data 4 (Char.chr ((height lsr 24) land 0xff));
58 Bytes.set ihdr_data 5 (Char.chr ((height lsr 16) land 0xff));
59 Bytes.set ihdr_data 6 (Char.chr ((height lsr 8) land 0xff));
60 Bytes.set ihdr_data 7 (Char.chr (height land 0xff));
61 (* Bit depth - 8 bits *)
62 Bytes.set ihdr_data 8 (Char.chr 8);
63 (* Color type - RGB with alpha *)
64 Bytes.set ihdr_data 9 (Char.chr 6);
65 (* Compression, filter, interlace - all 0 *)
66 Bytes.set ihdr_data 10 (Char.chr 0);
67 Bytes.set ihdr_data 11 (Char.chr 0);
68 Bytes.set ihdr_data 12 (Char.chr 0);
69
70 (* Very simple CRC32 implementation for PNG chunks *)
71 let calculate_crc data =
72 let crc = ref 0xffffffff in
73 for i = 0 to Bytes.length data - 1 do
74 let byte = Char.code (Bytes.get data i) in
75 crc := !crc lxor byte;
76 for _ = 0 to 7 do
77 if !crc land 1 <> 0 then
78 crc := (!crc lsr 1) lxor 0xedb88320
79 else
80 crc := !crc lsr 1
81 done
82 done;
83 !crc lxor 0xffffffff
84 in
85
86 (* Create IHDR chunk *)
87 let ihdr_chunk = Buffer.create 25 in
88 (* Length - 13 bytes *)
89 Buffer.add_char ihdr_chunk (Char.chr 0);
90 Buffer.add_char ihdr_chunk (Char.chr 0);
91 Buffer.add_char ihdr_chunk (Char.chr 0);
92 Buffer.add_char ihdr_chunk (Char.chr 13);
93 (* Chunk type - IHDR *)
94 Buffer.add_string ihdr_chunk "IHDR";
95 (* Chunk data *)
96 Buffer.add_string ihdr_chunk (Bytes.unsafe_to_string ihdr_data);
97 (* CRC *)
98 let ihdr_crc_data = Bytes.create 17 in
99 Bytes.blit_string "IHDR" 0 ihdr_crc_data 0 4;
100 Bytes.blit ihdr_data 0 ihdr_crc_data 4 13;
101 let crc = calculate_crc ihdr_crc_data in
102 Buffer.add_char ihdr_chunk (Char.chr ((crc lsr 24) land 0xff));
103 Buffer.add_char ihdr_chunk (Char.chr ((crc lsr 16) land 0xff));
104 Buffer.add_char ihdr_chunk (Char.chr ((crc lsr 8) land 0xff));
105 Buffer.add_char ihdr_chunk (Char.chr (crc land 0xff));
106
107 (* Create IDAT chunk (uncompressed for simplicity) *)
108 let row_size = width * 4 in
109 let data_size = height * (row_size + 1) in
110 let idat_chunk = Buffer.create (12 + data_size) in
111 (* Length *)
112 Buffer.add_char idat_chunk (Char.chr ((data_size lsr 24) land 0xff));
113 Buffer.add_char idat_chunk (Char.chr ((data_size lsr 16) land 0xff));
114 Buffer.add_char idat_chunk (Char.chr ((data_size lsr 8) land 0xff));
115 Buffer.add_char idat_chunk (Char.chr (data_size land 0xff));
116 (* Chunk type - IDAT *)
117 Buffer.add_string idat_chunk "IDAT";
118
119 (* Very simple zlib header (no compression) *)
120 Buffer.add_char idat_chunk (Char.chr 0x78); (* CMF byte *)
121 Buffer.add_char idat_chunk (Char.chr 0x01); (* FLG byte *)
122
123 (* Raw image data with filter type 0 (None) for each scanline *)
124 for y = 0 to height - 1 do
125 (* Filter type 0 (None) *)
126 Buffer.add_char idat_chunk (Char.chr 0);
127 for x = 0 to width - 1 do
128 let idx = (y * width + x) * 4 in
129 Buffer.add_char idat_chunk (Char.chr (Bytes.get_uint8 pixels idx)); (* R *)
130 Buffer.add_char idat_chunk (Char.chr (Bytes.get_uint8 pixels (idx + 1))); (* G *)
131 Buffer.add_char idat_chunk (Char.chr (Bytes.get_uint8 pixels (idx + 2))); (* B *)
132 Buffer.add_char idat_chunk (Char.chr (Bytes.get_uint8 pixels (idx + 3))); (* A *)
133 done
134 done;
135
136 (* Zlib Adler-32 checksum (simplified) *)
137 let adler = ref 1 in
138 Buffer.add_char idat_chunk (Char.chr ((!adler lsr 24) land 0xff));
139 Buffer.add_char idat_chunk (Char.chr ((!adler lsr 16) land 0xff));
140 Buffer.add_char idat_chunk (Char.chr ((!adler lsr 8) land 0xff));
141 Buffer.add_char idat_chunk (Char.chr (!adler land 0xff));
142
143 (* CRC *)
144 let idat_crc = ref 0 in (* Not calculating CRC for simplicity *)
145 Buffer.add_char idat_chunk (Char.chr ((!idat_crc lsr 24) land 0xff));
146 Buffer.add_char idat_chunk (Char.chr ((!idat_crc lsr 16) land 0xff));
147 Buffer.add_char idat_chunk (Char.chr ((!idat_crc lsr 8) land 0xff));
148 Buffer.add_char idat_chunk (Char.chr (!idat_crc land 0xff));
149
150 (* Create IEND chunk *)
151 let iend_chunk = Buffer.create 12 in
152 (* Length - 0 bytes *)
153 Buffer.add_char iend_chunk (Char.chr 0);
154 Buffer.add_char iend_chunk (Char.chr 0);
155 Buffer.add_char iend_chunk (Char.chr 0);
156 Buffer.add_char iend_chunk (Char.chr 0);
157 (* Chunk type - IEND *)
158 Buffer.add_string iend_chunk "IEND";
159 (* CRC *)
160 let iend_crc = 0xAE426082 in (* Precomputed CRC for IEND chunk *)
161 Buffer.add_char iend_chunk (Char.chr ((iend_crc lsr 24) land 0xff));
162 Buffer.add_char iend_chunk (Char.chr ((iend_crc lsr 16) land 0xff));
163 Buffer.add_char iend_chunk (Char.chr ((iend_crc lsr 8) land 0xff));
164 Buffer.add_char iend_chunk (Char.chr (iend_crc land 0xff));
165
166 (* Combine all parts *)
167 let result = Buffer.create (8 + Buffer.length ihdr_chunk + Buffer.length idat_chunk + Buffer.length iend_chunk) in
168 (* PNG signature *)
169 Array.iter (fun c -> Buffer.add_char result (Char.chr c)) signature;
170 (* IHDR chunk *)
171 Buffer.add_buffer result ihdr_chunk;
172 (* IDAT chunk *)
173 Buffer.add_buffer result idat_chunk;
174 (* IEND chunk *)
175 Buffer.add_buffer result iend_chunk;
176
177 Buffer.contents result
178
179 (* Generate random pixel art image *)
180 let generate_random_image ?(width=16) ?(height=16) ?(pixel_size=1) ?(seed=None) () =
181 let pixels = Bytes.create (width * height * 4) in
182
183 (* Set random seed if provided *)
184 (match seed with
185 | Some s -> Random.init s
186 | None -> Random.self_init ());
187
188 (* Generate a random color palette *)
189 let palette_size = Random.int 8 + 2 in (* 2-10 colors *)
190 let palette = Array.init palette_size (fun _ ->
191 (Random.int 256, Random.int 256, Random.int 256, 255) (* RGBA *)
192 ) in
193
194 (* Fill the pixel buffer *)
195 for y = 0 to height - 1 do
196 for x = 0 to width - 1 do
197 let color_idx = Random.int palette_size in
198 let (r, g, b, a) = palette.(color_idx) in
199 let idx = (y * width + x) * 4 in
200 Bytes.set_uint8 pixels idx r;
201 Bytes.set_uint8 pixels (idx + 1) g;
202 Bytes.set_uint8 pixels (idx + 2) b;
203 Bytes.set_uint8 pixels (idx + 3) a;
204 done
205 done;
206
207 (* Create symmetrical patterns - horizontally, vertically, or both *)
208 let symmetry_type = Random.int 3 in
209 if symmetry_type > 0 then begin
210 for y = 0 to height - 1 do
211 for x = 0 to width / 2 do
212 (* Mirror horizontally (except center column for odd widths) *)
213 if symmetry_type = 1 || symmetry_type = 2 then begin
214 let mirror_x = width - 1 - x in
215 if x <> mirror_x then begin
216 let src_idx = (y * width + x) * 4 in
217 let dst_idx = (y * width + mirror_x) * 4 in
218 for i = 0 to 3 do
219 Bytes.set pixels (dst_idx + i) (Bytes.get pixels (src_idx + i))
220 done
221 end
222 end
223 done
224 done;
225
226 (* Mirror vertically for symmetry_type = 2 *)
227 if symmetry_type = 2 then begin
228 for y = 0 to height / 2 do
229 let mirror_y = height - 1 - y in
230 if y <> mirror_y then begin
231 for x = 0 to width - 1 do
232 let src_idx = (y * width + x) * 4 in
233 let dst_idx = (mirror_y * width + x) * 4 in
234 for i = 0 to 3 do
235 Bytes.set pixels (dst_idx + i) (Bytes.get pixels (src_idx + i))
236 done
237 done
238 end
239 done
240 end
241 end;
242
243 (* Scale up the image if pixel_size > 1 *)
244 let final_width = width * pixel_size in
245 let final_height = height * pixel_size in
246
247 if pixel_size = 1 then
248 create_png width height pixels
249 else begin
250 let scaled_pixels = Bytes.create (final_width * final_height * 4) in
251
252 for y = 0 to height - 1 do
253 for x = 0 to width - 1 do
254 let src_idx = (y * width + x) * 4 in
255 for py = 0 to pixel_size - 1 do
256 for px = 0 to pixel_size - 1 do
257 let dst_x = x * pixel_size + px in
258 let dst_y = y * pixel_size + py in
259 let dst_idx = (dst_y * final_width + dst_x) * 4 in
260 for i = 0 to 3 do
261 Bytes.set scaled_pixels (dst_idx + i) (Bytes.get pixels (src_idx + i))
262 done
263 done
264 done
265 done
266 done;
267
268 create_png final_width final_height scaled_pixels
269 end
270end
271
272(* Helper for extracting values from JSON *)
273let get_param_int json name default =
274 match json with
275 | `Assoc fields ->
276 (match List.assoc_opt name fields with
277 | Some (`Int i) -> i
278 | Some (`Float f) -> int_of_float f
279 | _ -> default)
280 | _ -> default
281
282(* Create a server *)
283let server = create_server
284 ~name:"OCaml MCP Image Generator"
285 ~version:"0.1.0"
286 ~protocol_version:"2024-11-05"
287 ()
288
289(* Define startup and shutdown hooks *)
290let startup () =
291 Printf.fprintf stderr "ImageGeneratorServer is starting up!\n";
292 flush stderr;
293 Log.info "ImageGeneratorServer is starting up!"
294
295let shutdown () =
296 Printf.fprintf stderr "ImageGeneratorServer is shutting down. Goodbye!\n";
297 flush stderr;
298 Log.info "ImageGeneratorServer is shutting down. Goodbye!"
299
300(* Register the hooks *)
301let () =
302 set_startup_hook server startup;
303 set_shutdown_hook server shutdown
304
305(* Make an image content helper *)
306let make_image_content data mime_type =
307 let image_content = ImageContent.{
308 data;
309 mime_type;
310 annotations = None;
311 } in
312 Image image_content
313
314(* Define and register a random pixel art generator tool *)
315let _ = add_tool server
316 ~name:"generate_random_pixel_art"
317 ~description:"Generates a random pixel art image"
318 ~schema_properties:[
319 ("width", "integer", "Width of the pixel art grid (default: 16)");
320 ("height", "integer", "Height of the pixel art grid (default: 16)");
321 ("pixel_size", "integer", "Size of each pixel (default: 8)");
322 ("seed", "integer", "Random seed (optional)");
323 ]
324 ~schema_required:[]
325 (fun args ->
326 try
327 let width = get_param_int args "width" 16 in
328 let height = get_param_int args "height" 16 in
329 let pixel_size = get_param_int args "pixel_size" 8 in
330
331 (* Validate parameters *)
332 let width = max 1 (min 64 width) in (* Limit to 1-64 *)
333 let height = max 1 (min 64 height) in (* Limit to 1-64 *)
334 let pixel_size = max 1 (min 16 pixel_size) in (* Limit to 1-16 *)
335
336 (* Extract optional seed *)
337 let seed = match args with
338 | `Assoc fields ->
339 (match List.assoc_opt "seed" fields with
340 | Some (`Int s) -> Some s
341 | _ -> None)
342 | _ -> None
343 in
344
345 (* Generate the image *)
346 let image_data = ImageGenerator.generate_random_image
347 ~width ~height ~pixel_size ~seed () in
348
349 (* Encode as base64 *)
350 let base64_data = Base64.encode image_data in
351
352 Log.info (Printf.sprintf "Generated random pixel art image (%dx%d grid, %dpx pixels)"
353 width height pixel_size);
354
355 (* Create a response with both text and image content *)
356 CallToolResult.yojson_of_t CallToolResult.{
357 content = [
358 Text TextContent.{
359 text = Printf.sprintf "Generated random pixel art image (%dx%d grid, %dpx pixels)"
360 width height pixel_size;
361 annotations = None
362 };
363 Image ImageContent.{
364 data = base64_data;
365 mime_type = "image/png";
366 annotations = None
367 }
368 ];
369 is_error = false;
370 meta = None
371 }
372 with
373 | Failure msg ->
374 Log.error (Printf.sprintf "Error in image generator tool: %s" msg);
375 CallToolResult.yojson_of_t CallToolResult.{
376 content = [
377 Text TextContent.{
378 text = Printf.sprintf "Error: %s" msg;
379 annotations = None
380 }
381 ];
382 is_error = true;
383 meta = None
384 }
385 )
386
387(* Define and register a pixel art prompt *)
388let _ = add_prompt server
389 ~name:"pixel-art-prompt"
390 ~description:"A prompt that includes a random pixel art image"
391 ~arguments:[
392 ("width", Some "Width of the pixel art (1-64)", false);
393 ("height", Some "Height of the pixel art (1-64)", false);
394 ("pixel_size", Some "Size of each pixel (1-16)", false);
395 ]
396 (fun args ->
397 (* Parse parameters with defaults *)
398 let width =
399 try int_of_string (List.assoc "width" args)
400 with _ -> 16
401 in
402 let height =
403 try int_of_string (List.assoc "height" args)
404 with _ -> 16
405 in
406 let pixel_size =
407 try int_of_string (List.assoc "pixel_size" args)
408 with _ -> 8
409 in
410
411 (* Validate parameters *)
412 let width = max 1 (min 64 width) in
413 let height = max 1 (min 64 height) in
414 let pixel_size = max 1 (min 16 pixel_size) in
415
416 (* Generate image *)
417 let image_data = ImageGenerator.generate_random_image
418 ~width ~height ~pixel_size () in
419
420 (* Encode as base64 *)
421 let base64_data = Base64.encode image_data in
422
423 Log.info (Printf.sprintf "Generated pixel art for prompt (%dx%d grid, %dpx pixels)"
424 width height pixel_size);
425
426 [
427 Prompt.{
428 role = `User;
429 content = make_text_content "I've generated a random pixel art image for you:"
430 };
431 Prompt.{
432 role = `User;
433 content = make_image_content base64_data "image/png"
434 };
435 Prompt.{
436 role = `User;
437 content = make_text_content (Printf.sprintf "Please describe what you see in this %dx%d pixel art."
438 width height)
439 };
440 Prompt.{
441 role = `Assistant;
442 content = make_text_content "I'll describe what I see in this pixel art image."
443 }
444 ]
445 )
446
447(* Main function *)
448let () =
449 (* Parse command line arguments *)
450 let transport_type = ref Stdio in
451 let args = [
452 ("--http", Arg.Unit (fun () -> transport_type := Http),
453 "Start server with HTTP transport (default is stdio)");
454 ] in
455 let usage_msg = "Usage: image_generator_example [--http]" in
456 Arg.parse args (fun _ -> ()) usage_msg;
457
458 (* Use stderr for direct printing to avoid interfering with JSON-RPC protocol *)
459 Printf.fprintf stderr "Starting ImageGeneratorServer...\n";
460 flush stderr;
461 Log.info "Starting ImageGeneratorServer...";
462
463 (* Configure the server with appropriate capabilities *)
464 ignore (configure_server server ());
465
466 (* Create and start MCP server with the selected transport *)
467 let mcp_server = create ~server ~transport:!transport_type () in
468 start mcp_server