WIP: feat: add opt-in SSE keep-alive to streaming HTTP transports

docs/server/transport.md

@@ -45,6 +45,18 @@ The streaming transports have concrete implementations per effect system, in sep
 | Ox (direct style) | `OxServerHttpTransport` | `OxServerStdioTransport` |
 | Pekko | `PekkoServerHttpTransport` | `PekkoServerStdioTransport` |
 
+### Keep-alive
+
+Long-running tool calls hold the Server-Sent Event response stream open. Idle connections can be closed by proxies in between. Set `keepAlive` on a streaming HTTP transport to emit a data-less `ping` event at a fixed interval while the stream is open. The events carry no data and are ignored by MCP clients:
+
+```scala mdoc:compile-only
+import chimp.server.ox.OxServerHttpTransport
+
+import scala.concurrent.duration.*
+
+val transport = OxServerHttpTransport(List("mcp"), keepAlive = Some(15.seconds))
+```
+
 ## Backends
 
 - **HTTP** transports produce a Tapir `ServerEndpoint` that you run on any Tapir server interpreter. The streaming HTTP transport additionally requires an interpreter with streaming capability.

server-streaming/server-ox/src/main/scala/chimp/server/ox/OxServerHttpTransport.scala

@@ -13,7 +13,14 @@ import sttp.shared.Identity
 import sttp.tapir.StreamBodyIO
 import sttp.tapir.server.netty.sync.{serverSentEventsBody, OxStreams}
 
-final class OxServerHttpTransport(path: List[String]) extends ServerStreamingHttpTransport[Identity, OxStreams](path):
+import scala.concurrent.duration.FiniteDuration
+
+/** @param keepAlive
+  *   If set, a data-less `ping` Server-Sent Event is emitted on the response stream at this interval, to keep idle connections open through
+  *   proxies. The events carry no data and are ignored by MCP clients.
+  */
+final class OxServerHttpTransport(path: List[String], keepAlive: Option[FiniteDuration] = None)
+    extends ServerStreamingHttpTransport[Identity, OxStreams](path):
   val streams: OxStreams = OxStreams
 
   type EventStream = Flow[ServerSentEvent]
@@ -22,8 +29,10 @@ final class OxServerHttpTransport(path: List[String]) extends ServerStreamingHtt
 
   val emptyStream: EventStream = Flow.empty
 
+  private val pingEvent = ServerSentEvent(eventType = Some("ping"))
+
   def eventStream(handle: OutboundSink[Identity] => Option[Json]): Flow[ServerSentEvent] =
-    Flow.usingEmit: emit =>
+    val messages: Flow[ServerSentEvent] = Flow.usingEmit: emit =>
       supervised:
         val outbound = Channel.buffered[Json](64)
         val sink = new OutboundSink[Identity]:
@@ -32,3 +41,4 @@ final class OxServerHttpTransport(path: List[String]) extends ServerStreamingHtt
           try handle(sink).foreach(outbound.send)
           finally outbound.done()
         outbound.foreach(json => emit(ServerSentEvent(data = Some(json.noSpaces))))
+    keepAlive.fold(messages)(interval => messages.merge[ServerSentEvent](Flow.tick(interval, pingEvent), propagateDoneLeft = true))

server-streaming/server-ox/src/test/scala/chimp/server/ox/OxServerKeepAliveSpec.scala

@@ -0,0 +1,24 @@
+package chimp.server.ox
+
+import chimp.server.OutboundSink
+import io.circe.Json
+import org.scalatest.flatspec.AnyFlatSpec
+import org.scalatest.matchers.should.Matchers
+import ox.supervised
+import sttp.shared.Identity
+
+import scala.concurrent.duration.*
+
+class OxServerKeepAliveSpec extends AnyFlatSpec with Matchers:
+
+  it should "emit data-less ping events while a tool call is in flight" in:
+    val transport = OxServerHttpTransport(List("mcp"), keepAlive = Some(50.millis))
+    val handle: OutboundSink[Identity] => Option[Json] = _ =>
+      Thread.sleep(300)
+      None
+
+    val events = supervised(transport.eventStream(handle).take(1).runToList())
+
+    events should have size 1
+    events.head.eventType shouldBe Some("ping")
+    events.head.data shouldBe None

