pass through information about the connection for things like clientIp and server hostname (depending on the protocol)

Author: msiegenthaler

README.md

@@ -34,15 +34,15 @@ import cats.implicits.*
 import ch.linkyard.mcp.jsonrpc2.transport.StdioJsonRpcConnection
 import ch.linkyard.mcp.protocol.Initialize.PartyInfo
 import ch.linkyard.mcp.server.*
+import ch.linkyard.mcp.server.McpServer.Client
+import ch.linkyard.mcp.server.McpServer.ConnectionInfo
 import ch.linkyard.mcp.server.ToolFunction.Effect
 import com.melvinlow.json.schema.generic.auto.given
 import io.circe.generic.auto.given
 
 object SimpleEchoServer extends IOApp:
-  // Define the input/output types for your tool
   case class EchoInput(text: String)
 
-  // Create the echo tool function
   private def echoTool: ToolFunction[IO] = ToolFunction.text(
     ToolFunction.Info(
       "echo",
@@ -54,7 +54,6 @@ object SimpleEchoServer extends IOApp:
     (input: EchoInput, _) => IO(input.text),
   )
 
-  // Define your server session
   private class Session extends McpServer.Session[IO] with McpServer.ToolProvider[IO]:
     override val serverInfo: PartyInfo = PartyInfo(
       "Simple Echo MCP",
@@ -63,20 +62,18 @@ object SimpleEchoServer extends IOApp:
     override def instructions: IO[Option[String]] = None.pure
     override val tools: IO[List[ToolFunction[IO]]] = List(echoTool).pure
 
-  // Define your server
   private class Server extends McpServer[IO]:
-    override def connect(client: McpServer.Client[IO]): Resource[IO, McpServer.Session[IO]] =
-      Resource.pure(Session())
+    override def initialize(client: Client[IO], info: ConnectionInfo[IO]): Resource[IO, McpServer.Session[IO]] = ???
+    Resource.pure(Session())
 
   override def run(args: List[String]): IO[ExitCode] =
     // run with stdio transport
     Server().start(
-      StdioJsonRpcConnection.resource[IO],
+      StdioJsonRpcConnection.create[IO],
       e => IO(System.err.println(s"Error: $e")),
     ).useForever.as(ExitCode.Success)
   end run
 end SimpleEchoServer
-
 ```
 
 ### Running Your Server
@@ -175,7 +172,7 @@ A demonstration of authentication and authorization in MCP servers using Bearer
 
 - **Set up the Authentication**: Uses OAuthAuthorizationServer to guard the `/mcp` path and provide the `.well-known/oauth-protected-resource`. Pass it a token validator function to e.g. check the signature of the provided JWP.
 - **Proxy Authorization Server** (optional): Since not all OIDC IdP provide the necessary `.well-known/oauth-authorization-server` endpoint this route provides a proxy. The result will be the `.well-known/openid-configuration` of the IdP.
-- **Access the Token**: In the open session the authentication token is available as `client.authentication`. This is refreshed on every request, so it will be kept current (provided the client makes a request from time to time, eg a Ping).
+- **Access the Token**: In the open session the authentication token is available as `connectionInfo.authentication`, the connectionInfo is passed in when a session is started (McpServer.initialize). This is refreshed on every request, so it will be kept current (provided the client makes a request from time to time, eg a Ping).
 
 This example is useful for understanding how to build secure MCP servers that require user authentication and implement proper authorization controls and have access to the token.
 

build.sbt

@@ -78,6 +78,7 @@ lazy val jsonrpc2 = (project in file("jsonrpc2"))
     libraryDependencies ++= Seq(
       "co.fs2" %% "fs2-core" % Dependencies.fs2,
       "io.circe" %% "circe-core" % Dependencies.circe,
+      "com.comcast" %% "ip4s-core" % Dependencies.ip4s,
     ),
   )
 

example/demo/src/main/scala/ch/linkyard/mcp/example/demo/DemoServer.scala

@@ -15,13 +15,14 @@ import ch.linkyard.mcp.protocol.Resources.ReadResource
 import ch.linkyard.mcp.server.CallContext
 import ch.linkyard.mcp.server.McpError
 import ch.linkyard.mcp.server.McpServer
+import ch.linkyard.mcp.server.McpServer.ConnectionInfo
 import ch.linkyard.mcp.server.McpServer.Pageable
 import ch.linkyard.mcp.server.PromptFunction
 import ch.linkyard.mcp.server.ResourceTemplate
 import ch.linkyard.mcp.server.ToolFunction
 
 class DemoServer extends McpServer[IO]:
-  override def initialize(client: McpServer.Client[IO]): Resource[IO, McpServer.Session[IO]] =
+  override def initialize(client: McpServer.Client[IO], info: ConnectionInfo[IO]): Resource[IO, McpServer.Session[IO]] =
     Resource.pure(DemoSession(client))
 
   private class DemoSession(client: McpServer.Client[IO]) extends McpServer.Session[IO] with McpServer.ToolProvider[IO]

example/demo/src/main/scala/ch/linkyard/mcp/example/demo/StdioDemoMcpServer.scala

@@ -11,6 +11,6 @@ object StdioDemoMcpServer extends IOApp:
     // run with stdio transport
     IO(System.err.println("Welcome to Echo MCP")) >>
       DemoServer().start(
-        StdioJsonRpcConnection.resource[IO],
+        StdioJsonRpcConnection.create[IO],
         e => IO(System.err.println(s"Error: $e")),
       ).useForever.as(ExitCode.Success)

example/simple-authenticated/src/main/scala/ch/linkyard/mcp/example/simpleAuthenticated/TheServer.scala

@@ -7,6 +7,7 @@ import ch.linkyard.mcp.example.simpleAuthenticated.TheServer.Session
 import ch.linkyard.mcp.protocol.Initialize.PartyInfo
 import ch.linkyard.mcp.server.McpServer
 import ch.linkyard.mcp.server.McpServer.Client
+import ch.linkyard.mcp.server.McpServer.ConnectionInfo
 import ch.linkyard.mcp.server.ToolFunction
 import ch.linkyard.mcp.server.ToolFunction.Effect
 import com.melvinlow.json.schema.annotation.JsonSchemaField
@@ -15,16 +16,16 @@ import io.circe.generic.auto.given
 import io.circe.syntax.*
 
 class TheServer extends McpServer[IO]:
-  override def initialize(client: Client[IO]): Resource[IO, McpServer.Session[IO]] =
-    Resource.pure(Session(client))
+  override def initialize(client: Client[IO], info: ConnectionInfo[IO]): Resource[IO, McpServer.Session[IO]] =
+    Resource.pure(Session(info))
 
 object TheServer:
   case class HelloInput(
     @JsonSchemaField("description", "Your Name".asJson)
     name: String
   )
 
-  private def helloTool(client: Client[IO]): ToolFunction[IO] = ToolFunction.text(
+  private def helloTool(info: ConnectionInfo[IO]): ToolFunction[IO] = ToolFunction.text(
     ToolFunction.Info(
       "hello",
       "Say Hello".some,
@@ -34,14 +35,15 @@ object TheServer:
     ),
     (input: HelloInput, _) =>
       for
-        auth <- client.authentication
+        auth <- info.authentication
       yield s"Hello ${input.name}!\nYour authentication Token is $auth",
   )
 
-  private class Session(client: Client[IO]) extends McpServer.Session[IO] with McpServer.ToolProvider[IO]:
+  private class Session(info: ConnectionInfo[IO]) extends McpServer.Session[IO]
+      with McpServer.ToolProvider[IO]:
     override val serverInfo: PartyInfo = PartyInfo(
       "Simple Authenticated MCP",
       "1.0.0",
     )
     override def instructions: IO[Option[String]] = None.pure
-    override val tools: IO[List[ToolFunction[IO]]] = List(helloTool(client)).pure
+    override val tools: IO[List[ToolFunction[IO]]] = List(helloTool(info)).pure

example/simple-echo/src/main/scala/ch/linkyard/mcp/example/simpleEcho/SimpleEchoServer.scala

@@ -8,6 +8,8 @@ import cats.implicits.*
 import ch.linkyard.mcp.jsonrpc2.transport.StdioJsonRpcConnection
 import ch.linkyard.mcp.protocol.Initialize.PartyInfo
 import ch.linkyard.mcp.server.*
+import ch.linkyard.mcp.server.McpServer.Client
+import ch.linkyard.mcp.server.McpServer.ConnectionInfo
 import ch.linkyard.mcp.server.ToolFunction.Effect
 import com.melvinlow.json.schema.generic.auto.given
 import io.circe.generic.auto.given
@@ -35,13 +37,13 @@ object SimpleEchoServer extends IOApp:
     override val tools: IO[List[ToolFunction[IO]]] = List(echoTool).pure
 
   private class Server extends McpServer[IO]:
-    override def initialize(client: McpServer.Client[IO]): Resource[IO, McpServer.Session[IO]] =
-      Resource.pure(Session())
+    override def initialize(client: Client[IO], info: ConnectionInfo[IO]): Resource[IO, McpServer.Session[IO]] = ???
+    Resource.pure(Session())
 
   override def run(args: List[String]): IO[ExitCode] =
     // run with stdio transport
     Server().start(
-      StdioJsonRpcConnection.resource[IO],
+      StdioJsonRpcConnection.create[IO],
       e => IO(System.err.println(s"Error: $e")),
     ).useForever.as(ExitCode.Success)
   end run

jsonrpc2/src/main/scala/ch/linkyard/mcp/jsonrpc2/JsonRpcConnection.scala

@@ -1,5 +1,18 @@
 package ch.linkyard.mcp.jsonrpc2
 
+import com.comcast.ip4s.Host
+import com.comcast.ip4s.IpAddress
+import com.comcast.ip4s.Port
+import io.circe.Json
+
 trait JsonRpcConnection[F[_]]:
+  def info: JsonRpcConnection.Info
   def out: fs2.Pipe[F, JsonRpc.Message, Unit]
   def in: fs2.Stream[F, JsonRpc.MessageEnvelope]
+
+object JsonRpcConnection:
+  enum Info:
+    case Stdio(additional: Map[String, Json])
+    case Http(server: Option[(Host, Port)], client: Option[IpAddress], additional: Map[String, Json])
+    case Other(additional: Map[String, Json])
+    def additional: Map[String, Json]

jsonrpc2/src/main/scala/ch/linkyard/mcp/jsonrpc2/JsonRpcServer.scala

@@ -12,14 +12,12 @@ end JsonRpcServer
 object JsonRpcServer:
   def start[F[_]: Concurrent](
     server: JsonRpcServer[F],
-    connection: Resource[F, JsonRpcConnection[F]],
+    connection: JsonRpcConnection[F],
   ): Resource[F, Unit] =
-    connection.flatMap { conn =>
-      val in = conn.in
-        .through(server.handler)
-        .through(conn.out)
-      val out = server.out
-        .through(conn.out)
-      val merged = in.merge[F, Unit](out)
-      merged.compile.resource.drain
-    }
+    val in = connection.in
+      .through(server.handler)
+      .through(connection.out)
+    val out = server.out
+      .through(connection.out)
+    val merged = in.merge[F, Unit](out)
+    merged.compile.resource.drain

jsonrpc2/src/test/scala/ch/linkyard/mcp/jsonrpc2/JsonRpcServerSpec.scala

@@ -2,9 +2,9 @@ package ch.linkyard.mcp.jsonrpc2
 
 import cats.effect.IO
 import cats.effect.Ref
-import cats.effect.Resource
 import cats.effect.unsafe.implicits.global
 import ch.linkyard.mcp.jsonrpc2.JsonRpc.*
+import ch.linkyard.mcp.jsonrpc2.JsonRpcConnection.Info
 import fs2.Pipe
 import fs2.Stream
 import io.circe.JsonObject
@@ -20,6 +20,7 @@ class JsonRpcServerSpec extends AnyFunSpec with Matchers with EitherValues {
       outPipe: Pipe[IO, Message, Unit],
       outMessages: Ref[IO, List[Message]],
     ) extends JsonRpcConnection[IO] {
+      def info: Info = Info.Other(Map.empty)
       def in: Stream[IO, MessageEnvelope] = inStream
       def out: Pipe[IO, Message, Unit] = outPipe
     }
@@ -60,8 +61,7 @@ class JsonRpcServerSpec extends AnyFunSpec with Matchers with EitherValues {
         outMessages <- Ref[IO].of(List.empty[Message])
         connection = createMockConnection(List(request), outMessages)
         server = createMockServer(Map(request -> response))
-        connectionResource = Resource.pure[IO, JsonRpcConnection[IO]](connection)
-        _ <- JsonRpcServer.start(server, connectionResource).use(IO.pure)
+        _ <- JsonRpcServer.start(server, connection).use(IO.pure)
         finalMessages <- outMessages.get
       } yield finalMessages
       val messages = test.unsafeRunSync()
@@ -74,8 +74,7 @@ class JsonRpcServerSpec extends AnyFunSpec with Matchers with EitherValues {
         outMessages <- Ref[IO].of(List.empty[Message])
         connection = createMockConnection(Nil, outMessages)
         server = createMockServer(outMessages = List(serverOutMessage))
-        connectionResource = Resource.pure[IO, JsonRpcConnection[IO]](connection)
-        _ <- JsonRpcServer.start(server, connectionResource).use(IO.pure)
+        _ <- JsonRpcServer.start(server, connection).use(IO.pure)
         finalMessages <- outMessages.get
       } yield finalMessages
       val messages = test.unsafeRunSync()
@@ -93,8 +92,7 @@ class JsonRpcServerSpec extends AnyFunSpec with Matchers with EitherValues {
           handlerResponses = Map(request -> response),
           outMessages = List(serverOutMessage),
         )
-        connectionResource = Resource.pure[IO, JsonRpcConnection[IO]](connection)
-        _ <- JsonRpcServer.start(server, connectionResource).use(IO.pure)
+        _ <- JsonRpcServer.start(server, connection).use(IO.pure)
         finalMessages <- outMessages.get
       } yield finalMessages
       val messages = test.unsafeRunSync()

mcp/server/src/main/scala/ch/linkyard/mcp/server/ConnectionInfoRepr.scala

@@ -0,0 +1,30 @@
+package ch.linkyard.mcp.server
+
+import cats.MonadThrow
+import cats.effect.kernel.Async
+import cats.effect.kernel.Ref
+import cats.implicits.*
+import ch.linkyard.mcp.jsonrpc2.Authentication
+import ch.linkyard.mcp.jsonrpc2.JsonRpc.ErrorCode
+import ch.linkyard.mcp.jsonrpc2.JsonRpcConnection
+
+private class ConnectionInfoRepr[F[_]: MonadThrow] private (
+  authenticationRef: Ref[F, Authentication],
+  override val connection: JsonRpcConnection.Info,
+) extends McpServer.ConnectionInfo[F]:
+  def authentication: F[Authentication] = authenticationRef.get
+
+  private[server] def updateAuthentication(auth: Authentication): F[Unit] =
+    authenticationRef.modify(old =>
+      (old, auth) match
+        case (Authentication.Anonymous, Authentication.Anonymous)           => (auth, true)
+        case (Authentication.BearerToken(_), Authentication.BearerToken(_)) => (auth, true)
+        case _                                                              => (old, false)
+    ).ifM(
+      ().pure[F],
+      McpError.raise(ErrorCode.InvalidRequest, "Cannot change authentication mode after initialization").void,
+    )
+
+private object ConnectionInfoRepr:
+  def apply[F[_]: Async](authentication: Authentication, connection: JsonRpcConnection.Info): F[ConnectionInfoRepr[F]] =
+    Ref.of[F, Authentication](authentication).map(new ConnectionInfoRepr[F](_, connection))