From 1a6230ab56d4a7706562c05438257673edc25358 Mon Sep 17 00:00:00 2001
From: Daniel Vigovszky <daniel.vigovszky@gmail.com>
Date: Fri, 28 Jul 2023 12:22:03 +0100
Subject: [PATCH 1/4] SocketApp refactoring, ability to provide custom
 websocket config

---
 .../src/main/scala/zio/http/TestClient.scala  | 10 +--
 .../scala/zio/http/SocketContractSpec.scala   | 12 +--
 .../test/scala/zio/http/TestClientSpec.scala  |  4 +-
 .../src/main/scala/zio/http/Handler.scala     | 63 ++++----------
 .../src/main/scala/zio/http/Response.scala    |  8 --
 .../src/main/scala/zio/http/SocketApp.scala   | 84 +++++++++++++++++++
 .../src/main/scala/zio/http/ZClient.scala     |  2 +-
 .../http/netty/client/NettyClientDriver.scala |  5 +-
 .../netty/server/ServerInboundHandler.scala   | 12 ++-
 .../src/main/scala/zio/http/package.scala     |  2 -
 .../test/scala/zio/http/WebSocketSpec.scala   |  5 +-
 11 files changed, 129 insertions(+), 78 deletions(-)
 create mode 100644 zio-http/src/main/scala/zio/http/SocketApp.scala

