diff --git a/docs/server/transport.md b/docs/server/transport.md index 5a33b68..d7bc72e 100644 --- a/docs/server/transport.md +++ b/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. diff --git a/server-streaming/server-ox/src/main/scala/chimp/server/ox/OxServerHttpTransport.scala b/server-streaming/server-ox/src/main/scala/chimp/server/ox/OxServerHttpTransport.scala index c997361..86db8d6 100644 --- a/server-streaming/server-ox/src/main/scala/chimp/server/ox/OxServerHttpTransport.scala +++ b/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)) diff --git a/server-streaming/server-ox/src/test/scala/chimp/server/ox/OxServerKeepAliveSpec.scala b/server-streaming/server-ox/src/test/scala/chimp/server/ox/OxServerKeepAliveSpec.scala new file mode 100644 index 0000000..864926f --- /dev/null +++ b/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 diff --git a/server-streaming/server-pekko/src/main/scala/chimp/server/pekko/PekkoServerHttpTransport.scala b/server-streaming/server-pekko/src/main/scala/chimp/server/pekko/PekkoServerHttpTransport.scala index 99f42d2..39ebf9c 100644 --- a/server-streaming/server-pekko/src/main/scala/chimp/server/pekko/PekkoServerHttpTransport.scala +++ b/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,13 +15,19 @@ 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): @@ -29,6 +35,8 @@ final class 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))) diff --git a/server-streaming/server-pekko/src/test/scala/chimp/server/pekko/PekkoServerKeepAliveSpec.scala b/server-streaming/server-pekko/src/test/scala/chimp/server/pekko/PekkoServerKeepAliveSpec.scala new file mode 100644 index 0000000..d435fb7 --- /dev/null +++ b/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 diff --git a/server-streaming/server-zio/src/main/scala/chimp/server/zio/ZioServerHttpTransport.scala b/server-streaming/server-zio/src/main/scala/chimp/server/zio/ZioServerHttpTransport.scala index 9b2c4d6..e4bf43d 100644 --- a/server-streaming/server-zio/src/main/scala/chimp/server/zio/ZioServerHttpTransport.scala +++ b/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)) + ) } } diff --git a/server-streaming/server-zio/src/test/scala/chimp/server/zio/ZioServerKeepAliveSpec.scala b/server-streaming/server-zio/src/test/scala/chimp/server/zio/ZioServerKeepAliveSpec.scala new file mode 100644 index 0000000..04a8bca --- /dev/null +++ b/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