fix trailers on the first response

dialohq · Nov 17, 2024 · 2349843 · 2349843
1 parent b4e85db
commit 2349843
Showing 1 changed file with 13 additions and 4 deletions.
diff --git a/lib/eio/io-client-h2-ocaml-protoc/io_client_h2_ocaml_protoc.ml b/lib/eio/io-client-h2-ocaml-protoc/io_client_h2_ocaml_protoc.ml
@@ -120,10 +120,19 @@ module Make_net (Client : Client) :
     (* Allocate once, use a pool of these *)
     let errored = ref false in
     let response_handler response reader =
-      let trailers, trailers_u = Eio.Promise.create () in
-      let () =
-        trailers_handler :=
-          fun trailers -> Eio.Promise.resolve trailers_u trailers
+      let grpc_status_header =
+        H2.Headers.get response.H2.Response.headers "grpc-status"
+      in
+      let trailers =
+        match grpc_status_header with
+        | Some status ->
+            Eio.Promise.create_resolved
+            @@ H2.Headers.of_list [ ("grpc-status", status) ]
+        | None ->
+            let trailers', trailers_u = Eio.Promise.create () in
+            (trailers_handler :=
+               fun trailers -> Eio.Promise.resolve trailers_u trailers);
+            trailers'
       in
 
       let next =