diff --git a/zio-http-testkit/src/main/scala/zio/http/TestClient.scala b/zio-http-testkit/src/main/scala/zio/http/TestClient.scala
index d6a2599090..ef40307e09 100644
--- a/zio-http-testkit/src/main/scala/zio/http/TestClient.scala
+++ b/zio-http-testkit/src/main/scala/zio/http/TestClient.scala
@@ -127,13 +127,13 @@ final case class TestClient(
       promise               <- Promise.make[Nothing, Unit]
       testChannelClient     <- TestChannel.make(in, out, promise)
       testChannelServer     <- TestChannel.make(out, in, promise)
-      _                     <- currentSocketBehavior.runZIO(testChannelClient).forkDaemon
-      _                     <- app.provideEnvironment(env).runZIO(testChannelServer).forkDaemon
+      _                     <- currentSocketBehavior.handler.runZIO(testChannelClient).forkDaemon
+      _                     <- app.provideEnvironment(env).handler.runZIO(testChannelServer).forkDaemon
     } yield Response.status(Status.SwitchingProtocols)
   }
 
   def installSocketApp[Env1](
-    app: Handler[Any, Throwable, WebSocketChannel, Unit],
+    app: SocketApp[Any],
   ): ZIO[Env1, Nothing, Unit] =
     for {
       env <- ZIO.environment[Env1]
@@ -182,7 +182,7 @@ object TestClient {
     ZIO.serviceWithZIO[TestClient](_.addHandler(handler))
 
   def installSocketApp(
-    app: Handler[Any, Throwable, WebSocketChannel, Unit],
+    app: SocketApp[Any],
   ): ZIO[TestClient, Nothing, Unit] =
     ZIO.serviceWithZIO[TestClient](_.installSocketApp(app))
 
@@ -190,7 +190,7 @@ object TestClient {
     ZLayer.scopedEnvironment {
       for {
         behavior       <- Ref.make[PartialFunction[Request, ZIO[Any, Response, Response]]](PartialFunction.empty)
-        socketBehavior <- Ref.make[SocketApp[Any]](Handler.unit)
+        socketBehavior <- Ref.make[SocketApp[Any]](SocketApp.unit)
         driver = TestClient(behavior, socketBehavior)
       } yield ZEnvironment[TestClient, Client](driver, ZClient.fromDriver(driver))
     }
diff --git a/zio-http-testkit/src/test/scala/zio/http/SocketContractSpec.scala b/zio-http-testkit/src/test/scala/zio/http/SocketContractSpec.scala
index c2653745be..04eb0a28fe 100644
--- a/zio-http-testkit/src/test/scala/zio/http/SocketContractSpec.scala
+++ b/zio-http-testkit/src/test/scala/zio/http/SocketContractSpec.scala
@@ -14,7 +14,7 @@ object SocketContractSpec extends ZIOSpecDefault {
   def spec: Spec[Any, Any] =
     suite("SocketOps")(
       contract("Successful Multi-message application") { p =>
-        val socketServer: Handler[Any, Throwable, WebSocketChannel, Unit] =
+        val socketServer: SocketApp[Any] =
           Handler.webSocket { channel =>
             channel.receiveAll {
               case Read(WebSocketFrame.Text("Hi Server"))          =>
@@ -31,7 +31,7 @@ object SocketContractSpec extends ZIOSpecDefault {
 
         socketServer
       } { _ =>
-        val socketClient: Handler[Any, Throwable, WebSocketChannel, Unit] =
+        val socketClient: SocketApp[Any] =
           Handler.webSocket { channel =>
             channel.receiveAll {
               case ChannelEvent.Read(WebSocketFrame.Text("Hi Client")) =>
@@ -89,8 +89,8 @@ object SocketContractSpec extends ZIOSpecDefault {
   private def contract(
     name: String,
   )(
-    serverApp: Promise[Throwable, Unit] => Handler[Any, Throwable, WebSocketChannel, Unit],
-  )(clientApp: Promise[Throwable, Unit] => Handler[Any, Throwable, WebSocketChannel, Unit]) = {
+    serverApp: Promise[Throwable, Unit] => SocketApp[Any],
+  )(clientApp: Promise[Throwable, Unit] => SocketApp[Any]) = {
     suite(name)(
       test("Live") {
         for {
@@ -123,7 +123,7 @@ object SocketContractSpec extends ZIOSpecDefault {
   }
 
   private def liveServerSetup(
-    serverApp: Promise[Throwable, Unit] => Handler[Any, Throwable, WebSocketChannel, Unit],
+    serverApp: Promise[Throwable, Unit] => SocketApp[Any],
   ): ZIO[Server, Nothing, (RuntimeFlags, Promise[Throwable, Unit])] =
     ZIO.serviceWithZIO[Server](server =>
       for {
@@ -133,7 +133,7 @@ object SocketContractSpec extends ZIOSpecDefault {
     )
 
   private def testServerSetup(
-    serverApp: Promise[Throwable, Unit] => Handler[Any, Throwable, WebSocketChannel, Unit],
+    serverApp: Promise[Throwable, Unit] => SocketApp[Any],
   ): ZIO[TestClient, Nothing, (RuntimeFlags, Promise[Throwable, Unit])] =
     for {
       p <- Promise.make[Throwable, Unit]
diff --git a/zio-http-testkit/src/test/scala/zio/http/TestClientSpec.scala b/zio-http-testkit/src/test/scala/zio/http/TestClientSpec.scala
index 74d238a323..180a1a0427 100644
--- a/zio-http-testkit/src/test/scala/zio/http/TestClientSpec.scala
+++ b/zio-http-testkit/src/test/scala/zio/http/TestClientSpec.scala
@@ -66,7 +66,7 @@ object TestClientSpec extends ZIOSpecDefault {
       ),
       suite("socket ops")(
         test("happy path") {
-          val socketClient: Handler[Any, Throwable, WebSocketChannel, Unit] =
+          val socketClient: SocketApp[Any] =
             Handler.webSocket { channel =>
               channel.receiveAll {
                 case ChannelEvent.Read(WebSocketFrame.Text("Hi Client")) =>
@@ -77,7 +77,7 @@ object TestClientSpec extends ZIOSpecDefault {
               }
             }
 
-          val socketServer: Handler[Any, Throwable, WebSocketChannel, Unit] =
+          val socketServer: SocketApp[Any] =
             Handler.webSocket { channel =>
               channel.receiveAll {
                 case ChannelEvent.Read(WebSocketFrame.Text("Hi Server")) =>
diff --git a/zio-http/src/main/scala/zio/http/Handler.scala b/zio-http/src/main/scala/zio/http/Handler.scala
index 93927d4a99..ec395eae44 100644
--- a/zio-http/src/main/scala/zio/http/Handler.scala
+++ b/zio-http/src/main/scala/zio/http/Handler.scala
@@ -222,34 +222,6 @@ sealed trait Handler[-R, +Err, -In, +Out] { self =>
   )(implicit trace: Trace): Handler[R1, Err1, In1, Out] =
     that.andThen(self)
 
-  /**
-   * Creates a socket connection on the provided URL. Typically used to connect
-   * as a client.
-   */
-  def connect(
-    url: String,
-    headers: Headers = Headers.empty,
-  )(implicit
-    ev: Err <:< Throwable,
-    ev2: WebSocketChannel <:< In,
-    trace: Trace,
-  ): ZIO[R with Client with Scope, Throwable, Response] =
-    ZIO.fromEither(URL.decode(url)).orDie.flatMap(connect(_, headers))
-
-  def connect(
-    url: URL,
-    headers: Headers,
-  )(implicit
-    ev1: Err <:< Throwable,
-    ev2: WebSocketChannel <:< In,
-    trace: Trace,
-  ): ZIO[R with Client with Scope, Throwable, Response] =
-    ZIO.serviceWithZIO[Client] { client =>
-      val client2 = if (url.isAbsolute) client.url(url) else client.addUrl(url)
-
-      client2.addHeaders(headers).socket(self.asInstanceOf[SocketApp[R]])
-    }
-
   /**
    * Transforms the input of the handler before passing it on to the current
    * Handler
@@ -621,21 +593,6 @@ sealed trait Handler[-R, +Err, -In, +Out] { self =>
     HttpApp(Routes.singleton(handler.contramap[(Path, Request)](_._2)))
   }
 
-  def toHttpAppWS(implicit err: Err <:< Throwable, in: WebSocketChannel <:< In, trace: Trace): HttpApp[R] =
-    Handler.fromZIO(self.toResponse).toHttpApp
-
-  /**
-   * Creates a new response from a socket handler.
-   */
-  def toResponse(implicit
-    ev1: Err <:< Throwable,
-    ev2: WebSocketChannel <:< In,
-    trace: Trace,
-  ): ZIO[R, Nothing, Response] =
-    ZIO.environment[R].flatMap { env =>
-      Response.fromSocketApp(self.asInstanceOf[SocketApp[R]].provideEnvironment(env))
-    }
-
   /**
    * Takes some defects and converts them into failures.
    */
@@ -727,6 +684,18 @@ object Handler {
   def badRequest(message: => String): Handler[Any, Nothing, Any, Response] =
     error(Status.BadRequest, message)
 
+  /**
+   * Constructs a handler from two functions, one that configures web socket and
+   * another that uses a web socket.
+   *
+   * If the config function returns with None, the server configuration is used.
+   */
+  final def customWebSocket[Env](
+    config: Request => ZIO[Env, Throwable, Option[WebSocketConfig]],
+    f: WebSocketChannel => ZIO[Env, Throwable, Any],
+  ): SocketApp[Env] =
+    SocketApp(Handler.fromFunctionZIO(f), Handler.fromFunctionZIO(config))
+
   /**
    * Returns a handler that dies with the specified `Throwable`. This method can
    * be used for terminating an handler because a defect has been detected in
@@ -1102,10 +1071,10 @@ object Handler {
   /**
    * Constructs a handler from a function that uses a web socket.
    */
-  final def webSocket[Env, Err, Out](
-    f: WebSocketChannel => ZIO[Env, Err, Out],
-  ): Handler[Env, Err, WebSocketChannel, Out] =
-    Handler.fromFunctionZIO(f)
+  final def webSocket[Env](
+    f: WebSocketChannel => ZIO[Env, Throwable, Any],
+  ): SocketApp[Env] =
+    SocketApp(Handler.fromFunctionZIO(f))
 
   final implicit class RequestHandlerSyntax[-R, +Err](val self: RequestHandler[R, Err])
       extends HeaderModifier[RequestHandler[R, Err]] {
diff --git a/zio-http/src/main/scala/zio/http/Response.scala b/zio-http/src/main/scala/zio/http/Response.scala
index 4a608a31bc..daa250205a 100644
--- a/zio-http/src/main/scala/zio/http/Response.scala
+++ b/zio-http/src/main/scala/zio/http/Response.scala
@@ -357,14 +357,6 @@ object Response {
   def fromServerSentEvents(data: ZStream[Any, Nothing, ServerSentEvent])(implicit trace: Trace): Response =
     new BasicResponse(Body.fromStream(data.map(_.encode)), contentTypeEventStream, Status.Ok)
 
-  /**
-   * Creates a new response for the provided socket
-   */
-  def fromSocket[R](
-    http: Handler[R, Throwable, WebSocketChannel, Any],
-  )(implicit trace: Trace): ZIO[R, Nothing, Response] =
-    fromSocketApp(http)
-
   /**
    * Creates a new response for the provided socket app
    */
diff --git a/zio-http/src/main/scala/zio/http/SocketApp.scala b/zio-http/src/main/scala/zio/http/SocketApp.scala
new file mode 100644
index 0000000000..1b2a70b879
--- /dev/null
+++ b/zio-http/src/main/scala/zio/http/SocketApp.scala
@@ -0,0 +1,84 @@
+package zio.http
+
+import zio.{Cause, Scope, Tag, Trace, ZEnvironment, ZIO, ZLayer}
+
+final case class SocketApp[-R](
+  handler: Handler[R, Throwable, WebSocketChannel, Any],
+  customConfig: Handler[R, Throwable, Request, Option[WebSocketConfig]],
+) { self =>
+
+  /**
+   * Creates a socket connection on the provided URL. Typically used to connect
+   * as a client.
+   */
+  def connect(
+    url: String,
+    headers: Headers = Headers.empty,
+  )(implicit
+    trace: Trace,
+  ): ZIO[R with Client with Scope, Throwable, Response] =
+    ZIO.fromEither(URL.decode(url)).orDie.flatMap(connect(_, headers))
+
+  def connect(
+    url: URL,
+    headers: Headers,
+  )(implicit
+    trace: Trace,
+  ): ZIO[R with Client with Scope, Throwable, Response] =
+    ZIO.serviceWithZIO[Client] { client =>
+      val client2 = if (url.isAbsolute) client.url(url) else client.addUrl(url)
+
+      client2.addHeaders(headers).socket(self)
+    }
+
+  def provideEnvironment(r: ZEnvironment[R])(implicit trace: Trace): SocketApp[Any] =
+    SocketApp(handler.provideEnvironment(r), customConfig.provideEnvironment(r))
+
+  def provideLayer[R0](layer: ZLayer[R0, Throwable, R])(implicit
+    trace: Trace,
+  ): SocketApp[R0] =
+    SocketApp(handler.provideLayer(layer), customConfig.provideLayer(layer))
+
+  def provideSomeEnvironment[R1](f: ZEnvironment[R1] => ZEnvironment[R])(implicit
+    trace: Trace,
+  ): SocketApp[R1] =
+    SocketApp(handler.provideSomeEnvironment(f), customConfig.provideSomeEnvironment(f))
+
+  def provideSomeLayer[R0, R1: Tag](
+    layer: ZLayer[R0, Throwable, R1],
+  )(implicit ev: R0 with R1 <:< R, trace: Trace): SocketApp[R0] =
+    SocketApp(handler.provideSomeLayer(layer), customConfig.provideSomeLayer(layer))
+
+  def tapErrorCauseZIO[R1 <: R](
+    f: Cause[Throwable] => ZIO[R1, Throwable, Any],
+  )(implicit trace: Trace): SocketApp[R1] =
+    SocketApp(handler.tapErrorCauseZIO(f), customConfig.tapErrorCauseZIO(f))
+
+  /**
+   * Returns a Handler that effectfully peeks at the failure of this SocketApp.
+   */
+  def tapErrorZIO[R1 <: R](
+    f: Throwable => ZIO[R1, Throwable, Any],
+  )(implicit trace: Trace): SocketApp[R1] =
+    self.tapErrorCauseZIO(cause => cause.failureOption.fold[ZIO[R1, Throwable, Any]](ZIO.unit)(f))
+
+  /**
+   * Creates a new response from a socket handler.
+   */
+  def toResponse(implicit
+    trace: Trace,
+  ): ZIO[R, Nothing, Response] =
+    ZIO.environment[R].flatMap { env =>
+      Response.fromSocketApp(self.provideEnvironment(env))
+    }
+
+  def toHttpAppWS(implicit trace: Trace): HttpApp[R] =
+    Handler.fromZIO(self.toResponse).toHttpApp
+}
+
+object SocketApp {
+  def apply[R](handler: Handler[R, Throwable, WebSocketChannel, Any]): SocketApp[R] =
+    SocketApp(handler, Handler.succeed(None))
+
+  val unit: SocketApp[Any] = SocketApp(Handler.unit)
+}
diff --git a/zio-http/src/main/scala/zio/http/ZClient.scala b/zio-http/src/main/scala/zio/http/ZClient.scala
index 0f32687c7c..82284401ce 100644
--- a/zio-http/src/main/scala/zio/http/ZClient.scala
+++ b/zio-http/src/main/scala/zio/http/ZClient.scala
@@ -633,7 +633,7 @@ object ZClient {
       val request = Request(version, method, url, headers, body, None)
       val cfg     = sslConfig.fold(config)(config.ssl)
 
-      requestAsync(request, cfg, () => Handler.unit, None)
+      requestAsync(request, cfg, () => SocketApp.unit, None)
     }
 
     def socket[Env1](
diff --git a/zio-http/src/main/scala/zio/http/netty/client/NettyClientDriver.scala b/zio-http/src/main/scala/zio/http/netty/client/NettyClientDriver.scala
index 807baddadf..b3c3b187b1 100644
--- a/zio-http/src/main/scala/zio/http/netty/client/NettyClientDriver.scala
+++ b/zio-http/src/main/scala/zio/http/netty/client/NettyClientDriver.scala
@@ -68,7 +68,8 @@ final case class NettyClientDriver private (
         nettyChannel     = NettyChannel.make[JWebSocketFrame](channel)
         webSocketChannel = WebSocketChannel.make(nettyChannel, queue)
         app              = createSocketApp()
-        _ <- app.runZIO(webSocketChannel).ignoreLogged.interruptible.forkScoped
+        customConfig <- app.customConfig.runZIO(req)
+        _            <- app.handler.runZIO(webSocketChannel).ignoreLogged.interruptible.forkScoped
       } yield {
         val pipeline                              = channel.pipeline()
         val toRemove: mutable.Set[ChannelHandler] = new mutable.HashSet[ChannelHandler]()
@@ -85,7 +86,7 @@ final case class NettyClientDriver private (
 
           val headers = Conversions.headersToNetty(req.headers)
           val config  = NettySocketProtocol
-            .clientBuilder(webSocketConfig)
+            .clientBuilder(customConfig.getOrElse(webSocketConfig))
             .customHeaders(headers)
             .webSocketUri(req.url.encode)
             .build()
diff --git a/zio-http/src/main/scala/zio/http/netty/server/ServerInboundHandler.scala b/zio-http/src/main/scala/zio/http/netty/server/ServerInboundHandler.scala
index 3a59ffb705..757387c18f 100644
--- a/zio-http/src/main/scala/zio/http/netty/server/ServerInboundHandler.scala
+++ b/zio-http/src/main/scala/zio/http/netty/server/ServerInboundHandler.scala
@@ -288,20 +288,26 @@ private[zio] final case class ServerInboundHandler(
     runtime: NettyRuntime,
   ): Unit = {
     val app = res.socketApp
+
     jReq match {
       case jReq: FullHttpRequest =>
         val queue = runtime.runtime(ctx).unsafe.run(Queue.unbounded[WebSocketChannelEvent]).getOrThrowFiberFailure()
+        val webSocketApp = app.getOrElse(SocketApp.unit)
+        val request      = makeZioRequest(ctx, jReq)
+        val customConfig =
+          runtime.runtime(ctx).unsafe.run { webSocketApp.customConfig.runZIO(request) }.getOrThrowFiberFailure()
         runtime.runtime(ctx).unsafe.run {
           val nettyChannel     = NettyChannel.make[JWebSocketFrame](ctx.channel())
           val webSocketChannel = WebSocketChannel.make(nettyChannel, queue)
-          val webSocketApp     = app.getOrElse(Handler.unit)
-          webSocketApp.runZIO(webSocketChannel).ignoreLogged.forkDaemon
+          webSocketApp.handler.runZIO(webSocketChannel).ignoreLogged.forkDaemon
         }
         ctx
           .channel()
           .pipeline()
           .addLast(
-            new WebSocketServerProtocolHandler(NettySocketProtocol.serverBuilder(config.webSocketConfig).build()),
+            new WebSocketServerProtocolHandler(
+              NettySocketProtocol.serverBuilder(customConfig.getOrElse(config.webSocketConfig)).build(),
+            ),
           )
           .addLast(Names.WebSocketHandler, new WebSocketAppHandler(runtime, queue, None))
 
diff --git a/zio-http/src/main/scala/zio/http/package.scala b/zio-http/src/main/scala/zio/http/package.scala
index 7e23367b65..5ac8a96abe 100644
--- a/zio-http/src/main/scala/zio/http/package.scala
+++ b/zio-http/src/main/scala/zio/http/package.scala
@@ -57,8 +57,6 @@ package object http extends UrlInterpolator {
   type Client = ZClient[Any, Body, Throwable, Response]
   def Client: ZClient.type = ZClient
 
-  type SocketApp[-R] = Handler[R, Throwable, WebSocketChannel, Any]
-
   /**
    * A channel that allows websocket frames to be written to it.
    */
diff --git a/zio-http/src/test/scala/zio/http/WebSocketSpec.scala b/zio-http/src/test/scala/zio/http/WebSocketSpec.scala
index 73f7ef1c86..4c312ed4c2 100644
--- a/zio-http/src/test/scala/zio/http/WebSocketSpec.scala
+++ b/zio-http/src/test/scala/zio/http/WebSocketSpec.scala
@@ -111,9 +111,10 @@ object WebSocketSpec extends HttpRunnableSpec {
       } yield assertCompletes
     } @@ nonFlaky,
     test("Multiple websocket upgrades") {
-      val app   = Handler.succeed(WebSocketFrame.text("BAR")).toHttpAppWS.deployWS
+      val app   =
+        Handler.webSocket(channel => channel.send(ChannelEvent.Read(WebSocketFrame.text("BAR")))).toHttpAppWS.deployWS
       val codes = ZIO
-        .foreach(1 to 1024)(_ => app.runZIO(Handler.unit).map(_.status))
+        .foreach(1 to 1024)(_ => app.runZIO(SocketApp.unit).map(_.status))
         .map(_.count(_ == Status.SwitchingProtocols))
 
       assertZIO(codes)(equalTo(1024))

From eedf1da6d2ed1b8a59db1c338cd1b4a0dd17805e Mon Sep 17 00:00:00 2001
From: Daniel Vigovszky <daniel.vigovszky@gmail.com>
Date: Fri, 28 Jul 2023 14:54:57 +0100
Subject: [PATCH 2/4] Format

---
 zio-http/src/main/scala/zio/http/SocketApp.scala | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/zio-http/src/main/scala/zio/http/SocketApp.scala b/zio-http/src/main/scala/zio/http/SocketApp.scala
index 1b2a70b879..416ca3eb5c 100644
--- a/zio-http/src/main/scala/zio/http/SocketApp.scala
+++ b/zio-http/src/main/scala/zio/http/SocketApp.scala
@@ -1,6 +1,6 @@
 package zio.http
 
-import zio.{Cause, Scope, Tag, Trace, ZEnvironment, ZIO, ZLayer}
+import zio._
 
 final case class SocketApp[-R](
   handler: Handler[R, Throwable, WebSocketChannel, Any],

From 9b37750741ce40a875698b8539b54666d8315cf9 Mon Sep 17 00:00:00 2001
From: Daniel Vigovszky <daniel.vigovszky@gmail.com>
Date: Fri, 28 Jul 2023 16:00:59 +0100
Subject: [PATCH 3/4] Cleanup

---
 .../scala/example/WebSocketAdvanced.scala     |  2 +-
 .../main/scala/example/WebSocketEcho.scala    |  2 +-
 .../example/WebSocketReconnectingClient.scala |  2 +-
 .../scala/example/WebSocketSimpleClient.scala |  2 +-
 .../src/main/scala/zio/http/TestClient.scala  | 10 ++---
 .../scala/zio/http/SocketContractSpec.scala   | 12 +++---
 .../test/scala/zio/http/TestClientSpec.scala  |  4 +-
 .../main/scala/zio/http/ClientDriver.scala    |  2 +-
 .../src/main/scala/zio/http/Handler.scala     | 16 +-------
 .../src/main/scala/zio/http/Response.scala    |  8 ++--
 .../{SocketApp.scala => WebSocketApp.scala}   | 37 ++++++++++---------
 .../src/main/scala/zio/http/ZClient.scala     | 20 +++++-----
 .../main/scala/zio/http/ZClientAspect.scala   |  8 ++--
 .../http/netty/client/NettyClientDriver.scala |  7 ++--
 .../netty/server/ServerInboundHandler.scala   |  7 +---
 .../test/scala/zio/http/WebSocketSpec.scala   |  2 +-
 .../zio/http/internal/HttpRunnableSpec.scala  |  4 +-
 17 files changed, 66 insertions(+), 79 deletions(-)
 rename zio-http/src/main/scala/zio/http/{SocketApp.scala => WebSocketApp.scala} (68%)

diff --git a/zio-http-example/src/main/scala/example/WebSocketAdvanced.scala b/zio-http-example/src/main/scala/example/WebSocketAdvanced.scala
index 9d41d307b0..f806507a94 100644
--- a/zio-http-example/src/main/scala/example/WebSocketAdvanced.scala
+++ b/zio-http-example/src/main/scala/example/WebSocketAdvanced.scala
@@ -8,7 +8,7 @@ import zio.http.codec.PathCodec.string
 
 object WebSocketAdvanced extends ZIOAppDefault {
 
-  val socketApp: SocketApp[Any] =
+  val socketApp: WebSocketApp[Any] =
     Handler.webSocket { channel =>
       channel.receiveAll {
         case Read(WebSocketFrame.Text("end"))                =>
diff --git a/zio-http-example/src/main/scala/example/WebSocketEcho.scala b/zio-http-example/src/main/scala/example/WebSocketEcho.scala
index 075c6321f8..df3e34a379 100644
--- a/zio-http-example/src/main/scala/example/WebSocketEcho.scala
+++ b/zio-http-example/src/main/scala/example/WebSocketEcho.scala
@@ -7,7 +7,7 @@ import zio.http._
 import zio.http.codec.PathCodec.string
 
 object WebSocketEcho extends ZIOAppDefault {
-  private val socketApp: SocketApp[Any] =
+  private val socketApp: WebSocketApp[Any] =
     Handler.webSocket { channel =>
       channel.receiveAll {
         case Read(WebSocketFrame.Text("FOO")) =>
diff --git a/zio-http-example/src/main/scala/example/WebSocketReconnectingClient.scala b/zio-http-example/src/main/scala/example/WebSocketReconnectingClient.scala
index 60c93380ae..a33ef8499a 100644
--- a/zio-http-example/src/main/scala/example/WebSocketReconnectingClient.scala
+++ b/zio-http-example/src/main/scala/example/WebSocketReconnectingClient.scala
@@ -10,7 +10,7 @@ object WebSocketReconnectingClient extends ZIOAppDefault {
   val url = "ws://ws.vi-server.org/mirror"
 
   // A promise is used to be able to notify application about websocket errors
-  def makeSocketApp(p: Promise[Nothing, Throwable]): SocketApp[Any] =
+  def makeSocketApp(p: Promise[Nothing, Throwable]): WebSocketApp[Any] =
     Handler
 
       // Listen for all websocket channel events
diff --git a/zio-http-example/src/main/scala/example/WebSocketSimpleClient.scala b/zio-http-example/src/main/scala/example/WebSocketSimpleClient.scala
index e1f538cf42..ef7f4903ff 100644
--- a/zio-http-example/src/main/scala/example/WebSocketSimpleClient.scala
+++ b/zio-http-example/src/main/scala/example/WebSocketSimpleClient.scala
@@ -9,7 +9,7 @@ object WebSocketSimpleClient extends ZIOAppDefault {
 
   val url = "ws://ws.vi-server.org/mirror"
 
-  val socketApp: SocketApp[Any] =
+  val socketApp: WebSocketApp[Any] =
     Handler
 
       // Listen for all websocket channel events
diff --git a/zio-http-testkit/src/main/scala/zio/http/TestClient.scala b/zio-http-testkit/src/main/scala/zio/http/TestClient.scala
index ef40307e09..4c28c3fe81 100644
--- a/zio-http-testkit/src/main/scala/zio/http/TestClient.scala
+++ b/zio-http-testkit/src/main/scala/zio/http/TestClient.scala
@@ -14,7 +14,7 @@ import zio.http.{Headers, Method, Scheme, Status, Version}
  */
 final case class TestClient(
   behavior: Ref[PartialFunction[Request, ZIO[Any, Response, Response]]],
-  serverSocketBehavior: Ref[SocketApp[Any]],
+  serverSocketBehavior: Ref[WebSocketApp[Any]],
 ) extends ZClient.Driver[Any, Throwable] {
 
   /**
@@ -117,7 +117,7 @@ final case class TestClient(
     version: Version,
     url: URL,
     headers: Headers,
-    app: SocketApp[Env1],
+    app: WebSocketApp[Env1],
   )(implicit trace: Trace): ZIO[Env1 with Scope, Throwable, Response] = {
     for {
       env                   <- ZIO.environment[Env1]
@@ -133,7 +133,7 @@ final case class TestClient(
   }
 
   def installSocketApp[Env1](
-    app: SocketApp[Any],
+    app: WebSocketApp[Any],
   ): ZIO[Env1, Nothing, Unit] =
     for {
       env <- ZIO.environment[Env1]
@@ -182,7 +182,7 @@ object TestClient {
     ZIO.serviceWithZIO[TestClient](_.addHandler(handler))
 
   def installSocketApp(
-    app: SocketApp[Any],
+    app: WebSocketApp[Any],
   ): ZIO[TestClient, Nothing, Unit] =
     ZIO.serviceWithZIO[TestClient](_.installSocketApp(app))
 
@@ -190,7 +190,7 @@ object TestClient {
     ZLayer.scopedEnvironment {
       for {
         behavior       <- Ref.make[PartialFunction[Request, ZIO[Any, Response, Response]]](PartialFunction.empty)
-        socketBehavior <- Ref.make[SocketApp[Any]](SocketApp.unit)
+        socketBehavior <- Ref.make[WebSocketApp[Any]](WebSocketApp.unit)
         driver = TestClient(behavior, socketBehavior)
       } yield ZEnvironment[TestClient, Client](driver, ZClient.fromDriver(driver))
     }
diff --git a/zio-http-testkit/src/test/scala/zio/http/SocketContractSpec.scala b/zio-http-testkit/src/test/scala/zio/http/SocketContractSpec.scala
index 04eb0a28fe..880e76405d 100644
--- a/zio-http-testkit/src/test/scala/zio/http/SocketContractSpec.scala
+++ b/zio-http-testkit/src/test/scala/zio/http/SocketContractSpec.scala
@@ -14,7 +14,7 @@ object SocketContractSpec extends ZIOSpecDefault {
   def spec: Spec[Any, Any] =
     suite("SocketOps")(
       contract("Successful Multi-message application") { p =>
-        val socketServer: SocketApp[Any] =
+        val socketServer: WebSocketApp[Any] =
           Handler.webSocket { channel =>
             channel.receiveAll {
               case Read(WebSocketFrame.Text("Hi Server"))          =>
@@ -31,7 +31,7 @@ object SocketContractSpec extends ZIOSpecDefault {
 
         socketServer
       } { _ =>
-        val socketClient: SocketApp[Any] =
+        val socketClient: WebSocketApp[Any] =
           Handler.webSocket { channel =>
             channel.receiveAll {
               case ChannelEvent.Read(WebSocketFrame.Text("Hi Client")) =>
@@ -89,8 +89,8 @@ object SocketContractSpec extends ZIOSpecDefault {
   private def contract(
     name: String,
   )(
-    serverApp: Promise[Throwable, Unit] => SocketApp[Any],
-  )(clientApp: Promise[Throwable, Unit] => SocketApp[Any]) = {
+    serverApp: Promise[Throwable, Unit] => WebSocketApp[Any],
+  )(clientApp: Promise[Throwable, Unit] => WebSocketApp[Any]) = {
     suite(name)(
       test("Live") {
         for {
@@ -123,7 +123,7 @@ object SocketContractSpec extends ZIOSpecDefault {
   }
 
   private def liveServerSetup(
-    serverApp: Promise[Throwable, Unit] => SocketApp[Any],
+    serverApp: Promise[Throwable, Unit] => WebSocketApp[Any],
   ): ZIO[Server, Nothing, (RuntimeFlags, Promise[Throwable, Unit])] =
     ZIO.serviceWithZIO[Server](server =>
       for {
@@ -133,7 +133,7 @@ object SocketContractSpec extends ZIOSpecDefault {
     )
 
   private def testServerSetup(
-    serverApp: Promise[Throwable, Unit] => SocketApp[Any],
+    serverApp: Promise[Throwable, Unit] => WebSocketApp[Any],
   ): ZIO[TestClient, Nothing, (RuntimeFlags, Promise[Throwable, Unit])] =
     for {
       p <- Promise.make[Throwable, Unit]
diff --git a/zio-http-testkit/src/test/scala/zio/http/TestClientSpec.scala b/zio-http-testkit/src/test/scala/zio/http/TestClientSpec.scala
index 180a1a0427..0b0c6a8795 100644
--- a/zio-http-testkit/src/test/scala/zio/http/TestClientSpec.scala
+++ b/zio-http-testkit/src/test/scala/zio/http/TestClientSpec.scala
@@ -66,7 +66,7 @@ object TestClientSpec extends ZIOSpecDefault {
       ),
       suite("socket ops")(
         test("happy path") {
-          val socketClient: SocketApp[Any] =
+          val socketClient: WebSocketApp[Any] =
             Handler.webSocket { channel =>
               channel.receiveAll {
                 case ChannelEvent.Read(WebSocketFrame.Text("Hi Client")) =>
@@ -77,7 +77,7 @@ object TestClientSpec extends ZIOSpecDefault {
               }
             }
 
-          val socketServer: SocketApp[Any] =
+          val socketServer: WebSocketApp[Any] =
             Handler.webSocket { channel =>
               channel.receiveAll {
                 case ChannelEvent.Read(WebSocketFrame.Text("Hi Server")) =>
diff --git a/zio-http/src/main/scala/zio/http/ClientDriver.scala b/zio-http/src/main/scala/zio/http/ClientDriver.scala
index 48e2bdfe72..14831f94ff 100644
--- a/zio-http/src/main/scala/zio/http/ClientDriver.scala
+++ b/zio-http/src/main/scala/zio/http/ClientDriver.scala
@@ -31,7 +31,7 @@ trait ClientDriver {
     onResponse: Promise[Throwable, Response],
     onComplete: Promise[Throwable, ChannelState],
     enableKeepAlive: Boolean,
-    createSocketApp: () => SocketApp[Any],
+    createSocketApp: () => WebSocketApp[Any],
     webSocketConfig: WebSocketConfig,
   )(implicit trace: Trace): ZIO[Scope, Throwable, ChannelInterface]
 
diff --git a/zio-http/src/main/scala/zio/http/Handler.scala b/zio-http/src/main/scala/zio/http/Handler.scala
index ec395eae44..73fc66fa1d 100644
--- a/zio-http/src/main/scala/zio/http/Handler.scala
+++ b/zio-http/src/main/scala/zio/http/Handler.scala
@@ -684,18 +684,6 @@ object Handler {
   def badRequest(message: => String): Handler[Any, Nothing, Any, Response] =
     error(Status.BadRequest, message)
 
-  /**
-   * Constructs a handler from two functions, one that configures web socket and
-   * another that uses a web socket.
-   *
-   * If the config function returns with None, the server configuration is used.
-   */
-  final def customWebSocket[Env](
-    config: Request => ZIO[Env, Throwable, Option[WebSocketConfig]],
-    f: WebSocketChannel => ZIO[Env, Throwable, Any],
-  ): SocketApp[Env] =
-    SocketApp(Handler.fromFunctionZIO(f), Handler.fromFunctionZIO(config))
-
   /**
    * Returns a handler that dies with the specified `Throwable`. This method can
    * be used for terminating an handler because a defect has been detected in
@@ -1073,8 +1061,8 @@ object Handler {
    */
   final def webSocket[Env](
     f: WebSocketChannel => ZIO[Env, Throwable, Any],
-  ): SocketApp[Env] =
-    SocketApp(Handler.fromFunctionZIO(f))
+  ): WebSocketApp[Env] =
+    WebSocketApp(Handler.fromFunctionZIO(f))
 
   final implicit class RequestHandlerSyntax[-R, +Err](val self: RequestHandler[R, Err])
       extends HeaderModifier[RequestHandler[R, Err]] {
diff --git a/zio-http/src/main/scala/zio/http/Response.scala b/zio-http/src/main/scala/zio/http/Response.scala
index daa250205a..1def87457c 100644
--- a/zio-http/src/main/scala/zio/http/Response.scala
+++ b/zio-http/src/main/scala/zio/http/Response.scala
@@ -114,7 +114,7 @@ sealed trait Response extends HeaderOps[Response] { self =>
   final def status(status: Status): Response =
     self.copy(status = status)
 
-  private[zio] final def socketApp: Option[SocketApp[Any]] = self match {
+  private[zio] final def socketApp: Option[WebSocketApp[Any]] = self match {
     case Response.GetApp(app) => Some(app)
     case _                    => None
   }
@@ -140,7 +140,7 @@ object Response {
   }
 
   object GetApp {
-    def unapply(response: Response): Option[SocketApp[Any]] = response match {
+    def unapply(response: Response): Option[WebSocketApp[Any]] = response match {
       case resp: SocketAppResponse => Some(resp.socketApp0)
       case _                       => None
     }
@@ -186,7 +186,7 @@ object Response {
   private[zio] class SocketAppResponse(
     val body: Body,
     val headers: Headers,
-    val socketApp0: SocketApp[Any],
+    val socketApp0: WebSocketApp[Any],
     val status: Status,
   ) extends Response { self =>
 
@@ -360,7 +360,7 @@ object Response {
   /**
    * Creates a new response for the provided socket app
    */
-  def fromSocketApp[R](app: SocketApp[R])(implicit trace: Trace): ZIO[R, Nothing, Response] = {
+  def fromSocketApp[R](app: WebSocketApp[R])(implicit trace: Trace): ZIO[R, Nothing, Response] = {
     ZIO.environment[R].map { env =>
       new SocketAppResponse(
         Body.empty,
diff --git a/zio-http/src/main/scala/zio/http/SocketApp.scala b/zio-http/src/main/scala/zio/http/WebSocketApp.scala
similarity index 68%
rename from zio-http/src/main/scala/zio/http/SocketApp.scala
rename to zio-http/src/main/scala/zio/http/WebSocketApp.scala
index 416ca3eb5c..c914ca7ff0 100644
--- a/zio-http/src/main/scala/zio/http/SocketApp.scala
+++ b/zio-http/src/main/scala/zio/http/WebSocketApp.scala
@@ -2,9 +2,9 @@ package zio.http
 
 import zio._
 
-final case class SocketApp[-R](
+final case class WebSocketApp[-R](
   handler: Handler[R, Throwable, WebSocketChannel, Any],
-  customConfig: Handler[R, Throwable, Request, Option[WebSocketConfig]],
+  customConfig: Option[WebSocketConfig],
 ) { self =>
 
   /**
@@ -31,35 +31,35 @@ final case class SocketApp[-R](
       client2.addHeaders(headers).socket(self)
     }
 
-  def provideEnvironment(r: ZEnvironment[R])(implicit trace: Trace): SocketApp[Any] =
-    SocketApp(handler.provideEnvironment(r), customConfig.provideEnvironment(r))
+  def provideEnvironment(r: ZEnvironment[R])(implicit trace: Trace): WebSocketApp[Any] =
+    WebSocketApp(handler.provideEnvironment(r), customConfig)
 
   def provideLayer[R0](layer: ZLayer[R0, Throwable, R])(implicit
     trace: Trace,
-  ): SocketApp[R0] =
-    SocketApp(handler.provideLayer(layer), customConfig.provideLayer(layer))
+  ): WebSocketApp[R0] =
+    WebSocketApp(handler.provideLayer(layer), customConfig)
 
   def provideSomeEnvironment[R1](f: ZEnvironment[R1] => ZEnvironment[R])(implicit
     trace: Trace,
-  ): SocketApp[R1] =
-    SocketApp(handler.provideSomeEnvironment(f), customConfig.provideSomeEnvironment(f))
+  ): WebSocketApp[R1] =
+    WebSocketApp(handler.provideSomeEnvironment(f), customConfig)
 
   def provideSomeLayer[R0, R1: Tag](
     layer: ZLayer[R0, Throwable, R1],
-  )(implicit ev: R0 with R1 <:< R, trace: Trace): SocketApp[R0] =
-    SocketApp(handler.provideSomeLayer(layer), customConfig.provideSomeLayer(layer))
+  )(implicit ev: R0 with R1 <:< R, trace: Trace): WebSocketApp[R0] =
+    WebSocketApp(handler.provideSomeLayer(layer), customConfig)
 
   def tapErrorCauseZIO[R1 <: R](
     f: Cause[Throwable] => ZIO[R1, Throwable, Any],
-  )(implicit trace: Trace): SocketApp[R1] =
-    SocketApp(handler.tapErrorCauseZIO(f), customConfig.tapErrorCauseZIO(f))
+  )(implicit trace: Trace): WebSocketApp[R1] =
+    WebSocketApp(handler.tapErrorCauseZIO(f), customConfig)
 
   /**
    * Returns a Handler that effectfully peeks at the failure of this SocketApp.
    */
   def tapErrorZIO[R1 <: R](
     f: Throwable => ZIO[R1, Throwable, Any],
-  )(implicit trace: Trace): SocketApp[R1] =
+  )(implicit trace: Trace): WebSocketApp[R1] =
     self.tapErrorCauseZIO(cause => cause.failureOption.fold[ZIO[R1, Throwable, Any]](ZIO.unit)(f))
 
   /**
@@ -74,11 +74,14 @@ final case class SocketApp[-R](
 
   def toHttpAppWS(implicit trace: Trace): HttpApp[R] =
     Handler.fromZIO(self.toResponse).toHttpApp
+
+  def withConfig(config: WebSocketConfig): WebSocketApp[R] =
+    copy(customConfig = Some(config))
 }
 
-object SocketApp {
-  def apply[R](handler: Handler[R, Throwable, WebSocketChannel, Any]): SocketApp[R] =
-    SocketApp(handler, Handler.succeed(None))
+object WebSocketApp {
+  def apply[R](handler: Handler[R, Throwable, WebSocketChannel, Any]): WebSocketApp[R] =
+    WebSocketApp(handler, None)
 
-  val unit: SocketApp[Any] = SocketApp(Handler.unit)
+  val unit: WebSocketApp[Any] = WebSocketApp(Handler.unit)
 }
diff --git a/zio-http/src/main/scala/zio/http/ZClient.scala b/zio-http/src/main/scala/zio/http/ZClient.scala
index 82284401ce..4d84e85a0e 100644
--- a/zio-http/src/main/scala/zio/http/ZClient.scala
+++ b/zio-http/src/main/scala/zio/http/ZClient.scala
@@ -205,7 +205,7 @@ final case class ZClient[-Env, -In, +Err, +Out](
   def scheme(scheme: Scheme): ZClient[Env, In, Err, Out] =
     copy(url = url.scheme(scheme))
 
-  def socket[Env1 <: Env](app: SocketApp[Env1])(implicit trace: Trace): ZIO[Env1 & Scope, Err, Out] =
+  def socket[Env1 <: Env](app: WebSocketApp[Env1])(implicit trace: Trace): ZIO[Env1 & Scope, Err, Out] =
     driver
       .socket(
         Version.Default,
@@ -291,7 +291,7 @@ object ZClient {
   def request(request: Request): ZIO[Client & Scope, Throwable, Response] =
     ZIO.serviceWithZIO[Client](c => c(request))
 
-  def socket[R](socketApp: SocketApp[R]): ZIO[R with Client & Scope, Throwable, Response] =
+  def socket[R](socketApp: WebSocketApp[R]): ZIO[R with Client & Scope, Throwable, Response] =
     ZIO.serviceWithZIO[Client](c => c.socket(socketApp))
 
   trait BodyDecoder[-Env, +Err, +Out] { self =>
@@ -387,7 +387,7 @@ object ZClient {
           version: Version,
           url: URL,
           headers: Headers,
-          app: SocketApp[Env1],
+          app: WebSocketApp[Env1],
         )(implicit trace: Trace): ZIO[Env1 & Scope, Throwable, Response] =
           self0
             .socket(
@@ -420,7 +420,7 @@ object ZClient {
           version: Version,
           url: URL,
           headers: Headers,
-          app: SocketApp[Env1],
+          app: WebSocketApp[Env1],
         )(implicit trace: Trace): ZIO[Env1 & Scope, Err2, Response] =
           self
             .socket(
@@ -450,7 +450,7 @@ object ZClient {
           version: Version,
           url: URL,
           headers: Headers,
-          app: SocketApp[Env1],
+          app: WebSocketApp[Env1],
         )(implicit trace: Trace): ZIO[Env1 & Scope, Err2, Response] =
           self
             .socket(
@@ -490,7 +490,7 @@ object ZClient {
           version: Version,
           url: URL,
           headers: Headers,
-          app: SocketApp[Env2],
+          app: WebSocketApp[Env2],
         )(implicit trace: Trace): ZIO[Env2 & Scope, Err1, Response] =
           self
             .socket(
@@ -506,7 +506,7 @@ object ZClient {
       version: Version,
       url: URL,
       headers: Headers,
-      app: SocketApp[Env1],
+      app: WebSocketApp[Env1],
     )(implicit trace: Trace): ZIO[Env1 & Scope, Err, Response]
 
     def widenError[E1](implicit ev: Err <:< E1): Driver[Env, E1] = self.asInstanceOf[Driver[Env, E1]]
@@ -633,14 +633,14 @@ object ZClient {
       val request = Request(version, method, url, headers, body, None)
       val cfg     = sslConfig.fold(config)(config.ssl)
 
-      requestAsync(request, cfg, () => SocketApp.unit, None)
+      requestAsync(request, cfg, () => WebSocketApp.unit, None)
     }
 
     def socket[Env1](
       version: Version,
       url: URL,
       headers: Headers,
-      app: SocketApp[Env1],
+      app: WebSocketApp[Env1],
     )(implicit trace: Trace): ZIO[Env1 & Scope, Throwable, Response] =
       for {
         env <- ZIO.environment[Env1]
@@ -668,7 +668,7 @@ object ZClient {
     private def requestAsync(
       request: Request,
       clientConfig: Config,
-      createSocketApp: () => SocketApp[Any],
+      createSocketApp: () => WebSocketApp[Any],
       outerScope: Option[Scope],
     )(implicit
       trace: Trace,
diff --git a/zio-http/src/main/scala/zio/http/ZClientAspect.scala b/zio-http/src/main/scala/zio/http/ZClientAspect.scala
index 054a38e61b..c6f48a6be5 100644
--- a/zio-http/src/main/scala/zio/http/ZClientAspect.scala
+++ b/zio-http/src/main/scala/zio/http/ZClientAspect.scala
@@ -155,8 +155,8 @@ object ZClientAspect {
               .flatMap(_._2)
               .unsandbox
 
-          override def socket[Env1 <: Env](version: Version, url: URL, headers: Headers, app: SocketApp[Env1])(implicit
-            trace: Trace,
+          override def socket[Env1 <: Env](version: Version, url: URL, headers: Headers, app: WebSocketApp[Env1])(
+            implicit trace: Trace,
           ): ZIO[Env1 with Scope, Err, Response] =
             client.driver.socket(version, url, headers, app)
         }
@@ -299,8 +299,8 @@ object ZClientAspect {
               .unsandbox
           }
 
-          override def socket[Env1 <: Env](version: Version, url: URL, headers: Headers, app: SocketApp[Env1])(implicit
-            trace: Trace,
+          override def socket[Env1 <: Env](version: Version, url: URL, headers: Headers, app: WebSocketApp[Env1])(
+            implicit trace: Trace,
           ): ZIO[Env1 with Scope, Err, Response] =
             client.driver.socket(version, url, headers, app)
         }
diff --git a/zio-http/src/main/scala/zio/http/netty/client/NettyClientDriver.scala b/zio-http/src/main/scala/zio/http/netty/client/NettyClientDriver.scala
index b3c3b187b1..0ee2a68316 100644
--- a/zio-http/src/main/scala/zio/http/netty/client/NettyClientDriver.scala
+++ b/zio-http/src/main/scala/zio/http/netty/client/NettyClientDriver.scala
@@ -49,7 +49,7 @@ final case class NettyClientDriver private (
     onResponse: Promise[Throwable, Response],
     onComplete: Promise[Throwable, ChannelState],
     enableKeepAlive: Boolean,
-    createSocketApp: () => SocketApp[Any],
+    createSocketApp: () => WebSocketApp[Any],
     webSocketConfig: WebSocketConfig,
   )(implicit trace: Trace): ZIO[Scope, Throwable, ChannelInterface] = {
     NettyRequestEncoder.encode(req).flatMap { jReq =>
@@ -68,8 +68,7 @@ final case class NettyClientDriver private (
         nettyChannel     = NettyChannel.make[JWebSocketFrame](channel)
         webSocketChannel = WebSocketChannel.make(nettyChannel, queue)
         app              = createSocketApp()
-        customConfig <- app.customConfig.runZIO(req)
-        _            <- app.handler.runZIO(webSocketChannel).ignoreLogged.interruptible.forkScoped
+        _ <- app.handler.runZIO(webSocketChannel).ignoreLogged.interruptible.forkScoped
       } yield {
         val pipeline                              = channel.pipeline()
         val toRemove: mutable.Set[ChannelHandler] = new mutable.HashSet[ChannelHandler]()
@@ -86,7 +85,7 @@ final case class NettyClientDriver private (
 
           val headers = Conversions.headersToNetty(req.headers)
           val config  = NettySocketProtocol
-            .clientBuilder(customConfig.getOrElse(webSocketConfig))
+            .clientBuilder(app.customConfig.getOrElse(webSocketConfig))
             .customHeaders(headers)
             .webSocketUri(req.url.encode)
             .build()
diff --git a/zio-http/src/main/scala/zio/http/netty/server/ServerInboundHandler.scala b/zio-http/src/main/scala/zio/http/netty/server/ServerInboundHandler.scala
index 757387c18f..dc685a855b 100644
--- a/zio-http/src/main/scala/zio/http/netty/server/ServerInboundHandler.scala
+++ b/zio-http/src/main/scala/zio/http/netty/server/ServerInboundHandler.scala
@@ -292,10 +292,7 @@ private[zio] final case class ServerInboundHandler(
     jReq match {
       case jReq: FullHttpRequest =>
         val queue = runtime.runtime(ctx).unsafe.run(Queue.unbounded[WebSocketChannelEvent]).getOrThrowFiberFailure()
-        val webSocketApp = app.getOrElse(SocketApp.unit)
-        val request      = makeZioRequest(ctx, jReq)
-        val customConfig =
-          runtime.runtime(ctx).unsafe.run { webSocketApp.customConfig.runZIO(request) }.getOrThrowFiberFailure()
+        val webSocketApp = app.getOrElse(WebSocketApp.unit)
         runtime.runtime(ctx).unsafe.run {
           val nettyChannel     = NettyChannel.make[JWebSocketFrame](ctx.channel())
           val webSocketChannel = WebSocketChannel.make(nettyChannel, queue)
@@ -306,7 +303,7 @@ private[zio] final case class ServerInboundHandler(
           .pipeline()
           .addLast(
             new WebSocketServerProtocolHandler(
-              NettySocketProtocol.serverBuilder(customConfig.getOrElse(config.webSocketConfig)).build(),
+              NettySocketProtocol.serverBuilder(webSocketApp.customConfig.getOrElse(config.webSocketConfig)).build(),
             ),
           )
           .addLast(Names.WebSocketHandler, new WebSocketAppHandler(runtime, queue, None))
diff --git a/zio-http/src/test/scala/zio/http/WebSocketSpec.scala b/zio-http/src/test/scala/zio/http/WebSocketSpec.scala
index aa1fd74638..1f28efd705 100644
--- a/zio-http/src/test/scala/zio/http/WebSocketSpec.scala
+++ b/zio-http/src/test/scala/zio/http/WebSocketSpec.scala
@@ -143,7 +143,7 @@ object WebSocketSpec extends HttpRunnableSpec {
       val app   =
         Handler.webSocket(channel => channel.send(ChannelEvent.Read(WebSocketFrame.text("BAR")))).toHttpAppWS.deployWS
       val codes = ZIO
-        .foreach(1 to 1024)(_ => app.runZIO(SocketApp.unit).map(_.status))
+        .foreach(1 to 1024)(_ => app.runZIO(WebSocketApp.unit).map(_.status))
         .map(_.count(_ == Status.SwitchingProtocols))
 
       assertZIO(codes)(equalTo(1024))
diff --git a/zio-http/src/test/scala/zio/http/internal/HttpRunnableSpec.scala b/zio-http/src/test/scala/zio/http/internal/HttpRunnableSpec.scala
index c4640f854b..a1133b934f 100644
--- a/zio-http/src/test/scala/zio/http/internal/HttpRunnableSpec.scala
+++ b/zio-http/src/test/scala/zio/http/internal/HttpRunnableSpec.scala
@@ -87,13 +87,13 @@ abstract class HttpRunnableSpec extends ZIOSpecDefault { self =>
       } yield response
 
     def deployWS
-      : Handler[R with Client with DynamicServer with Scope, Throwable, SocketApp[Client with Scope], Response] =
+      : Handler[R with Client with DynamicServer with Scope, Throwable, WebSocketApp[Client with Scope], Response] =
       for {
         id       <- Handler.fromZIO(DynamicServer.deploy[R](app))
         rawUrl   <- Handler.fromZIO(DynamicServer.wsURL)
         url      <- Handler.fromEither(URL.decode(rawUrl)).orDie
         client   <- Handler.fromZIO(ZIO.service[Client])
-        response <- Handler.fromFunctionZIO[SocketApp[Client with Scope]] { app =>
+        response <- Handler.fromFunctionZIO[WebSocketApp[Client with Scope]] { app =>
           ZIO.scoped[Client with Scope](
             client
               .url(url)

From dc2f3f92b39abe66ce4dabb3660acdc0addf04f4 Mon Sep 17 00:00:00 2001
From: Daniel Vigovszky <daniel.vigovszky@gmail.com>
Date: Fri, 28 Jul 2023 18:47:21 +0100
Subject: [PATCH 4/4] Doc fix

---
 docs/examples/advanced/websocket-server.md | 2 +-
 docs/examples/basic/websocket.md           | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/docs/examples/advanced/websocket-server.md b/docs/examples/advanced/websocket-server.md
index c72d0a311b..2534d5e9de 100644
--- a/docs/examples/advanced/websocket-server.md
+++ b/docs/examples/advanced/websocket-server.md
@@ -13,7 +13,7 @@ import zio.http.codec.PathCodec.string
 
 object WebSocketAdvanced extends ZIOAppDefault {
 
-  val socketApp: SocketApp[Any] =
+  val socketApp: WebSocketApp[Any] =
     Handler.webSocket { channel =>
       channel.receiveAll {
         case Read(WebSocketFrame.Text("end"))                =>
diff --git a/docs/examples/basic/websocket.md b/docs/examples/basic/websocket.md
index 21f6cdae62..045f3e4f6b 100644
--- a/docs/examples/basic/websocket.md
+++ b/docs/examples/basic/websocket.md
@@ -12,7 +12,7 @@ import zio.http._
 import zio.http.codec.PathCodec.string
 
 object WebSocketEcho extends ZIOAppDefault {
-  private val socketApp: SocketApp[Any] =
+  private val socketApp: WebSocketApp[Any] =
     Handler.webSocket { channel =>
       channel.receiveAll {
         case Read(WebSocketFrame.Text("FOO")) =>