fix(lsp): wait for cancellation responses

this is more spec complying Signed-off-by: Rudi Grinberg <[email protected]> ps-id: C29A2B64-D590-41EA-B93E-3BE361F24588
garnermt · Jun 13, 2022 · 7ee8f6b · 7ee8f6b
1 parent 398c31e
commit 7ee8f6b
Showing 1 changed file with 6 additions and 7 deletions.
diff --git a/lsp-fiber/src/rpc.ml b/lsp-fiber/src/rpc.ml
@@ -277,13 +277,12 @@ struct
       Fiber.Cancel.with_handler cancel
         ~on_cancel:(fun () -> on_cancel jsonrpc_req.id)
         (fun () ->
-          let _, req_f =
-            Session.request_with_cancel (Fdecl.get t.session) jsonrpc_req
-          in
-          let+ resp = req_f in
-          match resp with
-          | `Cancelled -> `Cancelled
-          | `Ok resp -> `Ok (receive_response req resp))
+          let+ resp = Session.request (Fdecl.get t.session) jsonrpc_req in
+          match resp.result with
+          | Error { code = RequestCancelled; _ } -> `Cancelled
+          | Ok _ when Fiber.Cancel.fired cancel -> `Cancelled
+          | Ok s -> `Ok (Out_request.response_of_json req s)
+          | Error e -> raise (Jsonrpc.Response.Error.E e))
     in
     match cancel_status with
     | Cancelled () -> `Cancelled