[lsp] Better JSON-RPC serialization

ejgallego · May 16, 2024 · 710097c · 710097c
1 parent 61add17
commit 710097c
Show file tree

Hide file tree

Showing 13 changed files with 231 additions and 171 deletions.
diff --git a/CHANGES.md b/CHANGES.md
@@ -128,6 +128,8 @@
  - Support Coq meta-commands (Reset, Reset Initial, Back) They are
    actually pretty useful to hint the incremental engine to ignore
    changes in some part of the document (@ejgallego, #709)
+ - JSON-RPC library now supports all kind of incoming messages
+   (@ejgallego, #713)
 
 # coq-lsp 0.1.8.1: Spring fix
 -----------------------------

diff --git a/controller/coq_lsp.ml b/controller/coq_lsp.ml
@@ -110,12 +110,19 @@ let lsp_main bt coqcorelib coqlib ocamlpath vo_load_path ml_include_path
   Stdlib.set_binary_mode_out stdout true;
 
   (* We output to stdout *)
-  let ifn () = LIO.read_request stdin in
+  let ifn () = LIO.read_message stdin in
+
   (* Set log channels *)
-  let ofn = LIO.send_json Format.std_formatter in
-  LIO.set_log_fn ofn;
+  let json_fn = LIO.send_json Format.std_formatter in
+
+  let ofn response =
+    let response = Lsp.Base.Response.to_yojson response in
+    LIO.send_json Format.std_formatter response
+  in
+
+  LIO.set_log_fn json_fn;
 
-  let io = lsp_cb ofn in
+  let io = lsp_cb json_fn in
   Fleche.Io.CallBack.set io;
 
   (* IMPORTANT: LSP spec forbids any message from server to client before
@@ -158,7 +165,7 @@ let lsp_main bt coqcorelib coqlib ocamlpath vo_load_path ml_include_path
         Fleche.Config.(
           v := { !v with send_diags = false; send_perf_data = false });
         LIO.set_log_fn (fun _obj -> ());
-        let io = concise_cb ofn in
+        let io = concise_cb json_fn in
         Fleche.Io.CallBack.set io;
         io)
       else io

diff --git a/controller/lsp_core.ml b/controller/lsp_core.ml
@@ -164,12 +164,17 @@ module Rq : sig
   end
 
   val serve :
-    ofn:(J.t -> unit) -> token:Coq.Limits.Token.t -> id:int -> Action.t -> unit
+       ofn:(LSP.Response.t -> unit)
+    -> token:Coq.Limits.Token.t
+    -> id:int
+    -> Action.t
+    -> unit
 
-  val cancel : ofn:(J.t -> unit) -> code:int -> message:string -> int -> unit
+  val cancel :
+    ofn:(LSP.Response.t -> unit) -> code:int -> message:string -> int -> unit
 
   val serve_postponed :
-       ofn:(J.t -> unit)
+       ofn:(LSP.Response.t -> unit)
     -> token:Coq.Limits.Token.t
     -> doc:Fleche.Doc.t
     -> Int.Set.t
@@ -178,8 +183,8 @@ end = struct
   (* Answer a request, private *)
   let answer ~ofn ~id result =
     (match result with
-    | Result.Ok result -> LSP.mk_reply ~id ~result
-    | Error (code, message) -> LSP.mk_request_error ~id ~code ~message)
+    | Result.Ok result -> LSP.Response.mk_ok ~id ~result
+    | Error (code, message) -> LSP.Response.mk_error ~id ~code ~message)
     |> ofn
 
   (* private to the Rq module, just used not to retrigger canceled requests *)
@@ -447,9 +452,8 @@ let lsp_init_process ~ofn ~cmdline ~debug msg : Init_effect.t =
       Format.asprintf "Initializing coq-lsp server %s" (version ())
     in
     LIO.logMessage ~lvl:Info ~message;
-    let result, dirs = Rq_init.do_initialize ~params in
-    (* We don't need to interrupt this *)
     let token = Coq.Limits.Token.create () in
+    let result, dirs = Rq_init.do_initialize ~params in
     Rq.Action.now (Ok result) |> Rq.serve ~ofn ~token ~id;
     LIO.logMessage ~lvl:Info ~message:"Server initialized";
     (* Workspace initialization *)
@@ -463,13 +467,16 @@ let lsp_init_process ~ofn ~cmdline ~debug msg : Init_effect.t =
     Success workspaces
   | LSP.Message.Request { id; _ } ->
     (* per spec *)
-    LSP.mk_request_error ~id ~code:(-32002) ~message:"server not initialized"
+    LSP.Response.mk_error ~id ~code:(-32002) ~message:"server not initialized"
     |> ofn;
     Loop
   | LSP.Message.Notification { method_ = "exit"; params = _ } -> Exit
   | LSP.Message.Notification _ ->
     (* We can't log before getting the initialize message *)
     Loop
+  | LSP.Message.Response _ ->
+    (* O_O *)
+    Loop
 
 (** Dispatching *)
 let dispatch_notification ~io ~ofn ~token ~state ~method_ ~params : unit =
@@ -535,10 +542,16 @@ let dispatch_request ~ofn ~token ~id ~method_ ~params =
 let dispatch_message ~io ~ofn ~token ~state (com : LSP.Message.t) : State.t =
   match com with
   | Notification { method_; params } ->
+    LIO.trace "process_queue" ("Serving notification: " ^ method_);
     dispatch_state_notification ~io ~ofn ~token ~state ~method_ ~params
   | Request { id; method_; params } ->
+    LIO.trace "process_queue" ("Serving Request: " ^ method_);
     dispatch_request ~ofn ~token ~id ~method_ ~params;
     state
+  | Response r ->
+    LIO.trace "process_queue"
+      ("Serving response for: " ^ string_of_int (Lsp.Base.Response.id r));
+    state
 
 (* Queue handling *)
 
@@ -611,7 +624,6 @@ let dispatch_or_resume_check ~io ~ofn ~state =
     let token = token_factory () in
     check_or_yield ~io ~ofn ~token ~state
   | Some com ->
-    LIO.trace "process_queue" ("Serving Request: " ^ LSP.Message.method_ com);
     (* We let Coq work normally now *)
     let token = token_factory () in
     Cont (dispatch_message ~io ~ofn ~token ~state com)

diff --git a/controller/lsp_core.mli b/controller/lsp_core.mli
@@ -41,7 +41,7 @@ module Init_effect : sig
 end
 
 val lsp_init_process :
-     ofn:(Yojson.Safe.t -> unit)
+     ofn:(Lsp.Base.Response.t -> unit)
   -> cmdline:Coq.Workspace.CmdLine.t
   -> debug:bool
   -> Lsp.Base.Message.t
@@ -56,7 +56,7 @@ type 'a cont =
     wake up pending requests *)
 val dispatch_or_resume_check :
      io:Fleche.Io.CallBack.t
-  -> ofn:(Yojson.Safe.t -> unit)
+  -> ofn:(Lsp.Base.Response.t -> unit)
   -> state:State.t
   -> State.t cont option
 

diff --git a/lsp/base.ml b/lsp/base.ml
@@ -25,43 +25,42 @@ module U = Yojson.Safe.Util
 let int_field name dict = U.to_int List.(assoc name dict)
 let string_field name dict = U.to_string List.(assoc name dict)
 
-let odict_field ~default name dict =
-  Option.cata U.to_assoc default (List.assoc_opt name dict)
+module Params = struct
+  type t = (string * Yojson.Safe.t) list
 
-module Message = struct
+  let to_yojson dict = `Assoc dict
+end
+
+module Notification = struct
   type t =
-    | Notification of
-        { method_ : string [@key "method"]
-        ; params : (string * Yojson.Safe.t) list
-        }
-    | Request of
-        { id : int
-        ; method_ : string [@key "method"]
-        ; params : (string * Yojson.Safe.t) list
-        }
+    { method_ : string
+    ; params : Params.t
+    }
 
-  (** Reify an incoming message *)
-  let from_yojson msg =
-    try
-      let dict = U.to_assoc msg in
-      let method_ = string_field "method" dict in
-      let params = odict_field ~default:[] "params" dict in
-      (match List.assoc_opt "id" dict with
-      | None -> Notification { method_; params }
-      | Some id ->
-        let id = U.to_int id in
-        Request { id; method_; params })
-      |> Result.ok
-    with
-    | Not_found -> Error ("missing parameter: " ^ J.to_string msg)
-    | U.Type_error (msg, obj) ->
-      Error (Format.asprintf "msg: %s; obj: %s" msg (J.to_string obj))
+  let make ~method_ ~params () = { method_; params }
+
+  let to_yojson { method_; params } =
+    let params = [ ("params", Params.to_yojson params) ] in
+    `Assoc ([ ("jsonrpc", `String "2.0"); ("method", `String method_) ] @ params)
+end
+
+module Request = struct
+  type t =
+    { id : int
+    ; method_ : string
+    ; params : Params.t
+    }
 
-  let method_ = function
-    | Notification { method_; _ } | Request { method_; _ } -> method_
+  let make ~method_ ~id ~params () = { id; method_; params }
 
-  let params = function
-    | Notification { params; _ } | Request { params; _ } -> params
+  let to_yojson { method_; id; params } =
+    let params = [ ("params", Params.to_yojson params) ] in
+    `Assoc
+      ([ ("jsonrpc", `String "2.0")
+       ; ("id", `Int id)
+       ; ("method", `String method_)
+       ]
+      @ params)
 end
 
 module Response = struct