server-streaming/server-pekko/src/main/scala/chimp/server/pekko/PekkoServerHttpTransport.scala

@@ -6,7 +6,7 @@ import chimp.server.transport.ServerStreamingHttpTransport
 import io.circe.Json
 import io.circe.syntax.*
 import org.apache.pekko.NotUsed
-import org.apache.pekko.stream.scaladsl.Source
+import org.apache.pekko.stream.scaladsl.{Flow, Source}
 import org.apache.pekko.stream.{Materializer, OverflowStrategy}
 import org.slf4j.LoggerFactory
 import sttp.capabilities.pekko.PekkoStreams
@@ -15,20 +15,28 @@ import sttp.tapir.server.pekkohttp.PekkoServerSentEvents
 import sttp.tapir.{streamTextBody, CodecFormat, StreamBodyIO}
 
 import java.nio.charset.StandardCharsets
+import scala.concurrent.duration.FiniteDuration
 import scala.concurrent.{ExecutionContext, Future}
 import scala.util.{Failure, Success, Try}
 
+/** @param keepAlive
+  *   If set, a data-less `ping` Server-Sent Event is emitted on the response stream at this interval, to keep idle connections open through
+  *   proxies. The events carry no data and are ignored by MCP clients.
+  */
 final class PekkoServerHttpTransport(
     path: List[String],
     outboundBufferSize: Int = PekkoOutbound.defaultBufferSize,
-    maxConcurrentSends: Int = PekkoOutbound.defaultMaxConcurrentSends
+    maxConcurrentSends: Int = PekkoOutbound.defaultMaxConcurrentSends,
+    keepAlive: Option[FiniteDuration] = None
 )(using mat: Materializer)
     extends ServerStreamingHttpTransport[Future, PekkoStreams](path):
 
   private val log = LoggerFactory.getLogger(classOf[PekkoServerHttpTransport])
 
   private given ExecutionContext = mat.executionContext
 
+  private val pingEvent = ServerSentEvent(eventType = Some("ping"))
+
   val streams: PekkoStreams = PekkoStreams
 
   type EventStream = Source[ServerSentEvent, Any]
