Source file headless_json_runner.ml

(******************************************************************************)
(*                                                                            *)
(* SPDX-License-Identifier: MIT                                               *)
(* Copyright (c) 2026 Mathias Bourgoin <mathias.bourgoin@atacama.tech>        *)
(*                                                                            *)
(******************************************************************************)

(** Headless JSON-over-stdio runner.  Activated when [MIAOU_DRIVER=headless].
    Reads newline-delimited JSON commands from stdin and writes JSON responses
    to stdout.  The app binary becomes a subprocess that an AI agent or CI
    script can drive without a real terminal. *)

[@@@warning "-32-34-37-69"]

module HD = Lib_miaou_internal.Headless_driver
module Tui_page = Miaou_core.Tui_page

(* ── ANSI strip ─────────────────────────────────────────────────────────── *)

(** Remove ANSI/VT escape sequences from [s], returning plain text. *)
let ansi_strip s =
  let buf = Buffer.create (String.length s) in
  let n = String.length s in
  let i = ref 0 in
  while !i < n do
    if s.[!i] = '\x1b' then (
      incr i ;
      if !i < n && s.[!i] = '[' then (
        (* CSI sequence: ESC [ ... <final byte 0x40-0x7e> *)
        incr i ;
        while !i < n && (s.[!i] < '@' || s.[!i] > '~') do
          incr i
        done ;
        if !i < n then incr i (* consume final byte *))
      else if !i < n then incr i (* skip single char after ESC (Fe sequence) *))
    else (
      Buffer.add_char buf s.[!i] ;
      incr i)
  done ;
  Buffer.contents buf

(* ── on_frame callback ──────────────────────────────────────────────────── *)

(** Optional callback invoked with the raw ANSI frame on every frame emit.
    Set by [run ~on_frame] before the command loop starts.  The callback
    receives the terminal dimensions (rows, cols) followed by the raw ANSI
    data so that a viewer can resize its terminal to match. *)
let on_frame_fn : (rows:int -> cols:int -> string -> unit) option ref = ref None

(* ── Frame helpers ───────────────────────────────────────────────────────── *)

let current_frame () =
  let size = HD.get_size () in
  let raw = HD.Screen.get () in
  (match !on_frame_fn with
  | Some f -> f ~rows:size.LTerm_geom.rows ~cols:size.LTerm_geom.cols raw
  | None -> ()) ;
  let text = ansi_strip raw in
  `Assoc
    [
      ("type", `String "frame");
      ("text", `String text);
      ("rows", `Int size.LTerm_geom.rows);
      ("cols", `Int size.LTerm_geom.cols);
    ]

let nav_response action =
  `Assoc [("type", `String "nav"); ("action", `String action)]

let error_response msg =
  `Assoc [("type", `String "error"); ("message", `String msg)]

let emit json =
  (print_string [@allow_forbidden "headless runner writes JSON to stdout"])
    (Yojson.Safe.to_string json) ;
  (print_char [@allow_forbidden "headless runner writes JSON to stdout"]) '\n' ;
  (flush [@allow_forbidden "headless runner flushes stdout"]) stdout

(* ── Command dispatch ────────────────────────────────────────────────────── *)

(** Process one parsed command.  Returns [true] to keep running, [false] to
    stop the loop. *)
let handle_cmd (cmd : (string * Yojson.Safe.t) list) : bool =
  let get_string key =
    match List.assoc_opt key cmd with Some (`String s) -> Some s | _ -> None
  in
  let get_int key =
    match List.assoc_opt key cmd with Some (`Int n) -> Some n | _ -> None
  in
  match get_string "cmd" with
  | None ->
      emit (error_response "Missing 'cmd' field") ;
      true
  | Some "render" ->
      emit (current_frame ()) ;
      true
  | Some "key" -> (
      match get_string "key" with
      | None ->
          emit (error_response "Missing 'key' field") ;
          true
      | Some k -> (
          let outcome = HD.Stateful.send_key k in
          (* Run a few idle ticks so timers/refresh settle *)
          let outcome =
            match outcome with
            | `Continue -> HD.Stateful.idle_wait ~iterations:3 ()
            | other -> other
          in
          match outcome with
          | `Quit ->
              emit (nav_response "quit") ;
              false
          | `Back ->
              emit (nav_response "back") ;
              false
          | `SwitchTo name ->
              emit (nav_response ("switch:" ^ name)) ;
              true
          | `Continue ->
              emit (current_frame ()) ;
              true))
  | Some "click" ->
      (* Click is not yet implemented in Stateful; just re-render *)
      ignore (get_int "row") ;
      ignore (get_int "col") ;
      emit (current_frame ()) ;
      true
  | Some "tick" -> (
      let n = Option.value ~default:1 (get_int "n") in
      let outcome = HD.Stateful.idle_wait ~iterations:n () in
      (match outcome with
      | `Quit -> emit (nav_response "quit")
      | `Back -> emit (nav_response "back")
      | `SwitchTo name -> emit (nav_response ("switch:" ^ name))
      | `Continue -> emit (current_frame ())) ;
      match outcome with `Quit | `Back -> false | _ -> true)
  | Some "resize" -> (
      match (get_int "rows", get_int "cols") with
      | Some rows, Some cols ->
          HD.set_size rows cols ;
          emit (current_frame ()) ;
          true
      | _ ->
          emit (error_response "Missing 'rows' or 'cols'") ;
          true)
  | Some "quit" ->
      emit (nav_response "quit") ;
      false
  | Some other ->
      emit (error_response (Printf.sprintf "Unknown command: %s" other)) ;
      true