@@ -77,51 +76,69 @@ module Response = struct
         ; data : Yojson.Safe.t option
         }
 
+  let mk_ok ~id ~result = Ok { id; result }
+  let mk_error ~id ~code ~message = Error { id; code; message; data = None }
+
+  let id = function
+    | Ok { id; _ } | Error { id; _ } -> id
+
+  let to_yojson = function
+    | Ok { id; result } ->
+      `Assoc [ ("jsonrpc", `String "2.0"); ("id", `Int id); ("result", result) ]
+    | Error { id; code; message; data = _ } ->
+      `Assoc
+        [ ("jsonrpc", `String "2.0")
+        ; ("id", `Int id)
+        ; ("error", `Assoc [ ("code", `Int code); ("message", `String message) ])
+        ]
+end
+
+module Message = struct
+  type t =
+    | Notification of Notification.t
+    | Request of Request.t
+    | Response of Response.t
+
+  let response_of_yojson dict =
+    let id = int_field "id" dict in
+    match List.assoc_opt "error" dict with
+    | None ->
+      let result = List.assoc "result" dict in
+      Response.Ok { id; result }
+    | Some error ->
+      let error = U.to_assoc error in
+      let code = int_field "message" error in
+      let message = string_field "message" error in
+      let data = None in
+      Error { id; code; message; data }
+
+  let request_of_yojson method_ dict =
+    let params =
+      List.assoc_opt "params" dict |> Option.map U.to_assoc |> Option.default []
+    in
+    match List.assoc_opt "id" dict with
+    | None -> Notification { Notification.method_; params }
+    | Some id ->
+      let id = U.to_int id in
+      Request { Request.id; method_; params }
+
+  let of_yojson msg =
+    let dict = U.to_assoc msg in
+    match List.assoc_opt "method" dict with
+    | None -> Response (response_of_yojson dict)
+    | Some method_ -> request_of_yojson (U.to_string method_) dict
+
   let of_yojson msg =
-    try
-      let dict = U.to_assoc msg in
-      let id = int_field "id" dict in
-      (match List.assoc_opt "error" dict with
-      | Some error ->
-        let error = U.to_assoc error in
-        let code = int_field "message" error in
-        let message = string_field "message" error in
-        let data = None in
-        Error { id; code; message; data }
-      | None ->
-        let result = List.assoc "result" dict in
-        Ok { id; result })
-      |> Result.ok
-    with
+    try of_yojson msg |> Result.ok with
     | Not_found -> Error ("missing parameter: " ^ J.to_string msg)
     | U.Type_error (msg, obj) ->
       Error (Format.asprintf "msg: %s; obj: %s" msg (J.to_string obj))
-end
 
-let mk_reply ~id ~result =
-  `Assoc [ ("jsonrpc", `String "2.0"); ("id", `Int id); ("result", result) ]
-
-let mk_request_error ~id ~code ~message =
-  `Assoc
-    [ ("jsonrpc", `String "2.0")
-    ; ("id", `Int id)
-    ; ("error", `Assoc [ ("code", `Int code); ("message", `String message) ])
-    ]
-
-let mk_request ~method_ ~id ~params =
-  `Assoc
-    [ ("jsonrpc", `String "2.0")
-    ; ("id", `Int id)
-    ; ("method", `String method_)
-    ; ("params", params)
-    ]
-
-let mk_notification ~method_ ~params =
-  `Assoc
-    [ ("jsonrpc", `String "2.0")
-    ; ("method", `String method_)
-    ; ("params", params)
-    ]
+  let to_yojson = function
+    | Notification n -> Notification.to_yojson n
+    | Request r -> Request.to_yojson r
+    | Response r -> Response.to_yojson r
+end
 
 module ProgressToken : sig
   type t =
@@ -154,7 +171,8 @@ end
 
 let mk_progress ~token ~value f =
   let params = ProgressParams.(to_yojson f { token; value }) in
-  mk_notification ~method_:"$/progress" ~params
+  let params = U.to_assoc params in
+  Notification.(to_yojson { method_ = "$/progress"; params })
 
 module WorkDoneProgressBegin = struct
   type t =