`SttpClientInterpreter` does not propagate headers in the http response for websockets

We are using tapir interceptor to enrich server responses with some contextual information (e.g. correlation id):

private val serverOptions = Http4sServerOptions
  .customiseInterceptors[IO]
  .prependInterceptor(LoggingContextInterceptor)
  .options

/** Tapir interceptor which extracts correlation ID from the request headers and puts correlation ID into Http response header. If
  * correlation ID is not present in the HTTP header, a new one is generated.
  */
object LoggingContextInterceptor extends RequestInterceptor[IO] 

On the client side, we wrap tapir interpreter to extract correlation id from the response header and lift the original

def SttpClientInterpreter#toClientThrowErrors: I => F[O]

to

def SttpClientInterpreterWithRandomCorrelationId#toClientThrowErrors: I => ResponseWithContext[F[O]]

/** Container holding
 *   - http response (e.g. tapir high-level case class),
 *   - logging context of the response (e.g correlation id extracted from http headers)
 *
 * Captured logging context allows to match server-side and client-side log (by correlation id).
 */
case class ResponseWithContext[A](ctx: LoggingContext, response: A)

In case of rest calls the wrapped tapir client properly extracts correlation id from the response.

In case of websockets, correlation id header is missing - the underlying SttpClientInterpreter returns a response without the header. A low-level sttp websocket client does receive the correlation id:

***** sttp websocket client
[2024-07-26 14:49:47,353] INFO  [io-compute-8] TapirWebSocketsCorrelationId:37 - [cid=XTM-QVL-HTR] sttp ws request
[2024-07-26 14:49:47,694] INFO  [io-compute-14] LoggingContextInterceptor:44 - RequestMetadata(GET,/ping,List(x-correlation-id: XTM-QVL-HTR, accept-encoding: gzip, deflate, upgrade: websocket, connection: upgrade, sec-websocket-key: C2/p0NQzhCMWpTQO5jXngQ==, sec-websocket-version: 13, origin: http://localhost:8081, host: localhost:8081, accept: */*, user-agent: AHC/2.1)), List(sttp.tapir.server.ServerEndpoint$$anon$3@249b2b39)
[2024-07-26 14:49:47,719] INFO  [io-compute-14] LoggingContextInterceptor:52 - response: Response(ServerResponse(200,Vector(x-correlation-id: XTM-QVL-HTR)))
[2024-07-26 14:49:47,741] INFO  [io-compute-22] TapirWebSocketsCorrelationId:51 - [cid=XTM-QVL-HTR] Sending 1
[2024-07-26 14:49:47,782] INFO  [io-compute-6] TapirWebSocketsCorrelationId:53 - [cid=XTM-QVL-HTR] Received 1
[2024-07-26 14:49:47,795] INFO  [io-compute-19] TapirWebSocketsCorrelationId:80 - [cid=unset] 

***** tapir websocket client
[2024-07-26 14:49:47,827] INFO  [io-compute-2] SttpClientInterpreterWithRandomCorrelationId:25 - [cid=LZE-JTD-WQP] tapir ws request
[2024-07-26 14:49:47,855] INFO  [io-compute-20] LoggingContextInterceptor:44 - RequestMetadata(GET,/ping,List(Connection: Upgrade, Host: localhost:8081, Upgrade: websocket, User-Agent: Java-http-client/21.0.3, Accept-Encoding: gzip, deflate, Sec-WebSocket-Key: sZ/mmEyrbVLoResXK3uMnA==, Sec-WebSocket-Version: 13, x-correlation-id: LZE-JTD-WQP)), List(sttp.tapir.server.ServerEndpoint$$anon$3@249b2b39)
[2024-07-26 14:49:47,856] INFO  [io-compute-20] LoggingContextInterceptor:52 - response: Response(ServerResponse(200,Vector(x-correlation-id: LZE-JTD-WQP)))
[2024-07-26 14:49:47,870] INFO  [io-compute-0] SttpClientInterpreterWithRandomCorrelationId:28 - [cid=unset] Response(sttp.tapir.client.sttp.ws.fs2.WebSocketToFs2Pipe$$Lambda/0x000071f2346ac678@70e0de77,101,,List(),List(),RequestMetadata(GET,ws://localhost:8081/ping,Vector(Accept-Encoding: gzip, deflate, x-correlation-id: LZE-JTD-WQP)))
[2024-07-26 14:49:47,875] INFO  [io-compute-12] TapirWebSocketsCorrelationId:64 - [cid=unset] Sending 1
[2024-07-26 14:49:47,882] INFO  [io-compute-1] TapirWebSocketsCorrelationId:66 - [cid=unset] Received 1

Repo with the code: https://github.com/kamilkloch/tapir-correlation-id

Start by examining SttpClientInterpreter#toClientThrowErrors and the websocket client path, then compare their response handling with regular REST calls. Use the tapir-correlation-id reproduction repository to confirm the response header is available and identify where it is lost; done means websocket responses expose the correlation-id header to the wrapper.