coop/rose-ash
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix JIT compilation cascade + MCP robustness

Browse Source 
Three interacting JIT bugs caused infinite loops and server hangs:

1. _jit_compiling cascade: the re-entrancy flag was local to each
   binary's hook. When vm_call triggered JIT compilation internally,
   compiler functions got JIT-compiled during compilation, creating
   infinite cascades. Fix: shared _jit_compiling flag in sx_vm.ml,
   set in jit_compile_lambda itself.

2. call_closure always created new VMs: every HO primitive callback
   (for-each, map, filter) allocated a fresh VM. With 43K+ calls
   during compilation, this was the direct cause of hangs. Fix:
   call_closure_reuse reuses the active VM by isolating frames and
   running re-entrantly. VmSuspended is handled by merging frames
   for proper IO resumption.

3. vm_call for compiled Lambdas: OP_CALL dispatching to a Lambda
   with cached bytecode created a new VM instead of pushing a frame
   on the current one. Fix: push_closure_frame directly.

Additional MCP server fixes:
- Hot-reload: auto-execv when binary on disk is newer (no restart needed)
- Robust JSON: to_int_safe/to_int_or handle null, string, int params
- sx_summarise depth now optional (default 2)
- Per-request error handling (malformed JSON doesn't crash server)
- sx_test uses pre-built binary (skips dune rebuild overhead)
- Timed module loading for startup diagnostics

sx_server.ml fixes:
- Uses shared _jit_compiling flag
- Marks lambdas as jit_failed_sentinel on compile failure (no retry spam)
- call_closure_reuse with VmSuspended frame merging for IO support

Compiled compiler bytecode bug: deeply nested cond/case/let forms
(e.g. tw-resolve-style) cause the compiled compiler to loop.
Workaround: _jit_compiling guard prevents compiled function execution
during compilation. Compilation uses CEK (slower but correct).
Test suite: 3127/3127 passed.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
giles
2026-04-07 21:59:31 +00:00
parent 75130876c7
commit 03278c640d
3 changed files with 193 additions and 124 deletions
Download Patch File Download Diff File Expand all files Collapse all files

220
hosts/ocaml/bin/mcp_tree.ml
View File

@@ -6,6 +6,41 @@
open Sx_types
(* ------------------------------------------------------------------ *)
(* Hot-reload: re-exec ourselves when the binary has been rebuilt *)
(* ------------------------------------------------------------------ *)
let exe_path =
try Unix.readlink "/proc/self/exe"
with _ -> Sys.executable_name
let binary_mtime =
ref (try (Unix.stat exe_path).Unix.st_mtime with _ -> 0.0)
let check_hot_reload () =
try
let cur = (Unix.stat exe_path).Unix.st_mtime in
if cur > !binary_mtime then begin
Printf.eprintf "[mcp] Binary updated (%.0f -> %.0f), hot-reloading...\n%!" !binary_mtime cur;
Unix.execv exe_path [| exe_path |]
end
with _ -> ()
(* ------------------------------------------------------------------ *)
(* Robust JSON helpers — MCP clients send ints as strings or null *)
(* ------------------------------------------------------------------ *)
let to_int_safe json =
match json with
| `Int n -> Some n
| `Float f -> Some (int_of_float f)
| `String s -> (try Some (int_of_string s) with _ -> None)
| `Null -> None
| _ -> None
let to_int_or ~default json =
match to_int_safe json with Some n -> n | None -> default
(* ------------------------------------------------------------------ *)
(* SX evaluator setup — minimal env for parser + tree-tools *)
(* ------------------------------------------------------------------ *)
@@ -29,44 +64,44 @@ let load_sx_file e path =
(* JIT infrastructure — shared VM globals table, kept in sync via env_bind hook *)
let _mcp_vm_globals : (string, value) Hashtbl.t = Hashtbl.create 2048
let _jit_compiling = ref false
let _jit_warned : (string, bool) Hashtbl.t = Hashtbl.create 32
let register_mcp_jit_hook () =
Sx_runtime._jit_try_call_fn := Some (fun f args ->
match f with
| Lambda l ->
let fn_name = match l.l_name with Some n -> n | None -> "" in
if fn_name <> "" && Hashtbl.mem _jit_warned fn_name then None
else
(match l.l_compiled with
| Some cl when not (Sx_vm.is_jit_failed cl) ->
(try Some (Sx_vm.call_closure cl args cl.vm_env_ref)
(* Already compiled — run via VM. Skip during compilation. *)
if !(Sx_vm._jit_compiling) then None
else
(try Some (Sx_vm.call_closure_reuse cl args)
with e ->
let fn_name = match l.l_name with Some n -> n | None -> "?" in
if not (Hashtbl.mem _jit_warned fn_name) then begin
Hashtbl.replace _jit_warned fn_name true;
Printf.eprintf "[mcp-jit] %s runtime fallback to CEK: %s\n%!" fn_name (Printexc.to_string e)
end;
l.l_compiled <- Some Sx_vm.jit_failed_sentinel;
None)
| Some _ -> None
| None ->
if !_jit_compiling then None
else begin
let fn_name = match l.l_name with Some n -> n | None -> "?" in
if Hashtbl.mem _jit_warned fn_name then None
else begin
_jit_compiling := true;
let compiled = Sx_vm.jit_compile_lambda l _mcp_vm_globals in
_jit_compiling := false;
match compiled with
| Some cl ->
l.l_compiled <- Some cl;
(try Some (Sx_vm.call_closure cl args cl.vm_env_ref)
with e ->
Printf.eprintf "[mcp-jit] %s first-call fallback: %s\n%!" fn_name (Printexc.to_string e);
Hashtbl.replace _jit_warned fn_name true;
None)
| None -> None
end
end)
(* Only block NEW compilations during _jit_compiling *)
if !(Sx_vm._jit_compiling) then None
else
let compiled = Sx_vm.jit_compile_lambda l _mcp_vm_globals in
(match compiled with
| Some cl ->
l.l_compiled <- Some cl;
(try Some (Sx_vm.call_closure_reuse cl args)
with e ->
Printf.eprintf "[mcp-jit] %s first-call fallback: %s\n%!" fn_name (Printexc.to_string e);
Hashtbl.replace _jit_warned fn_name true;
l.l_compiled <- Some Sx_vm.jit_failed_sentinel;
None)
| None -> None))
| _ -> None)
let setup_env () =
@@ -77,6 +112,13 @@ let setup_env () =
match v with
| Thunk (body, closure_env) -> Sx_ref.eval_expr body (Env closure_env)
| other -> other);
(* Seed VM globals with all primitives as NativeFn values.
Without this, OP_CALL_PRIM in JIT-compiled bytecode can't find
primitives and falls back to CEK — which causes cascading failures
for stdlib functions like nth/first that override native versions. *)
Hashtbl.iter (fun name fn ->
Hashtbl.replace _mcp_vm_globals name (NativeFn (name, fn))
) Sx_primitives.primitives;
(* JIT: mirror root-env bindings into shared VM globals table *)
Sx_types._env_bind_hook := Some (fun env name v ->
if env.parent = None then
@@ -304,51 +346,49 @@ let setup_env () =
| _ -> Nil);
bind "deftype" (fun _args -> Nil);
bind "defeffect" (fun _args -> Nil);
(* Load parser.sx for the SX-level sx-parse/sx-serialize *)
(* Load SX modules with timing *)
let spec_dir = try Sys.getenv "SX_SPEC_DIR" with Not_found -> "spec" in
let lib_dir = try Sys.getenv "SX_LIB_DIR" with Not_found -> "lib" in
(try load_sx_file e (Filename.concat spec_dir "parser.sx")
with exn -> Printf.eprintf "[mcp] Warning: parser.sx load failed: %s\n%!" (Printexc.to_string exn));
(* Load tree-tools *)
(try load_sx_file e (Filename.concat lib_dir "tree-tools.sx")
with exn -> Printf.eprintf "[mcp] Error: tree-tools.sx load failed: %s\n%!" (Printexc.to_string exn); exit 1);
(* Load signals — reactive signal primitives *)
(try load_sx_file e (Filename.concat spec_dir "signals.sx")
with exn -> Printf.eprintf "[mcp] Warning: signals.sx load failed: %s\n%!" (Printexc.to_string exn));
(* Load render + adapter-html for render-to-html *)
let web_dir = try Sys.getenv "SX_WEB_DIR" with Not_found -> "web" in
(try load_sx_file e (Filename.concat spec_dir "render.sx")
with exn -> Printf.eprintf "[mcp] Warning: render.sx load failed: %s\n%!" (Printexc.to_string exn));
(try load_sx_file e (Filename.concat web_dir "adapter-html.sx")
with exn -> Printf.eprintf "[mcp] Warning: adapter-html.sx load failed: %s\n%!" (Printexc.to_string exn));
(* Load harness *)
(try load_sx_file e (Filename.concat spec_dir "harness.sx")
with exn -> Printf.eprintf "[mcp] Warning: harness.sx load failed: %s\n%!" (Printexc.to_string exn));
(* Load eval-rules *)
(try load_sx_file e (Filename.concat spec_dir "eval-rules.sx")
with exn -> Printf.eprintf "[mcp] Warning: eval-rules.sx load failed: %s\n%!" (Printexc.to_string exn));
(* Load signals — reactive signal primitives *)
(try load_sx_file e (Filename.concat spec_dir "signals.sx")
with exn -> Printf.eprintf "[mcp] Warning: signals.sx load failed: %s\n%!" (Printexc.to_string exn));
(* Load render + adapter-html for render-to-html *)
let web_dir = try Sys.getenv "SX_WEB_DIR" with Not_found -> "web" in
(try load_sx_file e (Filename.concat spec_dir "render.sx")
with exn -> Printf.eprintf "[mcp] Warning: render.sx load failed: %s\n%!" (Printexc.to_string exn));
(try load_sx_file e (Filename.concat web_dir "adapter-html.sx")
with exn -> Printf.eprintf "[mcp] Warning: adapter-html.sx load failed: %s\n%!" (Printexc.to_string exn));
(* Load render pipeline — native OCaml renderer + HTML tag bindings *)
let t0 = Unix.gettimeofday () in
let timed name f =
let t = Unix.gettimeofday () in
f ();
Printf.eprintf "[mcp] %s: %.0fms\n%!" name ((Unix.gettimeofday () -. t) *. 1000.0)
in
timed "parser.sx" (fun () ->
try load_sx_file e (Filename.concat spec_dir "parser.sx")
with exn -> Printf.eprintf "[mcp] Warning: parser.sx load failed: %s\n%!" (Printexc.to_string exn));
timed "tree-tools.sx" (fun () ->
try load_sx_file e (Filename.concat lib_dir "tree-tools.sx")
with exn -> Printf.eprintf "[mcp] Error: tree-tools.sx load failed: %s\n%!" (Printexc.to_string exn); exit 1);
timed "signals.sx" (fun () ->
try load_sx_file e (Filename.concat spec_dir "signals.sx")
with exn -> Printf.eprintf "[mcp] Warning: signals.sx load failed: %s\n%!" (Printexc.to_string exn));
timed "render.sx" (fun () ->
try load_sx_file e (Filename.concat spec_dir "render.sx")
with exn -> Printf.eprintf "[mcp] Warning: render.sx load failed: %s\n%!" (Printexc.to_string exn));
timed "adapter-html.sx" (fun () ->
try load_sx_file e (Filename.concat web_dir "adapter-html.sx")
with exn -> Printf.eprintf "[mcp] Warning: adapter-html.sx load failed: %s\n%!" (Printexc.to_string exn));
timed "harness.sx" (fun () ->
try load_sx_file e (Filename.concat spec_dir "harness.sx")
with exn -> Printf.eprintf "[mcp] Warning: harness.sx load failed: %s\n%!" (Printexc.to_string exn));
timed "eval-rules.sx" (fun () ->
try load_sx_file e (Filename.concat spec_dir "eval-rules.sx")
with exn -> Printf.eprintf "[mcp] Warning: eval-rules.sx load failed: %s\n%!" (Printexc.to_string exn));
(* Render pipeline — native OCaml renderer + HTML tag bindings *)
Sx_render.setup_render_env e;
List.iter (fun tag ->
ignore (Sx_types.env_bind e tag
(NativeFn ("html:" ^ tag, fun args -> List (Symbol tag :: args))))
) Sx_render.html_tags;
ignore (Sx_types.env_bind e "island?" (NativeFn ("island?", fun args -> match args with [Island _] -> Bool true | _ -> Bool false)));
(* Load compiler + enable JIT *)
(try load_sx_file e (Filename.concat lib_dir "compiler.sx");
register_mcp_jit_hook ();
Printf.eprintf "[mcp] JIT enabled (compiler.sx loaded)\n%!"
with exn -> Printf.eprintf "[mcp] Warning: compiler.sx load failed (JIT disabled): %s\n%!" (Printexc.to_string exn));
Printf.eprintf "[mcp] SX tree-tools + harness + eval-rules + render loaded\n%!";
timed "compiler.sx+JIT" (fun () ->
try load_sx_file e (Filename.concat lib_dir "compiler.sx");
register_mcp_jit_hook ()
with exn -> Printf.eprintf "[mcp] Warning: compiler.sx load failed (JIT disabled): %s\n%!" (Printexc.to_string exn));
Printf.eprintf "[mcp] Ready in %.0fms\n%!" ((Unix.gettimeofday () -. t0) *. 1000.0);
env := e
(* ------------------------------------------------------------------ *)
@@ -616,9 +656,9 @@ let handle_tool name args =
let file = args |> member "file" |> to_string in
let tree, cst = parse_file_cst file in
let focus = args |> member "focus" |> to_string_option in
let max_depth = args |> member "max_depth" |> to_int_option in
let max_lines = args |> member "max_lines" |> to_int_option in
let offset = args |> member "offset" |> to_int_option |> Option.value ~default:0 in
let max_depth = to_int_safe (args |> member "max_depth") in
let max_lines = to_int_safe (args |> member "max_lines") in
let offset = to_int_or ~default:0 (args |> member "offset") in
(match focus with
| Some pattern ->
text_result (inject_cst_comments (value_to_string (call_sx "annotate-focused" [tree; String pattern])) (extract_cst_comments cst))
@@ -642,7 +682,7 @@ let handle_tool name args =
| "sx_summarise" ->
let file = args |> member "file" |> to_string in
let tree, cst = parse_file_cst file in
let depth = args |> member "depth" |> to_int in
let depth = to_int_or ~default:2 (args |> member "depth") in
text_result (inject_cst_comments (value_to_string (call_sx "summarise" [tree; Number (float_of_int depth)])) (extract_cst_comments cst))
| "sx_read_subtree" ->
@@ -697,7 +737,7 @@ let handle_tool name args =
let file = args |> member "file" |> to_string in
let tree, cst = parse_file_cst file in
let path = resolve_path tree (args |> member "path" |> to_string) in
let index = args |> member "index" |> to_int in
let index = to_int_or ~default:0 (args |> member "index") in
let src = args |> member "new_source" |> to_string in
write_edit_cst file cst (call_sx "insert-child" [tree; path; Number (float_of_int index); String src])
@@ -873,11 +913,18 @@ let handle_tool name args =
let spec_dir = try Sys.getenv "SX_SPEC_DIR" with Not_found -> "spec" in
Filename.dirname spec_dir
in
let timeout = args |> member "timeout" |> to_int_option |> Option.value ~default:300 in
let timeout = to_int_or ~default:300 (args |> member "timeout") in
let cmd = match host with
| "ocaml" ->
Printf.sprintf "cd %s/hosts/ocaml && eval $(opam env 2>/dev/null) && timeout %d dune exec bin/run_tests.exe%s 2>&1"
project_dir timeout (if full then " -- --full" else "")
(* Use pre-built binary directly — avoids dune rebuild delay.
Falls back to dune exec if the binary doesn't exist. *)
let exe = Printf.sprintf "%s/hosts/ocaml/_build/default/bin/run_tests.exe" project_dir in
if Sys.file_exists exe then
Printf.sprintf "cd %s/hosts/ocaml && timeout %d %s%s 2>&1"
project_dir timeout exe (if full then " --full" else "")
else
Printf.sprintf "cd %s/hosts/ocaml && eval $(opam env 2>/dev/null) && timeout %d dune exec bin/run_tests.exe%s 2>&1"
project_dir timeout (if full then " -- --full" else "")
| "js" | _ ->
Printf.sprintf "cd %s && timeout %d node hosts/javascript/run_tests.js%s 2>&1"
project_dir timeout (if full then " --full" else "")
@@ -1827,7 +1874,7 @@ let handle_tool name args =
| "sx_trace" ->
let expr_str = args |> member "expr" |> to_string in
let max_steps = (try args |> member "max_steps" |> to_int with _ -> 200) in
let max_steps = to_int_or ~default:200 (args |> member "max_steps") in
let file = try Some (args |> member "file" |> to_string) with _ -> None in
let components_only = (try args |> member "components_only" |> to_bool with _ -> false) in
let e = !env in
@@ -2364,7 +2411,7 @@ let tool_definitions = `List [
("offset", `Assoc [("type", `String "integer"); ("description", `String "Line offset for pagination (default 0)")])]
["file"];
tool "sx_summarise" "Folded structural overview of an .sx file. Use to orient before drilling into a region."
[file_prop; ("depth", `Assoc [("type", `String "integer"); ("description", `String "Max depth (0=heads only)")])] ["file"; "depth"];
[file_prop; ("depth", `Assoc [("type", `String "integer"); ("description", `String "Max depth (0=heads only, default 2)")])] ["file"];
tool "sx_read_subtree" "Expand a specific subtree by path. Use after summarise to drill in."
[file_prop; path_prop] ["file"; "path"];
tool "sx_get_context" "Show enclosing chain from root to a target node."
@@ -2542,21 +2589,28 @@ let () =
try while true do
let line = input_line stdin in
if String.length line > 0 then begin
let json = Yojson.Safe.from_string line in
let open Yojson.Safe.Util in
let meth = json |> member "method" |> to_string_option |> Option.value ~default:"" in
let params = json |> member "params" in
let id = json |> member "id" in
let result = dispatch meth params in
if id <> `Null then begin
let resp = `Assoc [
("jsonrpc", `String "2.0");
("id", id);
("result", result)] in
print_string (Yojson.Safe.to_string resp);
print_char '\n';
flush stdout
end
try
let json = Yojson.Safe.from_string line in
let open Yojson.Safe.Util in
let meth = json |> member "method" |> to_string_option |> Option.value ~default:"" in
let params = json |> member "params" in
let id = json |> member "id" in
let result =
try dispatch meth params
with e -> error_result ("Error: " ^ Printexc.to_string e)
in
if id <> `Null then begin
let resp = `Assoc [
("jsonrpc", `String "2.0");
("id", id);
("result", result)] in
print_string (Yojson.Safe.to_string resp);
print_char '\n';
flush stdout;
check_hot_reload ()
end
with e ->
Printf.eprintf "[mcp] Request error: %s\n%!" (Printexc.to_string e)
end
done
with End_of_file -> ()

