diff --git a/client-streaming/client-ox/src/main/scala/chimp/client/transport/ox/OxServerNotifications.scala b/client-streaming/client-ox/src/main/scala/chimp/client/transport/ox/OxServerNotifications.scala new file mode 100644 index 0000000..535c3ab --- /dev/null +++ b/client-streaming/client-ox/src/main/scala/chimp/client/transport/ox/OxServerNotifications.scala @@ -0,0 +1,22 @@ +package chimp.client.transport.ox + +import chimp.client.BidirectionalMcpClient +import chimp.client.notifications.{ServerNotification, ServerNotificationListener} +import ox.channels.Channel +import ox.discard +import ox.flow.Flow +import sttp.shared.Identity + +extension (client: BidirectionalMcpClient[Identity]) + /** A [[Flow]] of notifications pushed by the server. Each time the flow is run it registers a listener with the client, emits every + * notification the server sends, and removes the listener when the flow finishes. The backing channel is unbounded, so delivery never + * blocks the transport; run the flow with a bound (for example `.take`) or in its own scope if you do not want it to run for the whole + * lifetime of the client. + */ + def serverNotifications: Flow[ServerNotification] = + Flow.usingEmit: emit => + val channel = Channel.unlimited[ServerNotification] + val listener: ServerNotificationListener[Identity] = n => channel.sendOrClosed(n).discard + client.onServerNotification(listener) + try channel.foreach(n => emit(n)) + finally client.removeServerNotification(listener) diff --git a/client-streaming/client-ox/src/test/scala/chimp/client/transport/ox/OxServerNotificationsSpec.scala b/client-streaming/client-ox/src/test/scala/chimp/client/transport/ox/OxServerNotificationsSpec.scala new file mode 100644 index 0000000..e5f9890 --- /dev/null +++ b/client-streaming/client-ox/src/test/scala/chimp/client/transport/ox/OxServerNotificationsSpec.scala @@ -0,0 +1,34 @@ +package chimp.client.transport.ox + +import chimp.client.FakeNotificationClient +import chimp.client.notifications.ServerNotification +import chimp.protocol.{ProgressParams, ProgressToken} +import org.scalatest.flatspec.AnyFlatSpec +import org.scalatest.matchers.should.Matchers +import ox.* +import sttp.monad.{IdentityMonad, MonadError} +import sttp.shared.Identity + +class OxServerNotificationsSpec extends AnyFlatSpec with Matchers: + + private given MonadError[Identity] = IdentityMonad + + private def progress(p: Double): ServerNotification = + ServerNotification.Progress(ProgressParams(progressToken = ProgressToken("t"), progress = p)) + + it should "emit notifications pushed by the server and remove the listener when done" in: + val client = FakeNotificationClient[Identity]() + val n1 = progress(0.1) + val n2 = progress(0.2) + val n3 = progress(0.3) + + val result = supervised: + val collecting = fork(client.serverNotifications.take(3).runToList()) + while client.listenerCount == 0 do Thread.sleep(5) + client.emit(n1) + client.emit(n2) + client.emit(n3) + collecting.join() + + result shouldBe List(n1, n2, n3) + client.listenerCount shouldBe 0 diff --git a/client-streaming/client-pekko/src/main/scala/chimp/client/transport/pekko/PekkoServerNotifications.scala b/client-streaming/client-pekko/src/main/scala/chimp/client/transport/pekko/PekkoServerNotifications.scala new file mode 100644 index 0000000..bbde0ed --- /dev/null +++ b/client-streaming/client-pekko/src/main/scala/chimp/client/transport/pekko/PekkoServerNotifications.scala @@ -0,0 +1,26 @@ +package chimp.client.transport.pekko + +import chimp.client.BidirectionalMcpClient +import chimp.client.notifications.{ServerNotification, ServerNotificationListener} +import org.apache.pekko.NotUsed +import org.apache.pekko.stream.Materializer +import org.apache.pekko.stream.scaladsl.Source + +import scala.concurrent.Future + +extension (client: BidirectionalMcpClient[Future]) + /** A [[Source]] of notifications pushed by the server. When the source is materialized it registers a listener with the client, emits + * every notification the server sends in order, and removes the listener when the stream terminates. The buffer holds up to `bufferSize` + * notifications; when it is full the newest notifications are dropped, so delivery never blocks the transport. + */ + def serverNotifications(bufferSize: Int = 1024)(using mat: Materializer): Source[ServerNotification, NotUsed] = + given scala.concurrent.ExecutionContext = mat.executionContext + Source + .queue[ServerNotification](bufferSize) + .watchTermination(): (queue, done) => + val listener: ServerNotificationListener[Future] = n => + val _ = queue.offer(n) + Future.unit + val _ = client.onServerNotification(listener) + val _ = done.onComplete(_ => client.removeServerNotification(listener)) + NotUsed diff --git a/client-streaming/client-pekko/src/test/scala/chimp/client/transport/pekko/PekkoServerNotificationsSpec.scala b/client-streaming/client-pekko/src/test/scala/chimp/client/transport/pekko/PekkoServerNotificationsSpec.scala new file mode 100644 index 0000000..5366eda --- /dev/null +++ b/client-streaming/client-pekko/src/test/scala/chimp/client/transport/pekko/PekkoServerNotificationsSpec.scala @@ -0,0 +1,48 @@ +package chimp.client.transport.pekko + +import chimp.client.FakeNotificationClient +import chimp.client.notifications.ServerNotification +import chimp.protocol.{ProgressParams, ProgressToken} +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.monad.{FutureMonad, MonadError} + +import scala.concurrent.duration.DurationInt +import scala.concurrent.{Await, Future} + +class PekkoServerNotificationsSpec extends AnyFlatSpec with Matchers with BeforeAndAfterAll: + + private given system: ActorSystem = ActorSystem("chimp-client-pekko-notif-test") + private given Materializer = Materializer.matFromSystem + private given MonadError[Future] = FutureMonad()(using system.dispatcher) + + override def afterAll(): Unit = + val _ = Await.result(system.terminate(), 30.seconds) + + private def progress(p: Double): ServerNotification = + ServerNotification.Progress(ProgressParams(progressToken = ProgressToken("t"), progress = p)) + + private def awaitCondition(cond: => Boolean): Unit = + val deadline = System.currentTimeMillis + 5000 + while !cond && System.currentTimeMillis < deadline do Thread.sleep(5) + + it should "emit notifications pushed by the server and remove the listener when done" in: + val client = FakeNotificationClient[Future]() + val n1 = progress(0.1) + val n2 = progress(0.2) + val n3 = progress(0.3) + + val collecting: Future[Seq[ServerNotification]] = client.serverNotifications().take(3).runWith(Sink.seq) + awaitCondition(client.listenerCount > 0) + val _ = client.emit(n1) + val _ = client.emit(n2) + val _ = client.emit(n3) + + Await.result(collecting, 5.seconds) shouldBe Seq(n1, n2, n3) + + awaitCondition(client.listenerCount == 0) + client.listenerCount shouldBe 0 diff --git a/client-streaming/client-zio/src/main/scala/chimp/client/transport/zio/ZioServerNotifications.scala b/client-streaming/client-zio/src/main/scala/chimp/client/transport/zio/ZioServerNotifications.scala new file mode 100644 index 0000000..cd8d998 --- /dev/null +++ b/client-streaming/client-zio/src/main/scala/chimp/client/transport/zio/ZioServerNotifications.scala @@ -0,0 +1,20 @@ +package chimp.client.transport.zio + +import chimp.client.BidirectionalMcpClient +import chimp.client.notifications.{ServerNotification, ServerNotificationListener} +import zio.stream.ZStream +import zio.{Queue, Task, ZIO} + +extension (client: BidirectionalMcpClient[Task]) + /** A [[ZStream]] of notifications pushed by the server. When the stream is run it registers a listener with the client, emits every + * notification the server sends, and removes the listener and shuts down the backing queue when the stream finishes. The queue is + * unbounded, so delivery never blocks the transport. + */ + def serverNotifications: ZStream[Any, Throwable, ServerNotification] = + ZStream.unwrapScoped: + for + queue <- ZIO.acquireRelease(Queue.unbounded[ServerNotification])(_.shutdown) + listener = new ServerNotificationListener[Task]: + def onNotification(n: ServerNotification): Task[Unit] = queue.offer(n).unit + _ <- ZIO.acquireRelease(client.onServerNotification(listener))(_ => client.removeServerNotification(listener).orDie) + yield ZStream.fromQueue(queue) diff --git a/client-streaming/client-zio/src/test/scala/chimp/client/transport/zio/ZioServerNotificationsSpec.scala b/client-streaming/client-zio/src/test/scala/chimp/client/transport/zio/ZioServerNotificationsSpec.scala new file mode 100644 index 0000000..e480a38 --- /dev/null +++ b/client-streaming/client-zio/src/test/scala/chimp/client/transport/zio/ZioServerNotificationsSpec.scala @@ -0,0 +1,38 @@ +package chimp.client.transport.zio + +import chimp.client.FakeNotificationClient +import chimp.client.notifications.ServerNotification +import chimp.protocol.{ProgressParams, ProgressToken} +import org.scalatest.flatspec.AnyFlatSpec +import org.scalatest.matchers.should.Matchers +import sttp.client4.impl.zio.RIOMonadAsyncError +import sttp.monad.MonadError +import zio.{Chunk, Runtime, Task, Unsafe, ZIO} + +class ZioServerNotificationsSpec extends AnyFlatSpec with Matchers: + + private given MonadError[Task] = new RIOMonadAsyncError[Any] + private val runtime: Runtime[Any] = Runtime.default + + private def run[A](task: Task[A]): A = + Unsafe.unsafe(implicit u => runtime.unsafe.run(task).getOrThrowFiberFailure()) + + private def progress(p: Double): ServerNotification = + ServerNotification.Progress(ProgressParams(progressToken = ProgressToken("t"), progress = p)) + + it should "emit notifications pushed by the server and remove the listener when done" in: + val client = FakeNotificationClient[Task]() + val n1 = progress(0.1) + val n2 = progress(0.2) + val n3 = progress(0.3) + + val program = + for + collecting <- client.serverNotifications.take(3).runCollect.fork + _ <- ZIO.succeed(client.listenerCount).repeatUntil(_ > 0) + _ <- client.emit(n1) *> client.emit(n2) *> client.emit(n3) + result <- collecting.join + yield result + + run(program) shouldBe Chunk(n1, n2, n3) + client.listenerCount shouldBe 0 diff --git a/client/src/main/scala/chimp/client/McpClient.scala b/client/src/main/scala/chimp/client/McpClient.scala index 8c33c40..9ebf69e 100644 --- a/client/src/main/scala/chimp/client/McpClient.scala +++ b/client/src/main/scala/chimp/client/McpClient.scala @@ -113,6 +113,12 @@ trait BidirectionalMcpClient[F[_]] extends McpClient[F]: /** Registers a listener for notifications pushed by the server (e.g. resource updates, tool/prompt list changes, log messages). */ def onServerNotification(listener: ServerNotificationListener[F]): F[Unit] + /** Removes a listener previously registered with [[onServerNotification]]. Listeners are compared by reference; removing a listener that + * was not registered has no effect. Used by the effect-specific `serverNotifications` streams to release their listener when the stream + * finishes. + */ + def removeServerNotification(listener: ServerNotificationListener[F]): F[Unit] + object McpClient: /** Creates an unidirectional [[McpClient]] over the given [[chimp.client.transport.ClientTransport]] and performs the initialization * handshake with the server. diff --git a/client/src/main/scala/chimp/client/McpClientImpl.scala b/client/src/main/scala/chimp/client/McpClientImpl.scala index 0b96663..26566b7 100644 --- a/client/src/main/scala/chimp/client/McpClientImpl.scala +++ b/client/src/main/scala/chimp/client/McpClientImpl.scala @@ -269,3 +269,7 @@ object McpClientImpl: override def onServerNotification(listener: ServerNotificationListener[F]): F[Unit] = val _ = serverNotificationListeners.updateAndGet(listeners => listeners :+ listener) monad.unit(()) + + override def removeServerNotification(listener: ServerNotificationListener[F]): F[Unit] = + val _ = serverNotificationListeners.updateAndGet(listeners => listeners.filterNot(_ eq listener)) + monad.unit(()) diff --git a/client/src/test/scala/chimp/client/CapabilityHandlerSpec.scala b/client/src/test/scala/chimp/client/CapabilityHandlerSpec.scala index 556a547..23c0a84 100644 --- a/client/src/test/scala/chimp/client/CapabilityHandlerSpec.scala +++ b/client/src/test/scala/chimp/client/CapabilityHandlerSpec.scala @@ -65,6 +65,23 @@ class CapabilityHandlerSpec extends AnyFlatSpec with Matchers: received shouldBe Some(ServerNotification.Progress(params)) + it should "stop delivering notifications after removeServerNotification" in: + val transport = InMemoryTransport() + planInitResponse(transport) + val client = McpClient.bidirectional[Identity](transport, clientInfo) + var count = 0 + val listener: ServerNotificationListener[Identity] = _ => { count += 1; () } + val _ = client.onServerNotification(listener) + + val params = ProgressParams(progressToken = ProgressToken("p1"), progress = 0.42) + val notification: JSONRPCMessage = JSONRPCMessage.Notification(method = "notifications/progress", params = Some(params.asJson)) + transport.simulateIncoming(notification) + count shouldBe 1 + + val _ = client.removeServerNotification(listener) + transport.simulateIncoming(notification) + count shouldBe 1 + it should "include opted-in capabilities on initialize" in: val transport = InMemoryTransport() planInitResponse(transport) diff --git a/client/src/test/scala/chimp/client/FakeNotificationClient.scala b/client/src/test/scala/chimp/client/FakeNotificationClient.scala new file mode 100644 index 0000000..0775def --- /dev/null +++ b/client/src/test/scala/chimp/client/FakeNotificationClient.scala @@ -0,0 +1,49 @@ +package chimp.client + +import chimp.client.notifications.{ServerNotification, ServerNotificationListener} +import chimp.protocol.* +import io.circe.Json +import sttp.monad.MonadError +import sttp.monad.syntax.* + +import java.util.concurrent.atomic.AtomicReference + +/** A minimal [[BidirectionalMcpClient]] for testing the effect-specific `serverNotifications` streams. Only the notification-listener + * methods are functional; every other method fails. Use [[emit]] to push a notification to all registered listeners, and [[listenerCount]] + * to observe registration and removal. + */ +final class FakeNotificationClient[F[_]](using val monad: MonadError[F]) extends BidirectionalMcpClient[F]: + private val listeners = AtomicReference[List[ServerNotificationListener[F]]](Nil) + + def listenerCount: Int = listeners.get().size + + def emit(n: ServerNotification): F[Unit] = + listeners.get().foldLeft(monad.unit(()))((acc, l) => acc.flatMap(_ => l.onNotification(n))) + + override def onServerNotification(listener: ServerNotificationListener[F]): F[Unit] = + val _ = listeners.updateAndGet(_ :+ listener) + monad.unit(()) + + override def removeServerNotification(listener: ServerNotificationListener[F]): F[Unit] = + val _ = listeners.updateAndGet(_.filterNot(_ eq listener)) + monad.unit(()) + + private def unsupported[A]: F[A] = monad.error(UnsupportedOperationException("not supported by FakeNotificationClient")) + + override val serverCapabilities: ServerCapabilities = ServerCapabilities() + override val serverInfo: Implementation = Implementation(name = "fake", version = "0.0.0") + override def ping(): F[Unit] = unsupported + override def close(): F[Unit] = monad.unit(()) + override def listTools(cursor: Option[Cursor]): F[ListToolsResponse] = unsupported + override def callTool(name: String, arguments: Json): F[CallToolResult] = unsupported + override def listPrompts(cursor: Option[Cursor]): F[ListPromptsResult] = unsupported + override def getPrompt(name: String, arguments: Map[String, String]): F[GetPromptResult] = unsupported + override def listResources(cursor: Option[Cursor]): F[ListResourcesResult] = unsupported + override def listResourceTemplates(cursor: Option[Cursor]): F[ListResourceTemplatesResult] = unsupported + override def readResource(uri: String): F[ReadResourceResult] = unsupported + override def complete(ref: CompleteRef, argument: CompleteArgument): F[CompleteResult] = unsupported + override def setLoggingLevel(level: LoggingLevel): F[Unit] = unsupported + override def sendProgress(token: ProgressToken, progress: Double, total: Option[Double], message: Option[String]): F[Unit] = unsupported + override def subscribeResource(uri: String): F[Unit] = unsupported + override def unsubscribeResource(uri: String): F[Unit] = unsupported + override def sendRootsListChanged(): F[Unit] = unsupported diff --git a/docs/client/capabilities.md b/docs/client/capabilities.md index bede01f..28f5606 100644 --- a/docs/client/capabilities.md +++ b/docs/client/capabilities.md @@ -43,3 +43,18 @@ def listen(client: BidirectionalMcpClient[Task]): Task[Unit] = case _ => ZIO.unit } ``` + +The streaming client modules also expose the notifications as a stream native to the effect backend, through the `serverNotifications` extension. The stream registers a listener when it is run and removes it when it finishes. With ZIO it is a `ZStream`: + +```scala mdoc:compile-only +import chimp.client.* +import chimp.client.notifications.ServerNotification +import chimp.client.transport.zio.* +import zio.Task +import zio.stream.ZStream + +def notifications(client: BidirectionalMcpClient[Task]): ZStream[Any, Throwable, ServerNotification] = + client.serverNotifications +``` + +The `chimp-client-ox` module gives an `ox.flow.Flow[ServerNotification]`, and `chimp-client-pekko` gives a Pekko `Source[ServerNotification, NotUsed]`.