WIP: feat: expose server notifications as an effect-specific stream

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)

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

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

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

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)

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

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.

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(())

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)

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

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]`.