29
hosts/ocaml/bin/sx_server.ml
View File

@@ -908,7 +908,7 @@ let sx_render_to_html expr env =
into the kernel env. The hook handles both cached execution (bytecode
already compiled) and first-call compilation (invoke compiler.sx via
CEK, cache result). cek_call checks this before CEK dispatch. *)
let _jit_compiling = ref false (* re-entrancy guard *)
(* Re-entrancy guard lives in Sx_vm._jit_compiling (shared with vm_call path) *)
(* JIT compilation is lazy-only: every named lambda gets one compile
attempt on first call. Failures are sentineled (never retried). *)
@@ -936,9 +936,11 @@ let register_jit_hook env =
| Lambda l ->
(match l.l_compiled with
| Some cl when not (Sx_vm.is_jit_failed cl) ->
(* Cached bytecode — run on VM, fall back to CEK on runtime error.
Log once per function name, then stay quiet. Don't disable. *)
(try Some (Sx_vm.call_closure cl args cl.vm_env_ref)
(* Cached bytecode — run via VM. Skip during compilation to avoid
compiled-compiler bytecode bugs on complex nested forms. *)
if !(Sx_vm._jit_compiling) then None
else
(try Some (Sx_vm.call_closure_reuse cl args)
with
| Sx_vm.VmSuspended (request, saved_vm) ->
Some (make_vm_suspend_marker request saved_vm)
@@ -948,33 +950,34 @@ let register_jit_hook env =
Hashtbl.replace _jit_warned fn_name true;
Printf.eprintf "[jit] %s runtime fallback to CEK: %s\n%!" fn_name (Printexc.to_string e)
end;
l.l_compiled <- Some Sx_vm.jit_failed_sentinel;
None)
| Some _ -> None (* compile failed or disabled — CEK handles *)
| Some _ -> None
| None ->
(* Only block NEW compilations during _jit_compiling, not execution *)
let fn_name = match l.l_name with Some n -> n | None -> "?" in
if !_jit_compiling then None
if !(Sx_vm._jit_compiling) then None
else if Hashtbl.mem _jit_warned fn_name then None
else begin
_jit_compiling := true;
let t0 = Unix.gettimeofday () in
let compiled = Sx_vm.jit_compile_lambda l (env_to_vm_globals env) in
let dt = Unix.gettimeofday () -. t0 in
_jit_compiling := false;
if dt > 0.5 || (match compiled with None -> true | _ -> false) then
Printf.eprintf "[jit] %s compile %s in %.3fs\n%!"
fn_name (match compiled with Some _ -> "OK" | None -> "FAIL") dt;
Printf.eprintf "[jit] %s compile in %.3fs\n%!" fn_name dt;
match compiled with
| Some cl ->
l.l_compiled <- Some cl;
(try Some (Sx_vm.call_closure cl args cl.vm_env_ref)
(try Some (Sx_vm.call_closure_reuse cl args)
with
| Sx_vm.VmSuspended (request, saved_vm) ->
Some (make_vm_suspend_marker request saved_vm)
| e ->
Printf.eprintf "[jit] %s first-call fallback to CEK: %s\n%!" fn_name (Printexc.to_string e);
Hashtbl.replace _jit_warned fn_name true;
l.l_compiled <- Some Sx_vm.jit_failed_sentinel;
None)
| None -> None
| None ->
l.l_compiled <- Some Sx_vm.jit_failed_sentinel;
None
end)
| _ -> None)