From e99d0d6cc46c11017f1006d7d9b504f8e7132d6c Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Wed, 8 May 2024 11:40:56 -0400
Subject: [PATCH 1/8] add more error handlers and logging around notif/request
 handlers

---
 src/blocking_IO.ml   | 11 ++++-
 src/common_.ml       |  1 +
 src/jsonrpc2.ml      | 64 +++++++++++++++++++----------
 src/lwt/linol_lwt.ml |  8 +++-
 src/server.ml        | 96 +++++++++++++++++++++++++++++++-------------
 src/sigs.ml          |  4 +-
 6 files changed, 130 insertions(+), 54 deletions(-)

diff --git a/src/blocking_IO.ml b/src/blocking_IO.ml
index 60560e44..af17c500 100644
--- a/src/blocking_IO.ml
+++ b/src/blocking_IO.ml
@@ -10,7 +10,9 @@ let ( let* ) x f = f x
 let ( and+ ) a b = a, b
 let return x = x
 let failwith = failwith
-let fail = raise
+
+let fail = Printexc.raise_with_backtrace
+
 let stdin = fun () -> stdin
 let stdout = fun () -> stdout
 
@@ -25,7 +27,12 @@ let default_spawn f =
   in
   ignore (Thread.create run ())
 
-let catch f g = try f () with e -> g e
+let catch f g =
+  try f ()
+  with e ->
+    let bt = Printexc.get_raw_backtrace () in
+    g e bt
+
 let n_bytes_written = Atomic.make 0
 let n_bytes_read = Atomic.make 0
 
diff --git a/src/common_.ml b/src/common_.ml
index c1d0ddc9..78a5810a 100644
--- a/src/common_.ml
+++ b/src/common_.ml
@@ -1,3 +1,4 @@
 module Trace = Trace_core
 
 let ( let@ ) = ( @@ )
+let spf = Printf.sprintf
diff --git a/src/jsonrpc2.ml b/src/jsonrpc2.ml
index 6fa1ca11..75237afd 100644
--- a/src/jsonrpc2.ml
+++ b/src/jsonrpc2.ml
@@ -42,8 +42,6 @@ module Make (IO : IO) : S with module IO = IO = struct
 
   type json = J.t
 
