Experimenting with the state monad to have a context in https4

StateTest.scala

import cats.data._
import cats.effect.ExitCode
import cats.implicits._
import cats.~>
import org.http4s.dsl.Http4sDsl2
import org.http4s.server.blaze.BlazeServerBuilder
import org.http4s.{ Http, Request, Response }
import zio._
import zio.clock.Clock
import zio.interop.catz._

object StateTest extends App {

  type StateEffect[A] = StateT[Task, Context, A]
  type StateRoutes    = Http[OptionT[StateEffect, *], Task]
  type StateApp       = Http[StateEffect, Task]

  val dsl: Http4sDsl2[StateEffect, Task] = new Http4sDsl2[StateEffect, Task] {
    override def liftG: Task ~> StateEffect = StateT.liftK
  }
  import dsl._

  case class Context(flashMessage: Option[String], crsfToken: Option[String])

  //controller
  def getS: StateEffect[Response[Task]] =
    for {
      ctx      <- StateT.get[Task, Context]
      response <- Ok(s"flashMessage: ${ctx.flashMessage} | crsfToken: ${ctx.crsfToken}")
    } yield response

  def routes(pf: PartialFunction[Request[Task], StateEffect[Response[Task]]]): StateRoutes =
    Kleisli(req => OptionT(pf.lift(req).sequence))

  val route1: StateRoutes = routes {
    case GET -> Root / "hello" => getS
  }

  val route2: StateRoutes = routes {
    case GET -> Root / "welcome" => getS
  }

  val route3: StateRoutes = routes {
    case GET -> Root / "goodbye" => getS
  }

  val routes: StateRoutes = flashMessageMiddleWare(route1) <+> crsfTokenMiddleWare(route2) <+> flashMessageMiddleWare(
          crsfTokenMiddleWare(route3)
        )

  def flashMessageMiddleWare(service: StateRoutes): StateRoutes = Kleisli { req: Request[Task] =>
    // Fake impl..read something from the req BEFORE servicing it
    req.cookies.headOption match {
      case None => service(req)
      case Some(cookie) =>
        for {
          _    <- OptionT.liftF(StateT.modify[Task, Context](ctx => ctx.copy(flashMessage = cookie.content.some)))
          resp <- service(req)
        } yield resp
    }
  }

  def crsfTokenMiddleWare(service: StateRoutes): StateRoutes = Kleisli { req: Request[Task] =>
    // Fake impl..read something from the req BEFORE servicing it
    req.cookies.headOption match {
      case None => service(req)
      case Some(cookie) =>
        for {
          _    <- OptionT.liftF(StateT.modify[Task, Context](ctx => ctx.copy(crsfToken = cookie.name.some)))
          resp <- service(req)
        } yield resp
    }
  }

  def orNotFound(stateRoutes: StateRoutes): StateApp = Kleisli(a => stateRoutes.run(a).getOrElse(Response.notFound))

  override def run(args: List[String]): ZIO[zio.ZEnv, Nothing, Int] =
    runServer(orNotFound(routes), "localhost", 1234).fold(_ => 1, _ => 0)

  private def runServer(httpApp: StateApp, host: String, port: Int): Task[Unit] =
    ZIO
      .runtime[Clock]
      .map { implicit rts =>
        import zio.interop.catz.implicits._
        BlazeServerBuilder[Task].bindHttp(port, host)
      }
      .provide(Clock.Live)
      .flatMap {
        _.withHttpApp(httpApp.mapK(λ[StateEffect ~> Task](_.runA(Context(none, none))))).serve
          .compile[Task, Task, ExitCode]
          .drain
      }
}