@@ -57,3 +65,4 @@ final class PekkoServerHttpTransport(
                 queue.complete()
           NotUsed
         .map(json => ServerSentEvent(data = Some(json.noSpaces)))
+        .via(keepAlive.fold(Flow[ServerSentEvent])(interval => Flow[ServerSentEvent].keepAlive(interval, () => pingEvent)))

server-streaming/server-pekko/src/test/scala/chimp/server/pekko/PekkoServerKeepAliveSpec.scala

@@ -0,0 +1,34 @@
+package chimp.server.pekko
+
+import chimp.server.OutboundSink
+import io.circe.Json
+import org.apache.pekko.actor.ActorSystem
+import org.apache.pekko.stream.Materializer
+import org.apache.pekko.stream.scaladsl.Sink
+import org.scalatest.BeforeAndAfterAll
+import org.scalatest.flatspec.AnyFlatSpec
+import org.scalatest.matchers.should.Matchers
+import sttp.model.sse.ServerSentEvent
+
+import scala.concurrent.duration.*
+import scala.concurrent.{Await, ExecutionContext, Future}
+
+class PekkoServerKeepAliveSpec extends AnyFlatSpec with Matchers with BeforeAndAfterAll:
+
+  private given system: ActorSystem = ActorSystem("chimp-server-pekko-keepalive-test")
+  private given ExecutionContext = system.dispatcher
+  private given Materializer = Materializer.matFromSystem
+
+  override def afterAll(): Unit =
+    val _ = Await.result(system.terminate(), 30.seconds)
+
+  it should "emit data-less ping events while a tool call is in flight" in:
+    val transport = PekkoServerHttpTransport(List("mcp"), keepAlive = Some(50.millis))
+    val handle: OutboundSink[Future] => Future[Option[Json]] = _ => Future { Thread.sleep(300); None }
+
+    val events: Seq[ServerSentEvent] =
+      Await.result(transport.eventStream(handle).flatMap(_.take(1).runWith(Sink.seq)), 5.seconds)
+
+    events should have size 1
+    events.head.eventType shouldBe Some("ping")
+    events.head.data shouldBe None

server-streaming/server-zio/src/main/scala/chimp/server/zio/ZioServerHttpTransport.scala

@@ -10,11 +10,17 @@ import sttp.model.sse.ServerSentEvent
 import sttp.tapir.*
 import sttp.tapir.ztapir.ZioServerSentEvents
 import zio.stream.{Stream, ZStream}
-import zio.{Queue, Task, ZIO}
+import zio.{Duration, Queue, Schedule, Task, ZIO}
 
 import java.nio.charset.StandardCharsets
-
-final class ZioServerHttpTransport(path: List[String]) extends ServerStreamingHttpTransport[Task, ZioStreams](path):
+import scala.concurrent.duration.FiniteDuration
+
+/** @param keepAlive
+  *   If set, a data-less `ping` Server-Sent Event is emitted on the response stream at this interval, to keep idle connections open through
+  *   proxies. The events carry no data and are ignored by MCP clients.
+  */
+final class ZioServerHttpTransport(path: List[String], keepAlive: Option[FiniteDuration] = None)
+    extends ServerStreamingHttpTransport[Task, ZioStreams](path):
   val streams: ZioStreams = ZioStreams
 
   type EventStream = Stream[Throwable, ServerSentEvent]
@@ -25,6 +31,8 @@ final class ZioServerHttpTransport(path: List[String]) extends ServerStreamingHt
 
   val emptyStream: EventStream = ZStream.empty
 
+  private val pingEvent = ServerSentEvent(eventType = Some("ping"))
+
   def eventStream(handle: OutboundSink[Task] => Task[Option[Json]]): Task[EventStream] =
     ZIO.succeed {
       ZStream.unwrapScoped {
@@ -38,7 +46,13 @@ final class ZioServerHttpTransport(path: List[String]) extends ServerStreamingHt
             .ensuring(queue.offer(Outbound.Close))
             .catchAllCause(_ => ZIO.unit)
             .forkScoped
-        yield ZStream.fromQueue(queue).collectWhile { case Outbound.Message(json) => ServerSentEvent(data = Some(json.noSpaces)) }
+        yield
+          val messages = ZStream.fromQueue(queue).collectWhile { case Outbound.Message(json) =>
+            ServerSentEvent(data = Some(json.noSpaces))
+          }
+          keepAlive.fold(messages)(interval =>
+            messages.mergeHaltLeft(ZStream.fromSchedule(Schedule.spaced(Duration.fromScala(interval))).as(pingEvent))
+          )
       }
     }
 

server-streaming/server-zio/src/test/scala/chimp/server/zio/ZioServerKeepAliveSpec.scala

@@ -0,0 +1,26 @@
+package chimp.server.zio
+
+import chimp.server.OutboundSink
+import io.circe.Json
+import org.scalatest.flatspec.AnyFlatSpec
+import org.scalatest.matchers.should.Matchers
+import zio.{Duration, Runtime, Task, Unsafe, ZIO}
+
+import scala.concurrent.duration.*
+
+class ZioServerKeepAliveSpec extends AnyFlatSpec with Matchers:
+
+  private val runtime: Runtime[Any] = Runtime.default
+
+  private def run[A](t: Task[A]): A =
+    Unsafe.unsafe(implicit u => runtime.unsafe.run(t).getOrThrowFiberFailure())
+
+  it should "emit data-less ping events while a tool call is in flight" in:
+    val transport = ZioServerHttpTransport(List("mcp"), keepAlive = Some(50.millis))
+    val handle: OutboundSink[Task] => Task[Option[Json]] = _ => ZIO.sleep(Duration.fromMillis(300)).as(None)
+
+    val events = run(transport.eventStream(handle).flatMap(_.take(1).runCollect))
+
+    events.size shouldBe 1
+    events.head.eventType shouldBe Some("ping")
+    events.head.data shouldBe None