-  let spf = Printf.sprintf
-
   module ErrorCode = Jsonrpc.Response.Error.Code
   (*
   module Err = struct
@@ -138,7 +136,7 @@ module Make (IO : IO) : S with module IO = IO = struct
       (fun () ->
         let+ x = f () in
         Ok x)
-      (fun e -> IO.return (Error e))
+      (fun e bt -> IO.return (Error (e, bt)))
 
   (** Sends a server notification to the LSP client. *)
   let send_server_notification (self : t) (n : Lsp.Server_notification.t) :
@@ -166,10 +164,15 @@ module Make (IO : IO) : S with module IO = IO = struct
   (** Wraps some action and, in case the [IO.t] request has failed, logs the
       failure to the LSP client. *)
   let with_error_handler (self : t) (action : unit -> unit IO.t) : unit IO.t =
-    IO.catch action (fun e ->
+    IO.catch action (fun exn bt ->
+        let message =
+          spf "LSP handler failed with %s\n%s" (Printexc.to_string exn)
+            (Printexc.raw_backtrace_to_string bt)
+        in
+        Log.err (fun k -> k "%s" message);
         let msg =
           Lsp.Types.LogMessageParams.create ~type_:Lsp.Types.MessageType.Error
-            ~message:(Printexc.to_string e)
+            ~message
         in
         let msg =
           Lsp.Server_notification.LogMessage msg
@@ -191,11 +194,12 @@ module Make (IO : IO) : S with module IO = IO = struct
 
   let handle_request (self : t) (r : Jsonrpc.Request.t) : unit IO.t =
     let protect ~id f =
-      IO.catch f (fun e ->
+      IO.catch f (fun e bt ->
           let message =
-            spf "%s\n%s" (Printexc.to_string e) (Printexc.get_backtrace ())
+            spf "%s\n%s" (Printexc.to_string e)
+              (Printexc.raw_backtrace_to_string bt)
           in
-          Log.err (fun k -> k "error: %s" message);
+          Log.err (fun k -> k "error in request handler: %s" message);
           let r =
             Jsonrpc.Response.error id
               (Jsonrpc.Response.Error.make
@@ -215,13 +219,23 @@ module Make (IO : IO) : S with module IO = IO = struct
               ~notify_back:(send_server_notification self)
               ~server_request:(server_request self)
           in
-          let reply_json = Lsp.Client_request.yojson_of_result r reply in
-          let response = Jsonrpc.Response.ok id reply_json in
+          let response =
+            match reply with
+            | Ok reply ->
+              let reply_json = Lsp.Client_request.yojson_of_result r reply in
+              Jsonrpc.Response.ok id reply_json
+            | Error message ->
+              Jsonrpc.Response.error id
+                (Jsonrpc.Response.Error.make
+                   ~code:Jsonrpc.Response.Error.Code.InternalError ~message ())
+          in
+
           send_response self response
         | Error e -> IO.failwith (spf "cannot decode request: %s" e))
-      (fun e ->
+      (fun e bt ->
         let message =
-          spf "%s\n%s" (Printexc.to_string e) (Printexc.get_backtrace ())
+          spf "%s\n%s" (Printexc.to_string e)
+            (Printexc.raw_backtrace_to_string bt)
         in
         Log.err (fun k -> k "error: %s" message);
         let r =
@@ -273,7 +287,8 @@ module Make (IO : IO) : S with module IO = IO = struct
     | _ -> j
 
   (* read a full message *)
-  let read_msg (self : t) : (Jsonrpc.Packet.t, exn) result IO.t =
+  let read_msg (self : t) :
+      (Jsonrpc.Packet.t, exn * Printexc.raw_backtrace) result IO.t =
     let rec read_headers acc =
       let*? line = try_ @@ fun () -> IO.read_line self.ic in
       match String.trim line with
@@ -292,8 +307,9 @@ module Make (IO : IO) : S with module IO = IO = struct
          with
         | pair -> read_headers (pair :: acc)
         | exception _ ->
-          IO.return
-            (Error (E (ErrorCode.ParseError, spf "invalid header: %S" line))))
+          let bt = Printexc.get_raw_backtrace () in
+          let exn = E (ErrorCode.ParseError, spf "invalid header: %S" line) in
+          IO.return (Error (exn, bt)))
     in
     let*? headers = read_headers [] in
     Log.debug (fun k ->
@@ -323,15 +339,21 @@ module Make (IO : IO) : S with module IO = IO = struct
         (match Jsonrpc.Packet.t_of_yojson @@ fix_null_in_params j with
         | m -> IO.return @@ Ok m
         | exception exn ->
+          let bt = Printexc.get_raw_backtrace () in
           Log.err (fun k ->
-              k "cannot decode json message: %s" (Printexc.to_string exn));
-          IO.return (Error (E (ErrorCode.ParseError, "cannot decode json"))))
+              k "cannot decode json message: %s\n%s" (Printexc.to_string exn)
+                (Printexc.raw_backtrace_to_string bt));
+          let exn = E (ErrorCode.ParseError, "cannot decode json") in
+          IO.return (Error (exn, bt)))
       | exception _ ->
+        let bt = Printexc.get_raw_backtrace () in
         IO.return
-        @@ Error (E (ErrorCode.ParseError, "missing content-length' header"))
-    ) else
+        @@ Error (E (ErrorCode.ParseError, "missing content-length' header"), bt)
+    ) else (
+      let bt = Printexc.get_callstack 10 in
       IO.return
-      @@ Error (E (ErrorCode.InvalidRequest, "content-type must be 'utf-8'"))
+      @@ Error (E (ErrorCode.InvalidRequest, "content-type must be 'utf-8'"), bt)
+    )
 
   let send_server_request (self : t) (req : 'from_server Lsp.Server_request.t)
       (cb : ('from_server, Jsonrpc.Response.Error.t) result -> unit IO.t) :
@@ -369,7 +391,7 @@ module Make (IO : IO) : S with module IO = IO = struct
         | Ok r ->
           let* () = process_msg r in
           loop ()
-        | Error e -> IO.fail e
+        | Error (e, bt) -> IO.fail e bt
     in
     loop ()
 end
diff --git a/src/lwt/linol_lwt.ml b/src/lwt/linol_lwt.ml
index dbfa814b..432527bf 100644
--- a/src/lwt/linol_lwt.ml
+++ b/src/lwt/linol_lwt.ml
@@ -29,8 +29,12 @@ module IO_lwt :
   let write = Lwt_io.write_from_exactly
   let read = Lwt_io.read_into_exactly
   let read_line = Lwt_io.read_line
-  let catch = Lwt.catch
-  let fail = Lwt.fail
+
+  let catch f g =
+    let bt = Printexc.get_callstack 10 in
+    Lwt.catch f (fun exn -> g exn bt)
+
+  let fail e _bt = Lwt.fail e
 end
 
 (** Spawn function.
diff --git a/src/server.ml b/src/server.ml
index 221a2d41..9b40160b 100644
--- a/src/server.ml
+++ b/src/server.ml
@@ -60,7 +60,7 @@ module Make (IO : IO) = struct
             server_request:send_request ->
             id:Req_id.t ->
             'a Lsp.Client_request.t ->
-            'a IO.t
+            ('a, string) result IO.t
       (** Method called to handle client requests.
         @param notify_back an object used to reply to the client, send progress
         messages, diagnostics, etc.
@@ -150,6 +150,11 @@ module Make (IO : IO) = struct
   }
   (** Current state of a document. *)
 
+  let[@inline] lift_ok x =
+    let open IO in
+    let+ x = x in
+    Ok x
+
   (** An easily overloadable class. Pick the methods you want to support.
       The user must provide at least the callbacks for document lifecycle:
       open, close, update. The most basic LSP server should check documents
@@ -337,10 +342,24 @@ module Make (IO : IO) = struct
           server_request:_ ->
           id:Req_id.t ->
           r Lsp.Client_request.t ->
-          r IO.t =
+          (r, string) result IO.t =
         fun ~notify_back ~server_request ~id (r : _ Lsp.Client_request.t) ->
           Trace.with_span ~__FILE__ ~__LINE__ "linol.on-request"
-          @@ fun _sp : r IO.t ->
+          @@ fun _sp : (r, string) result IO.t ->
+          (* handler to catch all errors *)
+          let try_catch : (unit -> (r, _) result IO.t) -> (r, _) result IO.t =
+           fun f ->
+            IO.catch f (fun exn bt ->
+                let msg =
+                  spf "LSP request handler failed with %s\n%s"
+                    (Printexc.to_string exn)
+                    (Printexc.raw_backtrace_to_string bt)
+                in
+                Log.err (fun k -> k "%s" msg);
+                IO.return @@ Error msg)
+          in
+
+          try_catch @@ fun () ->
           Log.debug (fun k ->
               k "handle request[id=%s] <opaque>" (Req_id.to_string id));
 
@@ -361,22 +380,23 @@ module Make (IO : IO) = struct
                 ~partialResultToken:None ~workDoneToken:i.workDoneToken
                 ~notify_back ~server_request ()
             in
-            self#on_req_initialize ~notify_back i
+            lift_ok @@ self#on_req_initialize ~notify_back i
           | Lsp.Client_request.TextDocumentHover
               { textDocument; position; workDoneToken } ->
             let uri = textDocument.uri in
             Log.debug (fun k -> k "req: hover '%s'" (DocumentUri.to_path uri));
 
             (match Hashtbl.find_opt docs uri with
-            | None -> IO.return None
+            | None -> IO.return @@ Ok None
             | Some doc_st ->
               let notify_back =
                 new notify_back
                   ~workDoneToken ~partialResultToken:None ~uri ~notify_back
                   ~server_request ()
               in
-              self#on_req_hover ~notify_back ~id ~uri ~pos:position
-                ~workDoneToken doc_st)
+              lift_ok
+              @@ self#on_req_hover ~notify_back ~id ~uri ~pos:position
+                   ~workDoneToken doc_st)
           | Lsp.Client_request.TextDocumentCompletion
               {
                 textDocument;
@@ -389,15 +409,16 @@ module Make (IO : IO) = struct
             Log.debug (fun k ->
                 k "req: complete '%s'" (DocumentUri.to_path uri));
             (match Hashtbl.find_opt docs uri with
-            | None -> IO.return None
+            | None -> IO.return @@ Ok None
             | Some doc_st ->
               let notify_back =
                 new notify_back
                   ~partialResultToken ~workDoneToken ~uri ~notify_back
                   ~server_request ()
               in
-              self#on_req_completion ~notify_back ~id ~uri ~workDoneToken
-                ~partialResultToken ~pos:position ~ctx:context doc_st)
+              lift_ok
+              @@ self#on_req_completion ~notify_back ~id ~uri ~workDoneToken
+                   ~partialResultToken ~pos:position ~ctx:context doc_st)
           | Lsp.Client_request.TextDocumentDefinition
               { textDocument; position; workDoneToken; partialResultToken } ->
             let uri = textDocument.uri in
@@ -410,10 +431,11 @@ module Make (IO : IO) = struct
             in
 
             (match Hashtbl.find_opt docs uri with
-            | None -> IO.return None
+            | None -> IO.return @@ Ok None
             | Some doc_st ->
-              self#on_req_definition ~notify_back ~id ~workDoneToken
-                ~partialResultToken ~uri ~pos:position doc_st)
+              lift_ok
+              @@ self#on_req_definition ~notify_back ~id ~workDoneToken
+                   ~partialResultToken ~uri ~pos:position doc_st)
           | Lsp.Client_request.TextDocumentCodeLens
               { textDocument; workDoneToken; partialResultToken } ->
             let uri = textDocument.uri in
@@ -426,10 +448,11 @@ module Make (IO : IO) = struct
             in
 
             (match Hashtbl.find_opt docs uri with
-            | None -> IO.return []
+            | None -> IO.return @@ Ok []
             | Some doc_st ->
-              self#on_req_code_lens ~notify_back ~id ~uri ~workDoneToken
-                ~partialResultToken doc_st)
+              lift_ok
+              @@ self#on_req_code_lens ~notify_back ~id ~uri ~workDoneToken
+                   ~partialResultToken doc_st)
           | Lsp.Client_request.TextDocumentCodeLensResolve cl ->
             Log.debug (fun k -> k "req: codelens resolve");
             let notify_back =
@@ -437,7 +460,7 @@ module Make (IO : IO) = struct
                 ~workDoneToken:None ~partialResultToken:None ~notify_back
                 ~server_request ()
             in
-            self#on_req_code_lens_resolve ~notify_back ~id cl
+            lift_ok @@ self#on_req_code_lens_resolve ~notify_back ~id cl
           | Lsp.Client_request.ExecuteCommand
               { command; arguments; workDoneToken } ->
             Log.debug (fun k -> k "req: execute command '%s'" command);
@@ -446,8 +469,9 @@ module Make (IO : IO) = struct
                 ~workDoneToken ~partialResultToken:None ~notify_back
                 ~server_request ()
             in
-            self#on_req_execute_command ~notify_back ~id ~workDoneToken command
-              arguments
+            lift_ok
+            @@ self#on_req_execute_command ~notify_back ~id ~workDoneToken
+                 command arguments
           | Lsp.Client_request.DocumentSymbol
               { textDocument = d; workDoneToken; partialResultToken } ->
             let notify_back =
@@ -455,8 +479,9 @@ module Make (IO : IO) = struct
                 ~workDoneToken ~partialResultToken ~notify_back ~server_request
                 ()
             in
-            self#on_req_symbol ~notify_back ~id ~uri:d.uri ~workDoneToken
-              ~partialResultToken ()
+            lift_ok
+            @@ self#on_req_symbol ~notify_back ~id ~uri:d.uri ~workDoneToken
+                 ~partialResultToken ()
           | Lsp.Client_request.CodeAction a ->
             let notify_back =
               new notify_back
@@ -464,15 +489,16 @@ module Make (IO : IO) = struct
                 ~partialResultToken:a.partialResultToken ~notify_back
                 ~server_request ()
             in
-            self#on_req_code_action ~notify_back ~id a
+            lift_ok @@ self#on_req_code_action ~notify_back ~id a
           | Lsp.Client_request.InlayHint p ->
             let notify_back : notify_back =
               new notify_back
                 ~workDoneToken:p.workDoneToken ~partialResultToken:None
                 ~notify_back ~server_request ()
             in
-            self#on_req_inlay_hint ~notify_back ~id ~uri:p.textDocument.uri
-              ~range:p.range ()
+            lift_ok
+            @@ self#on_req_inlay_hint ~notify_back ~id ~uri:p.textDocument.uri
+                 ~range:p.range ()
           | Lsp.Client_request.CodeActionResolve _
           | Lsp.Client_request.LinkedEditingRange _
           | Lsp.Client_request.TextDocumentDeclaration _
@@ -512,15 +538,16 @@ module Make (IO : IO) = struct
                 ~workDoneToken:None ~partialResultToken:None ~notify_back
                 ~server_request ()
             in
-            self#on_request_unhandled ~notify_back ~id r
+            lift_ok @@ self#on_request_unhandled ~notify_back ~id r
           | Lsp.Client_request.UnknownRequest r ->
             let notify_back =
               new notify_back
                 ~workDoneToken:None ~partialResultToken:None ~notify_back
                 ~server_request ()
             in
-            self#on_unknown_request ~notify_back ~server_request ~id r.meth
-              r.params
+            lift_ok
+            @@ self#on_unknown_request ~notify_back ~server_request ~id r.meth
+                 r.params
 
       method virtual on_notif_doc_did_open
           : notify_back:notify_back ->
@@ -559,6 +586,21 @@ module Make (IO : IO) = struct
         let@ _sp =
           Trace.with_span ~__FILE__ ~__LINE__ "linol.on-notification"
         in
+
+        (* handler to catch all errors *)
+        let try_catch : (unit -> unit IO.t) -> unit IO.t =
+         fun f ->
+          IO.catch f (fun exn bt ->
+              let msg =
+                spf "LSP notification handler failed with %s\n%s"
+                  (Printexc.to_string exn)
+                  (Printexc.raw_backtrace_to_string bt)
+              in
+              Log.err (fun k -> k "%s" msg);
+              IO.return ())
+        in
+
+        try_catch @@ fun () ->
         let open Lsp.Types in
         match n with
         | Lsp.Client_notification.TextDocumentDidOpen
diff --git a/src/sigs.ml b/src/sigs.ml
index 125be840..784f27c6 100644
--- a/src/sigs.ml
+++ b/src/sigs.ml
@@ -18,6 +18,6 @@ module type IO = sig
   val read_line : in_channel -> string t
   val write : out_channel -> bytes -> int -> int -> unit t
   val write_string : out_channel -> string -> unit t
-  val fail : exn -> unit t
-  val catch : (unit -> 'a t) -> (exn -> 'a t) -> 'a t
+  val fail : exn -> Printexc.raw_backtrace -> unit t
+  val catch : (unit -> 'a t) -> (exn -> Printexc.raw_backtrace -> 'a t) -> 'a t
 end

From f2e3c4b369e97bd6dce24bc2d0966abad0c45f93 Mon Sep 17 00:00:00 2001
From: "Christoph M. Wintersteiger" <christoph@imandra.ai>
Date: Tue, 30 Jul 2024 16:42:39 +0100
Subject: [PATCH 2/8] Update return from on_req_shutdown

---
 src/server.ml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/server.ml b/src/server.ml
index 9b40160b..97700568 100644
--- a/src/server.ml
+++ b/src/server.ml
@@ -372,7 +372,7 @@ module Make (IO : IO) = struct
                 ~server_request ()
             in
             status <- `ReceivedShutdown;
-            self#on_req_shutdown ~notify_back ~id
+            lift_ok @@ self#on_req_shutdown ~notify_back ~id
           | Lsp.Client_request.Initialize i ->
             Log.debug (fun k -> k "req: initialize");
             let notify_back =

From 845371e3a0b6a67f5c5fe8f56f3c2a7e0ffea69c Mon Sep 17 00:00:00 2001
From: "Christoph M. Wintersteiger" <christoph@imandra.ai>
Date: Tue, 30 Jul 2024 16:50:35 +0100
Subject: [PATCH 3/8] Rebase-fix

---
 src/eio/linol_eio.ml | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/eio/linol_eio.ml b/src/eio/linol_eio.ml
index 8ecc2b1d..a3cb606c 100644
--- a/src/eio/linol_eio.ml
+++ b/src/eio/linol_eio.ml
@@ -14,7 +14,10 @@ module IO_eio :
   let return x = x
   let failwith = failwith
   let fail = raise
-  let catch f handler = try f () with exn -> handler exn
+
+  let catch f handler = try f () with exn ->
+    let bt = Printexc.get_raw_backtrace () in
+    handler exn bt
 
   let stdin env = Eio.Buf_read.of_flow ~max_size:1_000_000 (Eio.Stdenv.stdin env)
   let stdout = Eio.Stdenv.stdout

From 7ccdfff203e5336a175dbb12f4a6587c41455771 Mon Sep 17 00:00:00 2001
From: "Christoph M. Wintersteiger" <christoph@imandra.ai>
Date: Tue, 30 Jul 2024 21:01:05 +0100
Subject: [PATCH 4/8] I've gone too far, but it's too late to go back.

---
 example/template-eio/main.ml |  8 ++--
 example/template-lwt/main.ml |  6 +--
 src/eio/linol_eio.ml         |  9 ++--
 src/jsonrpc2.ml              | 41 ++++++++++--------
 src/jsonrpc2.mli             |  2 +-
 src/lwt/linol_lwt.ml         | 16 +++++---
 src/server.ml                | 80 ++++++++++++++++++------------------
 7 files changed, 86 insertions(+), 76 deletions(-)

diff --git a/example/template-eio/main.ml b/example/template-eio/main.ml
index 932e0cfa..c89697c8 100644
--- a/example/template-eio/main.ml
+++ b/example/template-eio/main.ml
@@ -51,7 +51,7 @@ class lsp_server =
        - return the diagnostics from the new state
     *)
     method private _on_doc ~(notify_back : Linol_eio.Jsonrpc2.notify_back)
-        (uri : Lsp.Types.DocumentUri.t) (contents : string) =
+        (uri : Lsp.Types.DocumentUri.t) (contents : string) : (unit, string) result =
       let new_state = process_some_input_file contents in
       Hashtbl.replace buffers uri new_state;
       let diags = diagnostics new_state in
@@ -59,7 +59,7 @@ class lsp_server =
 
     (* We now override the [on_notify_doc_did_open] method that will be called
        by the server each time a new document is opened. *)
-    method on_notif_doc_did_open ~notify_back d ~content : unit Linol_eio.t =
+    method on_notif_doc_did_open ~notify_back d ~content : (unit, string) result Linol_eio.t =
       self#_on_doc ~notify_back d.uri content
 
     (* Similarly, we also override the [on_notify_doc_did_change] method that will be called
@@ -70,9 +70,9 @@ class lsp_server =
 
     (* On document closes, we remove the state associated to the file from the global
        hashtable state, to avoid leaking memory. *)
-    method on_notif_doc_did_close ~notify_back:_ d : unit Linol_eio.t =
+    method on_notif_doc_did_close ~notify_back:_ d : (unit, string) result Linol_eio.t =
       Hashtbl.remove buffers d.uri;
-      ()
+      Ok ()
   end
 
 (* Main code
diff --git a/example/template-lwt/main.ml b/example/template-lwt/main.ml
index 42e0ea26..78083a1c 100644
--- a/example/template-lwt/main.ml
+++ b/example/template-lwt/main.ml
@@ -59,7 +59,7 @@ class lsp_server =
 
     (* We now override the [on_notify_doc_did_open] method that will be called
        by the server each time a new document is opened. *)
-    method on_notif_doc_did_open ~notify_back d ~content : unit Linol_lwt.t =
+    method on_notif_doc_did_open ~notify_back d ~content : (unit, string) result Linol_lwt.t =
       self#_on_doc ~notify_back d.uri content
 
     (* Similarly, we also override the [on_notify_doc_did_change] method that will be called
@@ -70,9 +70,9 @@ class lsp_server =
 
     (* On document closes, we remove the state associated to the file from the global
        hashtable state, to avoid leaking memory. *)
-    method on_notif_doc_did_close ~notify_back:_ d : unit Linol_lwt.t =
+    method on_notif_doc_did_close ~notify_back:_ d : (unit, string) result Linol_lwt.t =
       Hashtbl.remove buffers d.uri;
-      Linol_lwt.return ()
+      Linol_lwt.return (Ok ())
   end
 
 (* Main code
diff --git a/src/eio/linol_eio.ml b/src/eio/linol_eio.ml
index a3cb606c..1aa8fd01 100644
--- a/src/eio/linol_eio.ml
+++ b/src/eio/linol_eio.ml
@@ -38,12 +38,13 @@ module IO_eio :
 end
 
 (** Spawn function. *)
-let spawn f =
+let spawn (f:unit -> (unit, string) result) : unit  =
   let promise, resolver = Eio.Promise.create () in
   begin
     try
-      f ();
-      Eio.Promise.resolve_ok resolver ()
+      match f () with
+      | Ok _ -> Eio.Promise.resolve_ok resolver ()
+      | Error _ -> ()
     with
       exn ->
       (Printf.eprintf "uncaught exception in `spawn`:\n%s\n%!"
@@ -51,7 +52,7 @@ let spawn f =
       Eio.Promise.resolve_error resolver exn
     end;
 
-  Eio.Promise.await_exn promise
+  (Eio.Promise.await_exn promise)
 
 include Lsp.Types
 include IO_eio
diff --git a/src/jsonrpc2.ml b/src/jsonrpc2.ml
index 75237afd..33261c5b 100644
--- a/src/jsonrpc2.ml
+++ b/src/jsonrpc2.ml
@@ -24,7 +24,7 @@ module type S = sig
   val create_stdio :
     ?on_received:(json -> unit) -> ?on_sent:(json -> unit) -> env:IO.env -> server -> t
 
-  val send_server_notification : t -> Lsp.Server_notification.t -> unit IO.t
+  val send_server_notification : t -> Lsp.Server_notification.t -> (unit, string) result IO.t
 
   val send_server_request :
     t ->
@@ -56,6 +56,11 @@ module Make (IO : IO) : S with module IO = IO = struct
 
   exception E of ErrorCode.t * string
 
+  let[@inline] lift_ok x =
+    let open IO in
+    let+ x = x in
+    Ok x
+
   (* bind on IO+result *)
   let ( let*? ) x f =
     let* x = x in
@@ -110,9 +115,9 @@ module Make (IO : IO) : S with module IO = IO = struct
       [register_server_request_response_handler] before calling this method to
       ensure that [handle_response] will have a registered handler for this
       response. *)
-  let send_server_req (self : t) (m : Jsonrpc.Request.t) : unit IO.t =
+  let send_server_req (self : t) (m : Jsonrpc.Request.t) : (unit, string) result IO.t =
     let json = Jsonrpc.Request.yojson_of_t m in
-    send_json_ self json
+    lift_ok @@ send_json_ self json
 
   (** Returns a new, unused [Req_id.t] to send a server request. *)
   let fresh_lsp_id (self : t) : Req_id.t =
@@ -140,9 +145,9 @@ module Make (IO : IO) : S with module IO = IO = struct
 
   (** Sends a server notification to the LSP client. *)
   let send_server_notification (self : t) (n : Lsp.Server_notification.t) :
-      unit IO.t =
+      (unit, string) result IO.t =
     let msg = Lsp.Server_notification.to_jsonrpc n in
-    send_server_notif self msg
+    lift_ok @@ (send_server_notif self msg)
 
   (** Given a [server_request_handler_pair] consisting of some server request
       and its handler, sends this request to the LSP client and adds the handler
@@ -156,7 +161,7 @@ module Make (IO : IO) : S with module IO = IO = struct
     let msg = Lsp.Server_request.to_jsonrpc_request r ~id in
     let has_inserted = register_server_request_response_handler self id req in
     if has_inserted then
-      let* () = send_server_req self msg in
+      let* _res = send_server_req self msg in
       return id
     else
       IO.failwith "failed to register server request: id was already used"
@@ -180,19 +185,19 @@ module Make (IO : IO) : S with module IO = IO = struct
         in
         send_server_notif self msg)
 
-  let handle_notification (self : t) (n : Jsonrpc.Notification.t) : unit IO.t =
+  let handle_notification (self : t) (n : Jsonrpc.Notification.t) : (unit, string) result IO.t =
     let@ _sp =
       Trace.with_span ~__FILE__ ~__LINE__ "linol.handle-notification"
     in
     match Lsp.Client_notification.of_jsonrpc n with
     | Ok n ->
-      let@ () = with_error_handler self in
+      (* let@ () = with_error_handler self in *)
       self.s#on_notification n
         ~notify_back:(send_server_notification self)
         ~server_request:(server_request self)
     | Error e -> IO.failwith (spf "cannot decode notification: %s" e)
 
-  let handle_request (self : t) (r : Jsonrpc.Request.t) : unit IO.t =
+  let handle_request (self : t) (r : Jsonrpc.Request.t) : (unit, string) result IO.t =
     let protect ~id f =
       IO.catch f (fun e bt ->
           let message =
@@ -209,7 +214,7 @@ module Make (IO : IO) : S with module IO = IO = struct
     in
     (* request, so we need to reply *)
     let id = r.id in
-    IO.catch
+     lift_ok @@ IO.catch
       (fun () ->
         match Lsp.Client_request.of_jsonrpc r with
         | Ok (Lsp.Client_request.E r) ->
@@ -245,7 +250,7 @@ module Make (IO : IO) : S with module IO = IO = struct
         in
         send_response self r)
 
-  let handle_response (self : t) (r : Jsonrpc.Response.t) : unit IO.t =
+  let handle_response (self : t) (r : Jsonrpc.Response.t) : (unit, string) result IO.t =
     match Hashtbl.find_opt self.pending_responses r.id with
     | None ->
       IO.failwith
@@ -253,23 +258,23 @@ module Make (IO : IO) : S with module IO = IO = struct
       @@ Req_id.to_string r.id
     | Some (Request_and_handler (req, handler)) ->
       let () = Hashtbl.remove self.pending_responses r.id in
-      (match r.result with
+      lift_ok @@ (match r.result with
       | Error err -> with_error_handler self (fun () -> handler @@ Error err)
       | Ok json ->
         let r = Lsp.Server_request.response_of_json req json in
         with_error_handler self (fun () -> handler @@ Ok r))
 
   let handle_batch_response (_self : t) (_rs : Jsonrpc.Response.t list) :
-      unit IO.t =
-    IO.failwith "Unhandled: jsonrpc batch response"
+      (unit, string) result IO.t =
+      lift_ok @@ IO.failwith "Unhandled: jsonrpc batch response"
 
   let handle_batch_call (_self : t)
       (_cs :
         [ `Notification of Jsonrpc.Notification.t
         | `Request of Jsonrpc.Request.t
         ]
-        list) : unit IO.t =
-    IO.failwith "Unhandled: jsonrpc batch call"
+        list) : (unit, string) result IO.t =
+    lift_ok @@  IO.failwith "Unhandled: jsonrpc batch call"
 
   (* As in [https://github.com/c-cube/linol/issues/20],
      Jsonrpc expect "params" to be object or array,
@@ -367,7 +372,7 @@ module Make (IO : IO) : S with module IO = IO = struct
   let run ?(shutdown = fun _ -> false) (self : t) : unit IO.t =
     let async f =
       self.s#spawn_query_handler f;
-      IO.return ()
+      IO.return (Ok ())
     in
 
     let process_msg r =
@@ -389,7 +394,7 @@ module Make (IO : IO) : S with module IO = IO = struct
         let* r = read_msg self in
         match r with
         | Ok r ->
-          let* () = process_msg r in
+          let* _res = process_msg r in
           loop ()
         | Error (e, bt) -> IO.fail e bt
     in
diff --git a/src/jsonrpc2.mli b/src/jsonrpc2.mli
index 2ed46b04..68efa11c 100644
--- a/src/jsonrpc2.mli
+++ b/src/jsonrpc2.mli
@@ -27,7 +27,7 @@ module type S = sig
     ?on_received:(json -> unit) -> ?on_sent:(json -> unit) -> env:IO.env -> server -> t
   (** Create a connection using stdin/stdout *)
 
-  val send_server_notification : t -> Lsp.Server_notification.t -> unit IO.t
+  val send_server_notification : t -> Lsp.Server_notification.t -> (unit, string) result IO.t
   (** Send a notification from the server.
       @since 0.5 *)
 
diff --git a/src/lwt/linol_lwt.ml b/src/lwt/linol_lwt.ml
index 432527bf..60963389 100644
--- a/src/lwt/linol_lwt.ml
+++ b/src/lwt/linol_lwt.ml
@@ -39,12 +39,16 @@ end
 
 (** Spawn function.
     @since 0.5 *)
-let spawn f =
-  Lwt.async (fun () ->
-      Lwt.catch f (fun exn ->
-          Printf.eprintf "uncaught exception in `spawn`:\n%s\n%!"
-            (Printexc.to_string exn);
-          Lwt.return ()))
+let spawn (f: unit -> (unit, string) result Lwt.t) : unit =
+  let g = (fun () ->
+    let _ = Lwt.catch f (fun exn ->
+      Printf.eprintf "uncaught exception in `spawn`:\n%s\n%!"
+        (Printexc.to_string exn);
+      Lwt.return (Error (Printexc.to_string exn)))
+    in
+    Lwt.return ())
+  in
+  Lwt.async g
 
 include Lsp.Types
 include IO_lwt
diff --git a/src/server.ml b/src/server.ml
index 97700568..f581ba32 100644
--- a/src/server.ml
+++ b/src/server.ml
@@ -32,6 +32,11 @@ module Make (IO : IO) = struct
   module DiagnosticSeverity = DiagnosticSeverity
   module Req_id = Req_id
 
+  let[@inline] lift_ok x =
+    let open IO in
+    let+ x = x in
+    Ok x
+
   (** A variant carrying a [Lsp.Server_request.t] and a handler for its return
       value. The request is stored in order to allow us to discriminate its
       existential variable. *)
@@ -49,14 +54,14 @@ module Make (IO : IO) = struct
   class virtual base_server =
     object
       method virtual on_notification
-          : notify_back:(Lsp.Server_notification.t -> unit IO.t) ->
+          : notify_back:(Lsp.Server_notification.t -> (unit, string) result IO.t) ->
             server_request:send_request ->
             Lsp.Client_notification.t ->
-            unit IO.t
+            (unit, string) result IO.t
 
       method virtual on_request
           : 'a.
-            notify_back:(Lsp.Server_notification.t -> unit IO.t) ->
+            notify_back:(Lsp.Server_notification.t -> (unit, string) result IO.t) ->
             server_request:send_request ->
             id:Req_id.t ->
             'a Lsp.Client_request.t ->
@@ -69,30 +74,30 @@ module Make (IO : IO) = struct
       method must_quit = false
       (** Set to true if the client requested to exit *)
 
-      method virtual spawn_query_handler : (unit -> unit IO.t) -> unit
+      method virtual spawn_query_handler : (unit -> (unit, string) result IO.t) -> unit
       (** How to start a new future/task/thread concurrently. This is used
           to process incoming user queries.
           @since 0.5 *)
     end
 
-  let async (self : #base_server) f : unit IO.t =
+  let async (self : #base_server) f : (unit, string) result IO.t =
     self#spawn_query_handler f;
-    IO.return ()
+    IO.return (Ok ())
 
   (** A wrapper to more easily reply to notifications *)
-  class notify_back ~notify_back ~server_request ~workDoneToken
+  class notify_back ~(notify_back: Lsp.Server_notification.t -> (unit, string) result IO.t) ~server_request ~workDoneToken
     ~partialResultToken:_ ?version ?(uri : DocumentUri.t option) () =
     object
       val mutable uri = uri
       method set_uri u = uri <- Some u
       method get_uri = uri
 
-      method send_log_msg ~type_ msg : unit IO.t =
+      method send_log_msg ~type_ msg : (unit, string) result IO.t =
         let params = LogMessageParams.create ~type_ ~message:msg in
         notify_back (Lsp.Server_notification.LogMessage params)
       (** Send a log message to the editor *)
 
-      method send_diagnostic (l : Diagnostic.t list) : unit IO.t =
+      method send_diagnostic (l : Diagnostic.t list) : (unit, string) result IO.t =
         match uri with
         | None ->
           IO.failwith "notify_back: cannot publish diagnostics, no URI given"
@@ -103,33 +108,33 @@ module Make (IO : IO) = struct
           notify_back (Lsp.Server_notification.PublishDiagnostics params)
       (** Send diagnostics for the current document *)
 
-      method telemetry json : unit IO.t =
+      method telemetry json : (unit, string) result IO.t =
         notify_back @@ Lsp.Server_notification.TelemetryNotification json
 
-      method cancel_request (id : Jsonrpc.Id.t) : unit IO.t =
+      method cancel_request (id : Jsonrpc.Id.t) : (unit, string) result IO.t =
         notify_back @@ CancelRequest id
 
       method work_done_progress_begin (p : Lsp.Types.WorkDoneProgressBegin.t)
-          : unit IO.t =
+          : (unit, string) result IO.t =
         match workDoneToken with
         | Some token ->
           notify_back @@ WorkDoneProgress { token; value = Begin p }
-        | None -> IO.return ()
+        | None -> lift_ok @@ IO.return ()
 
       method work_done_progress_report (p : Lsp.Types.WorkDoneProgressReport.t)
-          : unit IO.t =
+          : (unit, string) result IO.t =
         match workDoneToken with
         | Some token ->
           notify_back @@ WorkDoneProgress { value = Report p; token }
-        | None -> IO.return ()
+        | None -> lift_ok @@ IO.return ()
 
       method work_done_progress_end (p : Lsp.Types.WorkDoneProgressEnd.t)
-          : unit IO.t =
+          : (unit, string) result IO.t =
         match workDoneToken with
         | Some token -> notify_back @@ WorkDoneProgress { value = End p; token }
-        | None -> IO.return ()
+        | None -> lift_ok @@ IO.return ()
 
-      method send_notification (n : Lsp.Server_notification.t) : unit IO.t =
+      method send_notification (n : Lsp.Server_notification.t) : (unit, string) result IO.t =
         notify_back n
       (** Send a notification from the server to the client (general purpose method) *)
 
@@ -150,11 +155,6 @@ module Make (IO : IO) = struct
   }
   (** Current state of a document. *)
 
-  let[@inline] lift_ok x =
-    let open IO in
-    let+ x = x in
-    Ok x
-
   (** An easily overloadable class. Pick the methods you want to support.
       The user must provide at least the callbacks for document lifecycle:
       open, close, update. The most basic LSP server should check documents
@@ -338,7 +338,7 @@ module Make (IO : IO) = struct
         @since 0.7 *)
 
       method on_request : type r.
-          notify_back:_ ->
+          notify_back:(Lsp.Server_notification.t -> (unit, string) result IO.t) ->
           server_request:_ ->
           id:Req_id.t ->
           r Lsp.Client_request.t ->
@@ -553,11 +553,11 @@ module Make (IO : IO) = struct
           : notify_back:notify_back ->
             TextDocumentItem.t ->
             content:string ->
-            unit IO.t
+            (unit, string) result IO.t
       (** Called when a document is opened *)
 
       method virtual on_notif_doc_did_close
-          : notify_back:notify_back -> TextDocumentIdentifier.t -> unit IO.t
+          : notify_back:notify_back -> TextDocumentIdentifier.t -> (unit, string) result IO.t
 
       method virtual on_notif_doc_did_change
           : notify_back:notify_back ->
@@ -565,30 +565,30 @@ module Make (IO : IO) = struct
             TextDocumentContentChangeEvent.t list ->
             old_content:string ->
             new_content:string ->
-            unit IO.t
+            (unit, string) result IO.t
       (** Called when the document changes. *)
 
       method on_notif_doc_did_save ~notify_back:(_ : notify_back)
-          (_params : DidSaveTextDocumentParams.t ) : unit IO.t =
-        IO.return ()
+          (_params : DidSaveTextDocumentParams.t ) : (unit, string) result IO.t =
+        lift_ok @@ IO.return ()
 
       method on_unknown_notification ~notify_back:(_ : notify_back)
-          (_n : Jsonrpc.Notification.t) : unit IO.t =
-        IO.return ()
+          (_n : Jsonrpc.Notification.t) : (unit, string) result IO.t =
+        lift_ok @@ IO.return ()
 
       method on_notification_unhandled ~notify_back:(_ : notify_back)
-          (_n : Lsp.Client_notification.t) : unit IO.t =
-        IO.return ()
+          (_n : Lsp.Client_notification.t) : (unit, string) result IO.t =
+        lift_ok @@ IO.return ()
       (** Override to handle unprocessed notifications *)
 
-      method on_notification ~notify_back ~server_request
-          (n : Lsp.Client_notification.t) : unit IO.t =
+      method on_notification ~(notify_back:Lsp.Server_notification.t -> (unit, string) result IO.t) ~server_request
+          (n : Lsp.Client_notification.t) : (unit, string) result IO.t =
         let@ _sp =
           Trace.with_span ~__FILE__ ~__LINE__ "linol.on-notification"
         in
 
         (* handler to catch all errors *)
-        let try_catch : (unit -> unit IO.t) -> unit IO.t =
+        let try_catch : (unit -> (unit, string) result IO.t) -> (unit, string) result IO.t =
          fun f ->
           IO.catch f (fun exn bt ->
               let msg =
@@ -597,7 +597,7 @@ module Make (IO : IO) = struct
                   (Printexc.raw_backtrace_to_string bt)
               in
               Log.err (fun k -> k "%s" msg);
-              IO.return ())
+              lift_ok @@ IO.return ())
         in
 
         try_catch @@ fun () ->
@@ -688,11 +688,11 @@ module Make (IO : IO) = struct
           Hashtbl.replace docs doc.uri new_st;
 
           async self (fun () ->
-              self#on_notif_doc_did_change
+              (self#on_notif_doc_did_change
                 ~notify_back:(notify_back : notify_back)
                 doc c
                 ~old_content:(Lsp.Text_document.text old_doc)
-                ~new_content:new_st.content)
+                ~new_content:new_st.content))
         | Lsp.Client_notification.DidSaveTextDocument params ->
           let notify_back =
             new notify_back
@@ -706,7 +706,7 @@ module Make (IO : IO) = struct
                 params)
         | Lsp.Client_notification.Exit ->
           status <- `ReceivedExit;
-          IO.return ()
+          lift_ok @@ IO.return ()
         | Lsp.Client_notification.WillSaveTextDocument _
         | Lsp.Client_notification.ChangeWorkspaceFolders _
         | Lsp.Client_notification.ChangeConfiguration _

From 9a7e1c221ca0f54507939d929dc3a850d5eaedc0 Mon Sep 17 00:00:00 2001
From: "Christoph M. Wintersteiger" <christoph@imandra.ai>
Date: Wed, 31 Jul 2024 16:20:53 +0100
Subject: [PATCH 5/8] Revert "I've gone too far, but it's too late to go back."

This reverts commit 7ccdfff203e5336a175dbb12f4a6587c41455771.
---
 example/template-eio/main.ml |  8 ++--
 example/template-lwt/main.ml |  6 +--
 src/eio/linol_eio.ml         |  9 ++--
 src/jsonrpc2.ml              | 41 ++++++++----------
 src/jsonrpc2.mli             |  2 +-
 src/lwt/linol_lwt.ml         | 16 +++-----
 src/server.ml                | 80 ++++++++++++++++++------------------
 7 files changed, 76 insertions(+), 86 deletions(-)

diff --git a/example/template-eio/main.ml b/example/template-eio/main.ml
index c89697c8..932e0cfa 100644
--- a/example/template-eio/main.ml
+++ b/example/template-eio/main.ml
@@ -51,7 +51,7 @@ class lsp_server =
        - return the diagnostics from the new state
     *)
     method private _on_doc ~(notify_back : Linol_eio.Jsonrpc2.notify_back)
-        (uri : Lsp.Types.DocumentUri.t) (contents : string) : (unit, string) result =
+        (uri : Lsp.Types.DocumentUri.t) (contents : string) =
       let new_state = process_some_input_file contents in
       Hashtbl.replace buffers uri new_state;
       let diags = diagnostics new_state in
@@ -59,7 +59,7 @@ class lsp_server =
 
     (* We now override the [on_notify_doc_did_open] method that will be called
        by the server each time a new document is opened. *)
-    method on_notif_doc_did_open ~notify_back d ~content : (unit, string) result Linol_eio.t =
+    method on_notif_doc_did_open ~notify_back d ~content : unit Linol_eio.t =
       self#_on_doc ~notify_back d.uri content
 
     (* Similarly, we also override the [on_notify_doc_did_change] method that will be called
@@ -70,9 +70,9 @@ class lsp_server =
 
     (* On document closes, we remove the state associated to the file from the global
        hashtable state, to avoid leaking memory. *)
-    method on_notif_doc_did_close ~notify_back:_ d : (unit, string) result Linol_eio.t =
+    method on_notif_doc_did_close ~notify_back:_ d : unit Linol_eio.t =
       Hashtbl.remove buffers d.uri;
-      Ok ()
+      ()
   end
 
 (* Main code
diff --git a/example/template-lwt/main.ml b/example/template-lwt/main.ml
index 78083a1c..42e0ea26 100644
--- a/example/template-lwt/main.ml
+++ b/example/template-lwt/main.ml
@@ -59,7 +59,7 @@ class lsp_server =
 
     (* We now override the [on_notify_doc_did_open] method that will be called
        by the server each time a new document is opened. *)
-    method on_notif_doc_did_open ~notify_back d ~content : (unit, string) result Linol_lwt.t =
+    method on_notif_doc_did_open ~notify_back d ~content : unit Linol_lwt.t =
       self#_on_doc ~notify_back d.uri content
 
     (* Similarly, we also override the [on_notify_doc_did_change] method that will be called
@@ -70,9 +70,9 @@ class lsp_server =
 
     (* On document closes, we remove the state associated to the file from the global
        hashtable state, to avoid leaking memory. *)
-    method on_notif_doc_did_close ~notify_back:_ d : (unit, string) result Linol_lwt.t =
+    method on_notif_doc_did_close ~notify_back:_ d : unit Linol_lwt.t =
       Hashtbl.remove buffers d.uri;
-      Linol_lwt.return (Ok ())
+      Linol_lwt.return ()
   end
 
 (* Main code
diff --git a/src/eio/linol_eio.ml b/src/eio/linol_eio.ml
index 1aa8fd01..a3cb606c 100644
--- a/src/eio/linol_eio.ml
+++ b/src/eio/linol_eio.ml
@@ -38,13 +38,12 @@ module IO_eio :
 end
 
 (** Spawn function. *)
-let spawn (f:unit -> (unit, string) result) : unit  =
+let spawn f =
   let promise, resolver = Eio.Promise.create () in
   begin
     try
-      match f () with
-      | Ok _ -> Eio.Promise.resolve_ok resolver ()
-      | Error _ -> ()
+      f ();
+      Eio.Promise.resolve_ok resolver ()
     with
       exn ->
       (Printf.eprintf "uncaught exception in `spawn`:\n%s\n%!"
@@ -52,7 +51,7 @@ let spawn (f:unit -> (unit, string) result) : unit  =
       Eio.Promise.resolve_error resolver exn
     end;
 
-  (Eio.Promise.await_exn promise)
+  Eio.Promise.await_exn promise
 
 include Lsp.Types
 include IO_eio
diff --git a/src/jsonrpc2.ml b/src/jsonrpc2.ml
index 33261c5b..75237afd 100644
--- a/src/jsonrpc2.ml
+++ b/src/jsonrpc2.ml
@@ -24,7 +24,7 @@ module type S = sig
   val create_stdio :
     ?on_received:(json -> unit) -> ?on_sent:(json -> unit) -> env:IO.env -> server -> t
 
-  val send_server_notification : t -> Lsp.Server_notification.t -> (unit, string) result IO.t
+  val send_server_notification : t -> Lsp.Server_notification.t -> unit IO.t
 
   val send_server_request :
     t ->
@@ -56,11 +56,6 @@ module Make (IO : IO) : S with module IO = IO = struct
 
   exception E of ErrorCode.t * string
 
-  let[@inline] lift_ok x =
-    let open IO in
-    let+ x = x in
-    Ok x
-
   (* bind on IO+result *)
   let ( let*? ) x f =
     let* x = x in
@@ -115,9 +110,9 @@ module Make (IO : IO) : S with module IO = IO = struct
       [register_server_request_response_handler] before calling this method to
       ensure that [handle_response] will have a registered handler for this
       response. *)
-  let send_server_req (self : t) (m : Jsonrpc.Request.t) : (unit, string) result IO.t =
+  let send_server_req (self : t) (m : Jsonrpc.Request.t) : unit IO.t =
     let json = Jsonrpc.Request.yojson_of_t m in
-    lift_ok @@ send_json_ self json
+    send_json_ self json
 
   (** Returns a new, unused [Req_id.t] to send a server request. *)
   let fresh_lsp_id (self : t) : Req_id.t =
@@ -145,9 +140,9 @@ module Make (IO : IO) : S with module IO = IO = struct
 
   (** Sends a server notification to the LSP client. *)
   let send_server_notification (self : t) (n : Lsp.Server_notification.t) :
-      (unit, string) result IO.t =
+      unit IO.t =
     let msg = Lsp.Server_notification.to_jsonrpc n in
-    lift_ok @@ (send_server_notif self msg)
+    send_server_notif self msg
 
   (** Given a [server_request_handler_pair] consisting of some server request
       and its handler, sends this request to the LSP client and adds the handler
@@ -161,7 +156,7 @@ module Make (IO : IO) : S with module IO = IO = struct
     let msg = Lsp.Server_request.to_jsonrpc_request r ~id in
     let has_inserted = register_server_request_response_handler self id req in
     if has_inserted then
-      let* _res = send_server_req self msg in
+      let* () = send_server_req self msg in
       return id
     else
       IO.failwith "failed to register server request: id was already used"
@@ -185,19 +180,19 @@ module Make (IO : IO) : S with module IO = IO = struct
         in
         send_server_notif self msg)
 
-  let handle_notification (self : t) (n : Jsonrpc.Notification.t) : (unit, string) result IO.t =
+  let handle_notification (self : t) (n : Jsonrpc.Notification.t) : unit IO.t =
     let@ _sp =
       Trace.with_span ~__FILE__ ~__LINE__ "linol.handle-notification"
     in
     match Lsp.Client_notification.of_jsonrpc n with
     | Ok n ->
-      (* let@ () = with_error_handler self in *)
+      let@ () = with_error_handler self in
       self.s#on_notification n
         ~notify_back:(send_server_notification self)
         ~server_request:(server_request self)
     | Error e -> IO.failwith (spf "cannot decode notification: %s" e)
 
-  let handle_request (self : t) (r : Jsonrpc.Request.t) : (unit, string) result IO.t =
+  let handle_request (self : t) (r : Jsonrpc.Request.t) : unit IO.t =
     let protect ~id f =
       IO.catch f (fun e bt ->
           let message =
@@ -214,7 +209,7 @@ module Make (IO : IO) : S with module IO = IO = struct
     in
     (* request, so we need to reply *)
     let id = r.id in
-     lift_ok @@ IO.catch
+    IO.catch
       (fun () ->
         match Lsp.Client_request.of_jsonrpc r with
         | Ok (Lsp.Client_request.E r) ->
@@ -250,7 +245,7 @@ module Make (IO : IO) : S with module IO = IO = struct
         in
         send_response self r)
 
-  let handle_response (self : t) (r : Jsonrpc.Response.t) : (unit, string) result IO.t =
+  let handle_response (self : t) (r : Jsonrpc.Response.t) : unit IO.t =
     match Hashtbl.find_opt self.pending_responses r.id with
     | None ->
       IO.failwith
@@ -258,23 +253,23 @@ module Make (IO : IO) : S with module IO = IO = struct
       @@ Req_id.to_string r.id
     | Some (Request_and_handler (req, handler)) ->
       let () = Hashtbl.remove self.pending_responses r.id in
-      lift_ok @@ (match r.result with
+      (match r.result with
       | Error err -> with_error_handler self (fun () -> handler @@ Error err)
       | Ok json ->
         let r = Lsp.Server_request.response_of_json req json in
         with_error_handler self (fun () -> handler @@ Ok r))
 
   let handle_batch_response (_self : t) (_rs : Jsonrpc.Response.t list) :
-      (unit, string) result IO.t =
-      lift_ok @@ IO.failwith "Unhandled: jsonrpc batch response"
+      unit IO.t =
+    IO.failwith "Unhandled: jsonrpc batch response"
 
   let handle_batch_call (_self : t)
       (_cs :
         [ `Notification of Jsonrpc.Notification.t
         | `Request of Jsonrpc.Request.t
         ]
-        list) : (unit, string) result IO.t =
-    lift_ok @@  IO.failwith "Unhandled: jsonrpc batch call"
+        list) : unit IO.t =
+    IO.failwith "Unhandled: jsonrpc batch call"
 
   (* As in [https://github.com/c-cube/linol/issues/20],
      Jsonrpc expect "params" to be object or array,
@@ -372,7 +367,7 @@ module Make (IO : IO) : S with module IO = IO = struct
   let run ?(shutdown = fun _ -> false) (self : t) : unit IO.t =
     let async f =
       self.s#spawn_query_handler f;
-      IO.return (Ok ())
+      IO.return ()
     in
 
     let process_msg r =
@@ -394,7 +389,7 @@ module Make (IO : IO) : S with module IO = IO = struct
         let* r = read_msg self in
         match r with
         | Ok r ->
-          let* _res = process_msg r in
+          let* () = process_msg r in
           loop ()
         | Error (e, bt) -> IO.fail e bt
     in
diff --git a/src/jsonrpc2.mli b/src/jsonrpc2.mli
index 68efa11c..2ed46b04 100644
--- a/src/jsonrpc2.mli
+++ b/src/jsonrpc2.mli
@@ -27,7 +27,7 @@ module type S = sig
     ?on_received:(json -> unit) -> ?on_sent:(json -> unit) -> env:IO.env -> server -> t
   (** Create a connection using stdin/stdout *)
 
-  val send_server_notification : t -> Lsp.Server_notification.t -> (unit, string) result IO.t
+  val send_server_notification : t -> Lsp.Server_notification.t -> unit IO.t
   (** Send a notification from the server.
       @since 0.5 *)
 
diff --git a/src/lwt/linol_lwt.ml b/src/lwt/linol_lwt.ml
index 60963389..432527bf 100644
--- a/src/lwt/linol_lwt.ml
+++ b/src/lwt/linol_lwt.ml
@@ -39,16 +39,12 @@ end
 
 (** Spawn function.
     @since 0.5 *)
-let spawn (f: unit -> (unit, string) result Lwt.t) : unit =
-  let g = (fun () ->
-    let _ = Lwt.catch f (fun exn ->
-      Printf.eprintf "uncaught exception in `spawn`:\n%s\n%!"
-        (Printexc.to_string exn);
-      Lwt.return (Error (Printexc.to_string exn)))
-    in
-    Lwt.return ())
-  in
-  Lwt.async g
+let spawn f =
+  Lwt.async (fun () ->
+      Lwt.catch f (fun exn ->
+          Printf.eprintf "uncaught exception in `spawn`:\n%s\n%!"
+            (Printexc.to_string exn);
+          Lwt.return ()))
 
 include Lsp.Types
 include IO_lwt
diff --git a/src/server.ml b/src/server.ml
index f581ba32..97700568 100644
--- a/src/server.ml
+++ b/src/server.ml
@@ -32,11 +32,6 @@ module Make (IO : IO) = struct
   module DiagnosticSeverity = DiagnosticSeverity
   module Req_id = Req_id
 
-  let[@inline] lift_ok x =
-    let open IO in
-    let+ x = x in
-    Ok x
-
   (** A variant carrying a [Lsp.Server_request.t] and a handler for its return
       value. The request is stored in order to allow us to discriminate its
       existential variable. *)
@@ -54,14 +49,14 @@ module Make (IO : IO) = struct
   class virtual base_server =
     object
       method virtual on_notification
-          : notify_back:(Lsp.Server_notification.t -> (unit, string) result IO.t) ->
+          : notify_back:(Lsp.Server_notification.t -> unit IO.t) ->
             server_request:send_request ->
             Lsp.Client_notification.t ->
-            (unit, string) result IO.t
+            unit IO.t
 
       method virtual on_request
           : 'a.
-            notify_back:(Lsp.Server_notification.t -> (unit, string) result IO.t) ->
+            notify_back:(Lsp.Server_notification.t -> unit IO.t) ->
             server_request:send_request ->
             id:Req_id.t ->
             'a Lsp.Client_request.t ->
@@ -74,30 +69,30 @@ module Make (IO : IO) = struct
       method must_quit = false
       (** Set to true if the client requested to exit *)
 
-      method virtual spawn_query_handler : (unit -> (unit, string) result IO.t) -> unit
+      method virtual spawn_query_handler : (unit -> unit IO.t) -> unit
       (** How to start a new future/task/thread concurrently. This is used
           to process incoming user queries.
           @since 0.5 *)
     end
 
-  let async (self : #base_server) f : (unit, string) result IO.t =
+  let async (self : #base_server) f : unit IO.t =
     self#spawn_query_handler f;
-    IO.return (Ok ())
+    IO.return ()
 
   (** A wrapper to more easily reply to notifications *)
-  class notify_back ~(notify_back: Lsp.Server_notification.t -> (unit, string) result IO.t) ~server_request ~workDoneToken
+  class notify_back ~notify_back ~server_request ~workDoneToken
     ~partialResultToken:_ ?version ?(uri : DocumentUri.t option) () =
     object
       val mutable uri = uri
       method set_uri u = uri <- Some u
       method get_uri = uri
 
-      method send_log_msg ~type_ msg : (unit, string) result IO.t =
+      method send_log_msg ~type_ msg : unit IO.t =
         let params = LogMessageParams.create ~type_ ~message:msg in
         notify_back (Lsp.Server_notification.LogMessage params)
       (** Send a log message to the editor *)
 
-      method send_diagnostic (l : Diagnostic.t list) : (unit, string) result IO.t =
+      method send_diagnostic (l : Diagnostic.t list) : unit IO.t =
         match uri with
         | None ->
           IO.failwith "notify_back: cannot publish diagnostics, no URI given"
@@ -108,33 +103,33 @@ module Make (IO : IO) = struct
           notify_back (Lsp.Server_notification.PublishDiagnostics params)
       (** Send diagnostics for the current document *)
 
-      method telemetry json : (unit, string) result IO.t =
+      method telemetry json : unit IO.t =
         notify_back @@ Lsp.Server_notification.TelemetryNotification json
 
-      method cancel_request (id : Jsonrpc.Id.t) : (unit, string) result IO.t =
+      method cancel_request (id : Jsonrpc.Id.t) : unit IO.t =
         notify_back @@ CancelRequest id
 
       method work_done_progress_begin (p : Lsp.Types.WorkDoneProgressBegin.t)
-          : (unit, string) result IO.t =
+          : unit IO.t =
         match workDoneToken with
         | Some token ->
           notify_back @@ WorkDoneProgress { token; value = Begin p }
-        | None -> lift_ok @@ IO.return ()
+        | None -> IO.return ()
 
       method work_done_progress_report (p : Lsp.Types.WorkDoneProgressReport.t)
-          : (unit, string) result IO.t =
+          : unit IO.t =
         match workDoneToken with
         | Some token ->
           notify_back @@ WorkDoneProgress { value = Report p; token }
-        | None -> lift_ok @@ IO.return ()
+        | None -> IO.return ()
 
       method work_done_progress_end (p : Lsp.Types.WorkDoneProgressEnd.t)
-          : (unit, string) result IO.t =
+          : unit IO.t =
         match workDoneToken with
         | Some token -> notify_back @@ WorkDoneProgress { value = End p; token }
-        | None -> lift_ok @@ IO.return ()
+        | None -> IO.return ()
 
-      method send_notification (n : Lsp.Server_notification.t) : (unit, string) result IO.t =
+      method send_notification (n : Lsp.Server_notification.t) : unit IO.t =
         notify_back n
       (** Send a notification from the server to the client (general purpose method) *)
 
@@ -155,6 +150,11 @@ module Make (IO : IO) = struct
   }
   (** Current state of a document. *)
 
+  let[@inline] lift_ok x =
+    let open IO in
+    let+ x = x in
+    Ok x
+
   (** An easily overloadable class. Pick the methods you want to support.
       The user must provide at least the callbacks for document lifecycle:
       open, close, update. The most basic LSP server should check documents
@@ -338,7 +338,7 @@ module Make (IO : IO) = struct
         @since 0.7 *)
 
       method on_request : type r.
-          notify_back:(Lsp.Server_notification.t -> (unit, string) result IO.t) ->
+          notify_back:_ ->
           server_request:_ ->
           id:Req_id.t ->
           r Lsp.Client_request.t ->
@@ -553,11 +553,11 @@ module Make (IO : IO) = struct
           : notify_back:notify_back ->
             TextDocumentItem.t ->
             content:string ->
-            (unit, string) result IO.t
+            unit IO.t
       (** Called when a document is opened *)
 
       method virtual on_notif_doc_did_close
-          : notify_back:notify_back -> TextDocumentIdentifier.t -> (unit, string) result IO.t
+          : notify_back:notify_back -> TextDocumentIdentifier.t -> unit IO.t
 
       method virtual on_notif_doc_did_change
           : notify_back:notify_back ->
@@ -565,30 +565,30 @@ module Make (IO : IO) = struct
             TextDocumentContentChangeEvent.t list ->
             old_content:string ->
             new_content:string ->
-            (unit, string) result IO.t
+            unit IO.t
       (** Called when the document changes. *)
 
       method on_notif_doc_did_save ~notify_back:(_ : notify_back)
-          (_params : DidSaveTextDocumentParams.t ) : (unit, string) result IO.t =
-        lift_ok @@ IO.return ()
+          (_params : DidSaveTextDocumentParams.t ) : unit IO.t =
+        IO.return ()
 
       method on_unknown_notification ~notify_back:(_ : notify_back)
-          (_n : Jsonrpc.Notification.t) : (unit, string) result IO.t =
-        lift_ok @@ IO.return ()
+          (_n : Jsonrpc.Notification.t) : unit IO.t =
+        IO.return ()
 
       method on_notification_unhandled ~notify_back:(_ : notify_back)
-          (_n : Lsp.Client_notification.t) : (unit, string) result IO.t =
-        lift_ok @@ IO.return ()
+          (_n : Lsp.Client_notification.t) : unit IO.t =
+        IO.return ()
       (** Override to handle unprocessed notifications *)
 
-      method on_notification ~(notify_back:Lsp.Server_notification.t -> (unit, string) result IO.t) ~server_request
-          (n : Lsp.Client_notification.t) : (unit, string) result IO.t =
+      method on_notification ~notify_back ~server_request
+          (n : Lsp.Client_notification.t) : unit IO.t =
         let@ _sp =
           Trace.with_span ~__FILE__ ~__LINE__ "linol.on-notification"
         in
 
         (* handler to catch all errors *)
-        let try_catch : (unit -> (unit, string) result IO.t) -> (unit, string) result IO.t =
+        let try_catch : (unit -> unit IO.t) -> unit IO.t =
          fun f ->
           IO.catch f (fun exn bt ->
               let msg =
@@ -597,7 +597,7 @@ module Make (IO : IO) = struct
                   (Printexc.raw_backtrace_to_string bt)
               in
               Log.err (fun k -> k "%s" msg);
-              lift_ok @@ IO.return ())
+              IO.return ())
         in
 
         try_catch @@ fun () ->
@@ -688,11 +688,11 @@ module Make (IO : IO) = struct
           Hashtbl.replace docs doc.uri new_st;
 
           async self (fun () ->
-              (self#on_notif_doc_did_change
+              self#on_notif_doc_did_change
                 ~notify_back:(notify_back : notify_back)
                 doc c
                 ~old_content:(Lsp.Text_document.text old_doc)
-                ~new_content:new_st.content))
+                ~new_content:new_st.content)
         | Lsp.Client_notification.DidSaveTextDocument params ->
           let notify_back =
             new notify_back
@@ -706,7 +706,7 @@ module Make (IO : IO) = struct
                 params)
         | Lsp.Client_notification.Exit ->
           status <- `ReceivedExit;
-          lift_ok @@ IO.return ()
+          IO.return ()
         | Lsp.Client_notification.WillSaveTextDocument _
         | Lsp.Client_notification.ChangeWorkspaceFolders _
         | Lsp.Client_notification.ChangeConfiguration _

From 8123e95a5f0a225ce5750c39e5a89fea9cbcf0a5 Mon Sep 17 00:00:00 2001
From: "Christoph M. Wintersteiger" <christoph@imandra.ai>
Date: Wed, 31 Jul 2024 16:25:01 +0100
Subject: [PATCH 6/8] Log exceptions in async notification handler

---
 src/server.ml | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/server.ml b/src/server.ml
index 97700568..390960b8 100644
--- a/src/server.ml
+++ b/src/server.ml
@@ -76,7 +76,10 @@ module Make (IO : IO) = struct
     end
 
   let async (self : #base_server) f : unit IO.t =
-    self#spawn_query_handler f;
+    self#spawn_query_handler
+      (fun () -> try f () with err ->
+        Log.err (fun k -> k "LSP async notification handler failed: %s" (Printexc.to_string err));
+        IO.return ());
     IO.return ()
 
   (** A wrapper to more easily reply to notifications *)

From 262e57e7c23493f38daa6a72761c813d1fe7a8f1 Mon Sep 17 00:00:00 2001
From: "Christoph M. Wintersteiger" <christoph@imandra.ai>
Date: Wed, 31 Jul 2024 16:57:03 +0100
Subject: [PATCH 7/8] Use IO.catch in async

---
 src/server.ml | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/src/server.ml b/src/server.ml
index 390960b8..bc808733 100644
--- a/src/server.ml
+++ b/src/server.ml
@@ -76,10 +76,14 @@ module Make (IO : IO) = struct
     end
 
   let async (self : #base_server) f : unit IO.t =
-    self#spawn_query_handler
-      (fun () -> try f () with err ->
-        Log.err (fun k -> k "LSP async notification handler failed: %s" (Printexc.to_string err));
-        IO.return ());
+    self#spawn_query_handler (fun () ->
+      IO.catch f (fun exn bt ->
+        let msg =
+          spf "LSP async notification handler failed with %s\n%s"
+            (Printexc.to_string exn)
+            (Printexc.raw_backtrace_to_string bt)
+        in
+        IO.return @@ Log.err (fun k -> k "%s" msg)));
     IO.return ()
 
   (** A wrapper to more easily reply to notifications *)

From 55a28825a33aead4cb672d100b56aaa12c1d0213 Mon Sep 17 00:00:00 2001
From: "Christoph M. Wintersteiger" <christoph@imandra.ai>
Date: Wed, 31 Jul 2024 16:58:18 +0100
Subject: [PATCH 8/8] Formatting

---
 src/server.ml | 22 +++++++++++-----------
 1 file changed, 11 insertions(+), 11 deletions(-)

diff --git a/src/server.ml b/src/server.ml
index bc808733..78fde4c7 100644
--- a/src/server.ml
+++ b/src/server.ml
@@ -77,13 +77,13 @@ module Make (IO : IO) = struct
 
   let async (self : #base_server) f : unit IO.t =
     self#spawn_query_handler (fun () ->
-      IO.catch f (fun exn bt ->
-        let msg =
-          spf "LSP async notification handler failed with %s\n%s"
-            (Printexc.to_string exn)
-            (Printexc.raw_backtrace_to_string bt)
-        in
-        IO.return @@ Log.err (fun k -> k "%s" msg)));
+        IO.catch f (fun exn bt ->
+            let msg =
+              spf "LSP async notification handler failed with %s\n%s"
+                (Printexc.to_string exn)
+                (Printexc.raw_backtrace_to_string bt)
+            in
+            IO.return @@ Log.err (fun k -> k "%s" msg)));
     IO.return ()
 
   (** A wrapper to more easily reply to notifications *)
@@ -223,7 +223,7 @@ module Make (IO : IO) = struct
             | `InlayHintOptions of InlayHintOptions.t
             | `InlayHintRegistrationOptions of InlayHintRegistrationOptions.t
             ]
-          option =
+            option =
         None
       (** Configuration for the inlay hints API. *)
 
@@ -576,7 +576,7 @@ module Make (IO : IO) = struct
       (** Called when the document changes. *)
 
       method on_notif_doc_did_save ~notify_back:(_ : notify_back)
-          (_params : DidSaveTextDocumentParams.t ) : unit IO.t =
+          (_params : DidSaveTextDocumentParams.t) : unit IO.t =
         IO.return ()
 
       method on_unknown_notification ~notify_back:(_ : notify_back)
@@ -703,8 +703,8 @@ module Make (IO : IO) = struct
         | Lsp.Client_notification.DidSaveTextDocument params ->
           let notify_back =
             new notify_back
-              ~workDoneToken:None ~partialResultToken:None ~uri:params.textDocument.uri
-              ~notify_back ~server_request ()
+              ~workDoneToken:None ~partialResultToken:None
+              ~uri:params.textDocument.uri ~notify_back ~server_request ()
           in
 
           async self (fun () ->