(* ── Main entry point ────────────────────────────────────────────────────── *)

let run ?on_frame (page : (module Tui_page.PAGE_SIG)) =
  on_frame_fn := on_frame ;
  (* Redirect stderr to /dev/null unless verbose mode requested *)
  (match Sys.getenv_opt "MIAOU_HEADLESS_VERBOSE" with
  | None | Some "" -> (
      try
        let null_fd = Unix.openfile "/dev/null" [Unix.O_WRONLY] 0 in
        Unix.dup2 null_fd Unix.stderr ;
        Unix.close null_fd
      with _ -> ())
  | _ -> ()) ;
  (* Unpack and repack to give OCaml a concrete module with an abstract state
     type, avoiding first-class module type-path issues with Stateful.init. *)
  let module P = (val page : Tui_page.PAGE_SIG) in
  HD.Stateful.init (module P) ;
  let rec loop () =
    (* Use run_in_systhread so the blocking input_line runs in a system
       thread.  This keeps the eio event loop active, allowing background
       fibers (e.g. LLM subprocess I/O) to make progress between commands. *)
    match
      Eio_unix.run_in_systhread ~label:"headless-stdin" (fun () ->
          input_line stdin)
    with
    | exception End_of_file -> ()
    | line -> (
        let line = String.trim line in
        if line = "" then loop ()
        else
          match Yojson.Safe.from_string line with
          | exception Yojson.Json_error msg ->
              emit (error_response ("JSON parse error: " ^ msg)) ;
              loop ()
          | `Assoc pairs -> if handle_cmd pairs then loop ()
          | _ ->
              emit (error_response "Expected a JSON object") ;
              loop ())
  in
  (* When a viewer is attached, fork a daemon fiber that periodically
     re-renders the screen and broadcasts changes.  Eio's cooperative
     scheduling ensures the refresh fiber only runs while the main loop
     is blocked on stdin (run_in_systhread), so there is no concurrent
     access to the headless driver state. *)
  (match !on_frame_fn with
  | Some _ ->
      Eio.Switch.run @@ fun sw ->
      Eio.Fiber.fork_daemon ~sw (fun () ->
          let prev = ref "" in
          while true do
            Eio_unix.sleep 0.2 ;
            (* Read the cached screen content written by the command handler.
               Do NOT call idle_wait here — that would race with the command
               handler's own idle_wait (both mutate the shared page-state ref
               and double-tick clocks/timers whenever the command handler
               yields).  The cache is always fresh: every tick/key/render
               command updates it before returning a response. *)
            let size = HD.get_size () in
            let raw = HD.Screen.get () in
            if raw <> !prev then (
              prev := raw ;
              match !on_frame_fn with
              | Some f ->
                  f ~rows:size.LTerm_geom.rows ~cols:size.LTerm_geom.cols raw
              | None -> ())
          done ;
          `Stop_daemon) ;
      loop ()
  | None -> loop ()) ;
  (* Exit cleanly — same pattern as Tui_driver_common which calls exit 0
     after its event loop to avoid blocking on pending Eio fibers. *)
  exit 0