From 497dfbb564290a190194900de8d2e2c61ef25d27 Mon Sep 17 00:00:00 2001 From: Maxime P <7041978+max-peroch@users.noreply.github.com> Date: Tue, 8 Sep 2026 13:22:34 -0700 Subject: [PATCH 1/5] Give the tools a principal made by the server security logic --- docs/server/transport.md | 38 +++++- .../main/scala/chimp/server/McpServer.scala | 100 +++++++++++++++- .../scala/chimp/server/ServerContext.scala | 8 ++ server/src/main/scala/chimp/server/Tool.scala | 16 +++ .../SecuredServerHttpTransport.scala | 45 ++++++++ .../server/SecuredHttpMcpServerSpec.scala | 109 ++++++++++++++++++ 6 files changed, 313 insertions(+), 3 deletions(-) create mode 100644 server/src/main/scala/chimp/server/transport/SecuredServerHttpTransport.scala create mode 100644 server/src/test/scala/chimp/server/SecuredHttpMcpServerSpec.scala diff --git a/docs/server/transport.md b/docs/server/transport.md index 5a33b68..c7ffaea 100644 --- a/docs/server/transport.md +++ b/docs/server/transport.md @@ -93,7 +93,7 @@ import sttp.model.StatusCode import sttp.tapir.* import sttp.tapir.server.netty.sync.NettySyncServer -object SecuredMcpServer: +object McpServerWithPrependedSecurity: def main(args: Array[String]): Unit = val adder = tool("echo").input[String].handle(echo => ToolResult.text(echo)) val mcpEndpoint = McpServer(tools = List(adder)).endpoint(List("mcp")) @@ -106,6 +106,40 @@ object SecuredMcpServer: NettySyncServer().port(8080).addEndpoint(securedEndpoint).startAndWait() ``` -The result of the security logic does not reach the tool logic. If a tool needs data from the caller, read the request headers with `handleWithHeaders` (or `serverLogic` with headers). +The result of the security logic of `prependSecurity` does not reach the tool logic. Use `prependSecurity` if the tools do not need data from the caller. If a tool needs such data, read the request headers with `handleWithHeaders` (or `serverLogic` with headers), or give the tools a principal, as below. For all the security inputs - API keys, basic and bearer authorization, OAuth2 flows - see the [Tapir endpoint security documentation](https://tapir.softwaremill.com/en/latest/endpoint/security.html). + +### Giving the security result to the tools + +To validate the caller one time and give the result to the tool logic, use `serverSecurityLogic` (or `serverSecurityLogicPure`, if the logic needs no effect). It takes the same security input and error output as `prependSecurity`, and makes a principal - a value of your own type, such as the identity of the caller. + +The server gives the principal to the logic of each tool which you add to it. Define such a tool with `handleSecured`, or with `securedServerLogic` if the logic needs an effect. Tools which do not need the principal keep their usual logic. The other builders of `McpServer` stay available, so you can configure the server before or after you add the security logic: + +```scala mdoc:compile-only +import chimp.server.* +import sttp.model.StatusCode +import sttp.shared.Identity +import sttp.tapir.* +import sttp.tapir.server.netty.sync.NettySyncServer + +case class User(email: String) + +object McpServerWithPrincipal: + def main(args: Array[String]): Unit = + val echo = tool("echo").input[String].handle(message => ToolResult.text(message)) + val whoAmI = tool("whoAmI").input[String].handleSecured[User]((_, user) => ToolResult.text(user.email)) + + val securedEndpoint = McpServer[Identity]() + .serverSecurityLogicPure( + auth.bearer[String](), + statusCode(StatusCode.Unauthorized).and(stringBody) + )(token => if token == "s3cret" then Right(User("employee@example.com")) else Left("Invalid token")) + .name("my-mcp-server") + .addTools(echo, whoAmI) + .endpoint(List("mcp")) + + NettySyncServer().port(8080).addEndpoint(securedEndpoint).startAndWait() +``` + +The security logic runs one time for each request, before the server handles the MCP message. If it gives a rejection, the server sends the error output and no tool logic runs. This is necessary if the client must get an HTTP status code, because a tool which rejects a call can only give a JSON-RPC error with status code 200. diff --git a/server/src/main/scala/chimp/server/McpServer.scala b/server/src/main/scala/chimp/server/McpServer.scala index 75fca77..1e500fa 100644 --- a/server/src/main/scala/chimp/server/McpServer.scala +++ b/server/src/main/scala/chimp/server/McpServer.scala @@ -1,8 +1,10 @@ package chimp.server import chimp.protocol.* -import chimp.server.transport.ServerHttpTransport +import chimp.server.transport.{SecuredServerHttpTransport, ServerHttpTransport} +import sttp.monad.MonadError import sttp.tapir.server.ServerEndpoint +import sttp.tapir.{EndpointInput, EndpointOutput} type CompletionHandler[F[_]] = (CompleteRef, CompleteArgument, Option[CompleteContext]) => F[Completion] @@ -91,6 +93,20 @@ case class McpServer[F[_]]( def endpoint(path: List[String]): ServerEndpoint[Any, F] = ServerHttpTransport(path).serve(this) + /** Adds the security input, the error output which describes a rejection, and the logic which validates the security input and makes the + * principal. The principal is given to the logic of the tools which are added to the returned server. + */ + def serverSecurityLogic[S, E, P](securityInput: EndpointInput[S], errorOutput: EndpointOutput[E])( + logic: S => F[Either[E, P]] + ): SecuredMcpServer[F, S, E, P] = + SecuredMcpServer(this, securityInput, errorOutput, _ => logic) + + /** The same as [[serverSecurityLogic]], but for security logic which needs no effect. */ + def serverSecurityLogicPure[S, E, P](securityInput: EndpointInput[S], errorOutput: EndpointOutput[E])( + logic: S => Either[E, P] + ): SecuredMcpServer[F, S, E, P] = + SecuredMcpServer(this, securityInput, errorOutput, monad => input => monad.unit(logic(input))) + def streaming: StreamingMcpServer[F] = StreamingMcpServer( name, @@ -174,3 +190,85 @@ case class StreamingMcpServer[F[_]]( def withSubscriptions(handler: ResourceSubscriptions[F]): StreamingMcpServer[F] = copy(subscriptions = Some(handler)) + +/** An [[McpServer]] with security logic, which runs before the server handles an MCP message. The result of the security logic, the + * principal, is given to the logic of the tools which are added to this server. Tools of the initial server, which do not need the + * principal, are kept. + * + * @tparam S + * The type of the security input, for example a bearer token. + * @tparam E + * The type of the error output, which the server sends if the security logic gives a rejection. + * @tparam P + * The type of the principal, which the security logic makes from the security input. + */ +case class SecuredMcpServer[F[_], S, E, P]( + server: McpServer[F], + securityInput: EndpointInput[S], + errorOutput: EndpointOutput[E], + securityLogic: MonadError[F] => S => F[Either[E, P]], + securedTools: List[ServerTool[?, ?, F, SecuredServerContext[F, P]]] = Nil +) extends McpServerDef[F, SecuredServerContext[F, P]]: + def name: String = server.name + def version: String = server.version + def instructions: Option[String] = server.instructions + def showJsonSchemaMetadata: Boolean = server.showJsonSchemaMetadata + def originCheck: OriginCheck = server.originCheck + def prompts: List[ServerPrompt[F]] = server.prompts + def resources: List[ServerResource[F]] = server.resources + def resourceTemplates: List[ServerResourceTemplate[F]] = server.resourceTemplates + def completion: Option[CompletionHandler[F]] = server.completion + def loggingLevel: Option[SetLoggingLevelHandler[F]] = server.loggingLevel + def subscriptions: Option[ResourceSubscriptions[F]] = server.subscriptions + + def tools: List[ServerTool[?, ?, F, SecuredServerContext[F, P]]] = server.tools ++ securedTools + + def name(value: String): SecuredMcpServer[F, S, E, P] = + copy(server = server.name(value)) + + def version(value: String): SecuredMcpServer[F, S, E, P] = + copy(server = server.version(value)) + + def instructions(value: String): SecuredMcpServer[F, S, E, P] = + copy(server = server.instructions(value)) + + def withJsonSchemaMetadata(value: Boolean): SecuredMcpServer[F, S, E, P] = + copy(server = server.withJsonSchemaMetadata(value)) + + def withOriginCheck(value: OriginCheck): SecuredMcpServer[F, S, E, P] = + copy(server = server.withOriginCheck(value)) + + def addTool(tool: ServerTool[?, ?, F, SecuredServerContext[F, P]]): SecuredMcpServer[F, S, E, P] = + copy(securedTools = securedTools :+ tool) + + def addTools(tools: ServerTool[?, ?, F, SecuredServerContext[F, P]]*): SecuredMcpServer[F, S, E, P] = + copy(securedTools = this.securedTools ++ tools) + + def addPrompt(prompt: ServerPrompt[F]): SecuredMcpServer[F, S, E, P] = + copy(server = server.addPrompt(prompt)) + + def addPrompts(prompts: ServerPrompt[F]*): SecuredMcpServer[F, S, E, P] = + copy(server = server.addPrompts(prompts*)) + + def addResource(resource: ServerResource[F]): SecuredMcpServer[F, S, E, P] = + copy(server = server.addResource(resource)) + + def addResources(resources: ServerResource[F]*): SecuredMcpServer[F, S, E, P] = + copy(server = server.addResources(resources*)) + + def addResourceTemplate(resourceTemplate: ServerResourceTemplate[F]): SecuredMcpServer[F, S, E, P] = + copy(server = server.addResourceTemplate(resourceTemplate)) + + def addResourceTemplates(resourceTemplates: ServerResourceTemplate[F]*): SecuredMcpServer[F, S, E, P] = + copy(server = server.addResourceTemplates(resourceTemplates*)) + + def withCompletion(handler: CompletionHandler[F]): SecuredMcpServer[F, S, E, P] = + copy(server = server.withCompletion(handler)) + + def withLoggingLevel(handler: SetLoggingLevelHandler[F]): SecuredMcpServer[F, S, E, P] = + copy(server = server.withLoggingLevel(handler)) + + def withSubscriptions(handler: ResourceSubscriptions[F]): SecuredMcpServer[F, S, E, P] = + copy(server = server.withSubscriptions(handler)) + + def endpoint(path: List[String]): ServerEndpoint[Any, F] = SecuredServerHttpTransport[F, S, E, P](path).serve(this) diff --git a/server/src/main/scala/chimp/server/ServerContext.scala b/server/src/main/scala/chimp/server/ServerContext.scala index c764818..5ddd46e 100644 --- a/server/src/main/scala/chimp/server/ServerContext.scala +++ b/server/src/main/scala/chimp/server/ServerContext.scala @@ -10,6 +10,14 @@ trait ServerContext[F[_]] object ServerContext: def noop[F[_]]: ServerContext[F] = new ServerContext[F] {} +/** A context which also gives the principal that the security logic of a [[SecuredMcpServer]] made from the request. */ +trait SecuredServerContext[F[_], +P] extends ServerContext[F]: + def principal: P + +object SecuredServerContext: + def apply[F[_], P](value: P): SecuredServerContext[F, P] = new SecuredServerContext[F, P]: + def principal: P = value + trait StreamingServerContext[F[_]] extends ServerContext[F]: def reportProgress(progress: Double, total: Option[Double] = None, message: Option[String] = None): F[Unit] def log(level: LoggingLevel, data: Json, logger: Option[String] = None): F[Unit] diff --git a/server/src/main/scala/chimp/server/Tool.scala b/server/src/main/scala/chimp/server/Tool.scala index 1e4f489..6ee06f8 100644 --- a/server/src/main/scala/chimp/server/Tool.scala +++ b/server/src/main/scala/chimp/server/Tool.scala @@ -103,6 +103,18 @@ case class Tool[I, O]( def serverLogic[F[_]](logic: (I, Seq[Header]) => F[ToolResult[O]]): ServerTool[I, O, F, ServerContext[F]] = ServerTool(name, description, inputSchema, inputDecoder, outputSchema, annotations, (input, _, headers) => logic(input, headers)) + /** Attaches effectful logic, with access to the principal; usable only on a [[SecuredMcpServer]]. */ + def securedServerLogic[F[_], P](logic: (I, P, Seq[Header]) => F[ToolResult[O]]): ServerTool[I, O, F, SecuredServerContext[F, P]] = + ServerTool( + name, + description, + inputSchema, + inputDecoder, + outputSchema, + annotations, + (input, context, headers) => logic(input, context.principal, headers) + ) + /** Attaches effectful logic with access to the [[StreamingServerContext]]; usable only on a streaming server. */ def streamingServerLogic[F[_]]( logic: (I, StreamingServerContext[F], Seq[Header]) => F[ToolResult[O]] @@ -117,6 +129,10 @@ case class Tool[I, O]( def handle(logic: I => ToolResult[O]): ServerTool[I, O, Identity, ServerContext[Identity]] = handleWithHeaders((i, _) => logic(i)) + /** Attaches synchronous logic over the decoded input and the principal; usable only on a [[SecuredMcpServer]]. */ + def handleSecured[P](logic: (I, P) => ToolResult[O]): ServerTool[I, O, Identity, SecuredServerContext[Identity, P]] = + securedServerLogic[Identity, P]((i, principal, _) => logic(i, principal)) + /** A fully-defined tool: its metadata plus the logic handling a call, in effect `F` with context `C`. */ case class ServerTool[I, O, F[_], -C <: ServerContext[F]]( name: String, diff --git a/server/src/main/scala/chimp/server/transport/SecuredServerHttpTransport.scala b/server/src/main/scala/chimp/server/transport/SecuredServerHttpTransport.scala new file mode 100644 index 0000000..0d43250 --- /dev/null +++ b/server/src/main/scala/chimp/server/transport/SecuredServerHttpTransport.scala @@ -0,0 +1,45 @@ +package chimp.server.transport + +import chimp.server.* +import io.circe.Json +import sttp.model.{Header, HeaderNames, StatusCode} +import sttp.monad.MonadError +import sttp.monad.syntax.* +import sttp.tapir.* +import sttp.tapir.json.circe.* +import sttp.tapir.server.ServerEndpoint + +/** Implementation of unidirectional MCP server using Streamable HTTP, protected by the security logic of a [[SecuredMcpServer]]. The + * security logic runs before any MCP message is handled. If it gives a rejection, the server sends the error output and no tool logic + * runs. If it gives a principal, the principal goes to the tool logic. + * + * @param path + * The MCP endpoint path. + */ +final case class SecuredServerHttpTransport[F[_], S, E, P](path: List[String]): + def serve(server: SecuredMcpServer[F, S, E, P]): ServerEndpoint[Any, F] = + val handler = new McpHandler(server) + val mcpEndpoint = endpoint.post + .securityIn(server.securityInput) + .in(path.foldLeft(emptyInput)((inputSoFar, pathComponent) => inputSoFar / pathComponent)) + .in(extractFromRequest(_.headers)) + .in(jsonBody[Json]) + .errorOut(server.errorOutput) + .out(statusCode) + .out(jsonBody[Option[Json]]) + + ServerEndpoint( + mcpEndpoint, + server.securityLogic, + me => { (principal: P) => (input: (Seq[Header], Json)) => + val (headers, json) = input + given MonadError[F] = me + val host = headers.find(_.name.equalsIgnoreCase(HeaderNames.Host)).map(_.value) + val origin = headers.find(_.name.equalsIgnoreCase(HeaderNames.Origin)).map(_.value) + if !server.originCheck.validate(host, origin) then me.unit(Right((StatusCode.Forbidden, None))) + else + handler + .handleJsonRpc(json, headers, _ => SecuredServerContext[F, P](principal)) + .map(response => Right((response.statusCode, response.body))) + } + ) diff --git a/server/src/test/scala/chimp/server/SecuredHttpMcpServerSpec.scala b/server/src/test/scala/chimp/server/SecuredHttpMcpServerSpec.scala new file mode 100644 index 0000000..3154823 --- /dev/null +++ b/server/src/test/scala/chimp/server/SecuredHttpMcpServerSpec.scala @@ -0,0 +1,109 @@ +package chimp.server + +import chimp.client.{McpAuthorizationException, McpClient} +import chimp.client.transport.ClientHttpTransport +import chimp.protocol.{Implementation, ResourceContents, ToolContent} +import io.circe.{Codec, Json} +import org.scalatest.flatspec.AnyFlatSpec +import org.scalatest.matchers.should.Matchers +import ox.supervised +import sttp.client4.* +import sttp.model.{Header, StatusCode} +import sttp.shared.Identity +import sttp.tapir.* +import sttp.tapir.server.ServerEndpoint +import sttp.tapir.server.netty.sync.NettySyncServer + +class SecuredHttpMcpServerSpec extends AnyFlatSpec with Matchers: + private case class EchoInput(message: String) derives Codec, Schema + private case class User(email: String) + + private val clientInfo = Implementation("chimp-server-test", "0.0.1") + private val validToken = "s3cret" + + private val echoTool = tool("echo") + .description("Echoes a message.") + .input[EchoInput] + .handle(in => ToolResult.text(in.message)) + + private val whoAmITool = tool("whoAmI") + .description("Echoes a message and the caller's email.") + .input[EchoInput] + .handleSecured[User]((in, user) => ToolResult.text(s"${in.message} ${user.email}")) + + private val greetingResource = resource("test://greeting") + .name("greeting") + .mimeType("text/plain") + .handle(() => Right(List(ResourceContents.Text(uri = "test://greeting", text = "hello", mimeType = Some("text/plain"))))) + + private def securityLogic(token: String): Either[String, User] = + if token == validToken then Right(User("employee@example.com")) else Left("Invalid token") + + private val securedServer = McpServer[Identity]() + .addTool(echoTool) + .serverSecurityLogicPure(auth.bearer[String](), statusCode(StatusCode.Unauthorized).and(stringBody))(securityLogic) + .addTool(whoAmITool) + + private val serverConfiguredAfterSecurity = McpServer[Identity]() + .serverSecurityLogicPure(auth.bearer[String](), statusCode(StatusCode.Unauthorized).and(stringBody))(securityLogic) + .name("secured-server") + .version("2.0.0") + .addResource(greetingResource) + .addTool(whoAmITool) + + private def withServer[T](mcpEndpoint: ServerEndpoint[Any, Identity])(test: (Int, SyncBackend) => T): T = + supervised: + val binding = NettySyncServer().port(0).addEndpoint(mcpEndpoint).start() + try + val backend = DefaultSyncBackend() + try test(binding.port, backend) + finally backend.close() + finally binding.stop() + + private def withClient[T](port: Int, backend: SyncBackend, token: String)(test: McpClient[Identity] => T): T = + val transport = + ClientHttpTransport[Identity](backend, uri"http://localhost:$port/mcp", headers = List(Header.authorization("Bearer", token))) + try test(McpClient(transport, clientInfo)) + finally transport.close() + + private val securedEndpoint = securedServer.endpoint(List("mcp")) + private val endpointConfiguredAfterSecurity = serverConfiguredAfterSecurity.endpoint(List("mcp")) + + "a secured MCP server" should "give the principal to the tool logic" in withServer(securedEndpoint): (port, backend) => + withClient(port, backend, validToken): client => + val result = client.callTool("whoAmI", Json.obj("message" -> Json.fromString("hi"))) + result.isError shouldBe false + result.content shouldBe List(ToolContent.Text("text", "hi employee@example.com")) + + it should "also serve the tools which do not need the principal" in withServer(securedEndpoint): (port, backend) => + withClient(port, backend, validToken): client => + client.listTools().tools.map(_.name) should contain allOf ("echo", "whoAmI") + client.callTool("echo", Json.obj("message" -> Json.fromString("hi"))).content shouldBe List(ToolContent.Text("text", "hi")) + + it should "use the identity which the builders set after the security logic" in + withServer(endpointConfiguredAfterSecurity): (port, backend) => + withClient(port, backend, validToken): client => + client.serverInfo shouldBe Implementation("secured-server", "2.0.0") + + it should "serve a resource which was added after the security logic" in + withServer(endpointConfiguredAfterSecurity): (port, backend) => + withClient(port, backend, validToken): client => + client.serverCapabilities.resources shouldBe defined + client.listResources().resources.map(_.uri) shouldBe List("test://greeting") + client.readResource("test://greeting").contents.head match + case ResourceContents.Text(_, text, _, _) => text shouldBe "hello" + case other => fail(s"expected text contents, got $other") + + it should "reject an invalid security input with the error output, before any tool logic runs" in + withServer(securedEndpoint): (port, backend) => + val exception = intercept[McpAuthorizationException](withClient(port, backend, "wrong")(_ => ())) + exception.statusCode shouldBe StatusCode.Unauthorized.code + + it should "reject a request with no security input with the error output" in withServer(securedEndpoint): (port, backend) => + val response = basicRequest + .post(uri"http://localhost:$port/mcp") + .header("Content-Type", "application/json") + .body("""{"jsonrpc":"2.0","id":"1","method":"tools/call","params":{"name":"whoAmI","arguments":{"message":"hi"}}}""") + .send(backend) + response.code shouldBe StatusCode.Unauthorized + response.body shouldBe Left("Invalid value for: header Authorization (missing)") From b6ce4003f6679b44cfa2e91f3d5ae1d2660fe2fd Mon Sep 17 00:00:00 2001 From: Maxime Perocheau Date: Tue, 8 Sep 2026 21:31:13 -0700 Subject: [PATCH 2/5] Add streaming support to secured MCP servers --- docs/server/transport.md | 42 ++++++++++ .../chimp/server/ox/OxMcpServerHttpSpec.scala | 41 +++++++++- .../server/pekko/PekkoMcpServerHttpSpec.scala | 42 +++++++++- .../server/zio/ZioMcpServerHttpSpec.scala | 40 ++++++++- .../main/scala/chimp/server/McpServer.scala | 81 +++++++++++++++++++ .../scala/chimp/server/ServerContext.scala | 16 ++++ server/src/main/scala/chimp/server/Tool.scala | 16 ++++ .../transport/JsonRpcHttpResponses.scala | 33 ++++++++ .../SecuredServerHttpTransport.scala | 12 +-- .../SecuredServerStreamingHttpTransport.scala | 49 +++++++++++ .../transport/ServerHttpTransport.scala | 10 +-- .../ServerStreamingHttpTransport.scala | 46 ++++++----- .../SecuredMcpServerStreamingTests.scala | 60 ++++++++++++++ 13 files changed, 445 insertions(+), 43 deletions(-) create mode 100644 server/src/main/scala/chimp/server/transport/JsonRpcHttpResponses.scala create mode 100644 server/src/main/scala/chimp/server/transport/SecuredServerStreamingHttpTransport.scala create mode 100644 server/src/test/scala/chimp/server/SecuredMcpServerStreamingTests.scala diff --git a/docs/server/transport.md b/docs/server/transport.md index c7ffaea..b7b191a 100644 --- a/docs/server/transport.md +++ b/docs/server/transport.md @@ -143,3 +143,45 @@ object McpServerWithPrincipal: ``` The security logic runs one time for each request, before the server handles the MCP message. If it gives a rejection, the server sends the error output and no tool logic runs. This is necessary if the client must get an HTTP status code, because a tool which rejects a call can only give a JSON-RPC error with status code 200. + +### Combining security with streaming + +`.streaming` on a secured server gives a `SecuredStreamingMcpServer`, which accepts streaming tools alongside the plain secured ones. Define such a tool with `securedStreamingServerLogic`, which gives both the principal and the `StreamingServerContext`: + +```scala mdoc:compile-only +import chimp.protocol.LoggingLevel +import chimp.server.* +import chimp.server.ox.OxServerHttpTransport +import chimp.server.transport.SecuredServerStreamingHttpTransport +import io.circe.Json +import sttp.model.StatusCode +import sttp.shared.Identity +import sttp.tapir.* +import sttp.tapir.server.netty.sync.NettySyncServer + +case class User(email: String) + +object McpServerWithPrincipalAndStreaming: + def main(args: Array[String]): Unit = + val whoAmI = tool("whoAmI") + .input[String] + .securedStreamingServerLogic[Identity, User]: (_, user, ctx, _) => + ctx.log(LoggingLevel.Info, Json.fromString(s"called by ${user.email}")) + ToolResult.text(user.email) + + val securedStreamingServer = McpServer[Identity]() + .serverSecurityLogicPure( + auth.bearer[String](), + statusCode(StatusCode.Unauthorized).and(stringBody) + )(token => if token == "s3cret" then Right(User("employee@example.com")) else Left("Invalid token")) + .streaming + .addStreamingTool(whoAmI) + + // `OxServerHttpTransport` is one effect backend's streaming machinery; substitute your own (Pekko, ZIO, ...). + val backend = OxServerHttpTransport(List("mcp")) + val securedEndpoint = SecuredServerStreamingHttpTransport(List("mcp"), backend).serve(securedStreamingServer) + + NettySyncServer().port(8080).addEndpoint(securedEndpoint).startAndWait() +``` + +`SecuredServerStreamingHttpTransport` wraps a `StreamingBackend` rather than extending it, so the same effect-specific backend instance - `OxServerHttpTransport`, `PekkoServerHttpTransport`, `ZioServerHttpTransport` - serves both a plain `StreamingMcpServer` and a `SecuredStreamingMcpServer`. diff --git a/server-streaming/server-ox/src/test/scala/chimp/server/ox/OxMcpServerHttpSpec.scala b/server-streaming/server-ox/src/test/scala/chimp/server/ox/OxMcpServerHttpSpec.scala index de51b3b..2f946da 100644 --- a/server-streaming/server-ox/src/test/scala/chimp/server/ox/OxMcpServerHttpSpec.scala +++ b/server-streaming/server-ox/src/test/scala/chimp/server/ox/OxMcpServerHttpSpec.scala @@ -4,17 +4,31 @@ import chimp.client.transport.ClientTransport import chimp.client.transport.ox.OxClientHttpTransport import chimp.client.{BidirectionalMcpClient, McpClient} import chimp.protocol.{Implementation, ProtocolVersion} -import chimp.server.{McpServer, McpServerStreamingTests, McpServerTests, StreamingMcpServer, SyncToFuture} +import chimp.server.transport.SecuredServerStreamingHttpTransport +import chimp.server.{ + McpServer, + McpServerStreamingTests, + McpServerTests, + SecuredMcpServerStreamingTests, + SecuredStreamingMcpServer, + StreamingMcpServer, + SyncToFuture +} import org.scalatest.Assertion import ox.supervised import sttp.client4.DefaultSyncBackend +import sttp.model.Header import sttp.model.Uri.UriContext import sttp.shared.Identity import sttp.tapir.server.netty.sync.NettySyncServer import scala.concurrent.Future -class OxMcpServerHttpSpec extends McpServerTests[Identity] with McpServerStreamingTests[Identity] with SyncToFuture: +class OxMcpServerHttpSpec + extends McpServerTests[Identity] + with McpServerStreamingTests[Identity] + with SecuredMcpServerStreamingTests[Identity] + with SyncToFuture: private val clientInfo = Implementation("chimp-server-test", "0.0.1") override protected def withServer(server: McpServer[Identity])(test: McpClient[Identity] => Identity[Assertion]): Future[Assertion] = @@ -41,3 +55,26 @@ class OxMcpServerHttpSpec extends McpServerTests[Identity] with McpServerStreami finally transport.close() finally backend.close() finally binding.stop() + + override protected def withSecuredStreamingServer( + server: SecuredStreamingMcpServer[Identity, String, String, User] + )(test: BidirectionalMcpClient[Identity] => Identity[Assertion]): Future[Assertion] = + toFuture: + supervised: + val endpoint = SecuredServerStreamingHttpTransport(List("mcp"), OxServerHttpTransport(List("mcp"))).serve(server) + val binding = NettySyncServer().port(0).addEndpoint(endpoint).start() + try + val backend = DefaultSyncBackend() + try + val transport = + OxClientHttpTransport( + backend, + uri"http://localhost:${binding.port}/mcp", + ProtocolVersion.Latest, + ClientTransport.defaultTimeout, + headers = List(Header.authorization("Bearer", validToken)) + ) + try test(McpClient.bidirectional(transport, clientInfo)) + finally transport.close() + finally backend.close() + finally binding.stop() diff --git a/server-streaming/server-pekko/src/test/scala/chimp/server/pekko/PekkoMcpServerHttpSpec.scala b/server-streaming/server-pekko/src/test/scala/chimp/server/pekko/PekkoMcpServerHttpSpec.scala index 016c7c8..32895d9 100644 --- a/server-streaming/server-pekko/src/test/scala/chimp/server/pekko/PekkoMcpServerHttpSpec.scala +++ b/server-streaming/server-pekko/src/test/scala/chimp/server/pekko/PekkoMcpServerHttpSpec.scala @@ -3,17 +3,30 @@ package chimp.server.pekko import chimp.client.transport.pekko.PekkoClientHttpTransport import chimp.client.{BidirectionalMcpClient, McpClient} import chimp.protocol.Implementation -import chimp.server.{McpServer, McpServerStreamingTests, McpServerTests, StreamingMcpServer} +import chimp.server.transport.SecuredServerStreamingHttpTransport +import chimp.server.{ + McpServer, + McpServerStreamingTests, + McpServerTests, + SecuredMcpServerStreamingTests, + SecuredStreamingMcpServer, + StreamingMcpServer +} import org.apache.pekko.http.scaladsl.Http import org.scalatest.Assertion import sttp.client4.pekkohttp.PekkoHttpBackend +import sttp.model.Header import sttp.model.Uri.UriContext import sttp.tapir.server.pekkohttp.PekkoHttpServerInterpreter import scala.concurrent.duration.DurationInt import scala.concurrent.{ExecutionContext, Future} -class PekkoMcpServerHttpSpec extends McpServerTests[Future] with McpServerStreamingTests[Future] with PekkoToFuture: +class PekkoMcpServerHttpSpec + extends McpServerTests[Future] + with McpServerStreamingTests[Future] + with SecuredMcpServerStreamingTests[Future] + with PekkoToFuture: private val clientInfo = Implementation("chimp-server-test", "0.0.1") override protected def withServer(server: McpServer[Future])(test: McpClient[Future] => Future[Assertion]): Future[Assertion] = @@ -39,3 +52,28 @@ class PekkoMcpServerHttpSpec extends McpServerTests[Future] with McpServerStream .transformWith(_ => backend.close()) .transformWith(_ => binding.terminate(5.seconds)) .transform(_ => result) + + override protected def withSecuredStreamingServer( + server: SecuredStreamingMcpServer[Future, String, String, User] + )(test: BidirectionalMcpClient[Future] => Future[Assertion]): Future[Assertion] = + given ExecutionContext = actorSystem.dispatcher + val endpoint = SecuredServerStreamingHttpTransport(List("mcp"), PekkoServerHttpTransport(List("mcp"))).serve(server) + Http() + .newServerAt("localhost", 0) + .bind(PekkoHttpServerInterpreter().toRoute(endpoint)) + .flatMap: binding => + val backend = PekkoHttpBackend.usingActorSystem(actorSystem) + val transport = PekkoClientHttpTransport( + backend, + uri"http://localhost:${binding.localAddress.getPort}/mcp", + headers = List(Header.authorization("Bearer", validToken)) + ) + McpClient + .bidirectional(transport, clientInfo) + .flatMap(test) + .transformWith: result => + transport + .close() + .transformWith(_ => backend.close()) + .transformWith(_ => binding.terminate(5.seconds)) + .transform(_ => result) diff --git a/server-streaming/server-zio/src/test/scala/chimp/server/zio/ZioMcpServerHttpSpec.scala b/server-streaming/server-zio/src/test/scala/chimp/server/zio/ZioMcpServerHttpSpec.scala index 32525b8..f4da5c0 100644 --- a/server-streaming/server-zio/src/test/scala/chimp/server/zio/ZioMcpServerHttpSpec.scala +++ b/server-streaming/server-zio/src/test/scala/chimp/server/zio/ZioMcpServerHttpSpec.scala @@ -4,17 +4,30 @@ import chimp.client.transport.ClientTransport import chimp.client.transport.zio.ZioClientHttpTransport import chimp.client.{BidirectionalMcpClient, McpClient} import chimp.protocol.{Implementation, ProtocolVersion} -import chimp.server.{McpServer, McpServerStreamingTests, McpServerTests, StreamingMcpServer} +import chimp.server.transport.SecuredServerStreamingHttpTransport +import chimp.server.{ + McpServer, + McpServerStreamingTests, + McpServerTests, + SecuredMcpServerStreamingTests, + SecuredStreamingMcpServer, + StreamingMcpServer +} import org.scalatest.Assertion import sttp.client4.* import sttp.client4.httpclient.zio.HttpClientZioBackend +import sttp.model.Header import sttp.tapir.server.ziohttp.ZioHttpInterpreter import zio.http.Server import zio.{Scope, Task, ZIO} import scala.concurrent.Future -class ZioMcpServerHttpSpec extends McpServerTests[Task] with McpServerStreamingTests[Task] with ZioToFuture: +class ZioMcpServerHttpSpec + extends McpServerTests[Task] + with McpServerStreamingTests[Task] + with SecuredMcpServerStreamingTests[Task] + with ZioToFuture: private val clientInfo = Implementation("chimp-server-test", "0.0.1") override protected def withServer(server: McpServer[Task])(test: McpClient[Task] => Task[Assertion]): Future[Assertion] = @@ -35,3 +48,26 @@ class ZioMcpServerHttpSpec extends McpServerTests[Task] with McpServerStreamingT .flatMap(client => test(client)) .ensuring(backend.close().ignore) yield result).provideSome[Scope](Server.defaultWithPort(0)) + + override protected def withSecuredStreamingServer( + server: SecuredStreamingMcpServer[Task, String, String, User] + )(test: BidirectionalMcpClient[Task] => Task[Assertion]): Future[Assertion] = + toFuture: + val routes = + ZioHttpInterpreter().toHttp(SecuredServerStreamingHttpTransport(List("mcp"), ZioServerHttpTransport(List("mcp"))).serve(server)) + ZIO.scoped: + (for + port <- Server.install(routes) + result <- HttpClientZioBackend().flatMap: backend => + ZioClientHttpTransport + .scoped( + backend, + uri"http://localhost:$port/mcp", + ProtocolVersion.Latest, + ClientTransport.defaultTimeout, + headers = List(Header.authorization("Bearer", validToken)) + ) + .flatMap(transport => McpClient.bidirectional(transport, clientInfo)) + .flatMap(client => test(client)) + .ensuring(backend.close().ignore) + yield result).provideSome[Scope](Server.defaultWithPort(0)) diff --git a/server/src/main/scala/chimp/server/McpServer.scala b/server/src/main/scala/chimp/server/McpServer.scala index 1e500fa..3fe3ff0 100644 --- a/server/src/main/scala/chimp/server/McpServer.scala +++ b/server/src/main/scala/chimp/server/McpServer.scala @@ -272,3 +272,84 @@ case class SecuredMcpServer[F[_], S, E, P]( copy(server = server.withSubscriptions(handler)) def endpoint(path: List[String]): ServerEndpoint[Any, F] = SecuredServerHttpTransport[F, S, E, P](path).serve(this) + + def streaming: SecuredStreamingMcpServer[F, S, E, P] = SecuredStreamingMcpServer(this) + +/** A [[SecuredMcpServer]] which also accepts streaming tools, which are given a [[SecuredStreamingServerContext]] combining the + * principal with the [[StreamingServerContext]]. Tools of the initial secured server, which need only the principal, are kept. + */ +case class SecuredStreamingMcpServer[F[_], S, E, P]( + server: SecuredMcpServer[F, S, E, P], + streamingTools: List[ServerTool[?, ?, F, SecuredStreamingServerContext[F, P]]] = Nil +) extends McpServerDef[F, SecuredStreamingServerContext[F, P]]: + def name: String = server.name + def version: String = server.version + def instructions: Option[String] = server.instructions + def showJsonSchemaMetadata: Boolean = server.showJsonSchemaMetadata + def originCheck: OriginCheck = server.originCheck + def prompts: List[ServerPrompt[F]] = server.prompts + def resources: List[ServerResource[F]] = server.resources + def resourceTemplates: List[ServerResourceTemplate[F]] = server.resourceTemplates + def completion: Option[CompletionHandler[F]] = server.completion + def loggingLevel: Option[SetLoggingLevelHandler[F]] = server.loggingLevel + def subscriptions: Option[ResourceSubscriptions[F]] = server.subscriptions + + def securityInput: EndpointInput[S] = server.securityInput + def errorOutput: EndpointOutput[E] = server.errorOutput + def securityLogic: MonadError[F] => S => F[Either[E, P]] = server.securityLogic + + def tools: List[ServerTool[?, ?, F, SecuredStreamingServerContext[F, P]]] = server.tools ++ streamingTools + + def name(value: String): SecuredStreamingMcpServer[F, S, E, P] = + copy(server = server.name(value)) + + def version(value: String): SecuredStreamingMcpServer[F, S, E, P] = + copy(server = server.version(value)) + + def instructions(value: String): SecuredStreamingMcpServer[F, S, E, P] = + copy(server = server.instructions(value)) + + def withJsonSchemaMetadata(value: Boolean): SecuredStreamingMcpServer[F, S, E, P] = + copy(server = server.withJsonSchemaMetadata(value)) + + def withOriginCheck(value: OriginCheck): SecuredStreamingMcpServer[F, S, E, P] = + copy(server = server.withOriginCheck(value)) + + def addTool(tool: ServerTool[?, ?, F, SecuredServerContext[F, P]]): SecuredStreamingMcpServer[F, S, E, P] = + copy(server = server.addTool(tool)) + + def addTools(tools: ServerTool[?, ?, F, SecuredServerContext[F, P]]*): SecuredStreamingMcpServer[F, S, E, P] = + copy(server = server.addTools(tools*)) + + def addStreamingTool(tool: ServerTool[?, ?, F, SecuredStreamingServerContext[F, P]]): SecuredStreamingMcpServer[F, S, E, P] = + copy(streamingTools = streamingTools :+ tool) + + def addStreamingTools(tools: ServerTool[?, ?, F, SecuredStreamingServerContext[F, P]]*): SecuredStreamingMcpServer[F, S, E, P] = + copy(streamingTools = this.streamingTools ++ tools) + + def addPrompt(prompt: ServerPrompt[F]): SecuredStreamingMcpServer[F, S, E, P] = + copy(server = server.addPrompt(prompt)) + + def addPrompts(prompts: ServerPrompt[F]*): SecuredStreamingMcpServer[F, S, E, P] = + copy(server = server.addPrompts(prompts*)) + + def addResource(resource: ServerResource[F]): SecuredStreamingMcpServer[F, S, E, P] = + copy(server = server.addResource(resource)) + + def addResources(resources: ServerResource[F]*): SecuredStreamingMcpServer[F, S, E, P] = + copy(server = server.addResources(resources*)) + + def addResourceTemplate(resourceTemplate: ServerResourceTemplate[F]): SecuredStreamingMcpServer[F, S, E, P] = + copy(server = server.addResourceTemplate(resourceTemplate)) + + def addResourceTemplates(resourceTemplates: ServerResourceTemplate[F]*): SecuredStreamingMcpServer[F, S, E, P] = + copy(server = server.addResourceTemplates(resourceTemplates*)) + + def withCompletion(handler: CompletionHandler[F]): SecuredStreamingMcpServer[F, S, E, P] = + copy(server = server.withCompletion(handler)) + + def withLoggingLevel(handler: SetLoggingLevelHandler[F]): SecuredStreamingMcpServer[F, S, E, P] = + copy(server = server.withLoggingLevel(handler)) + + def withSubscriptions(handler: ResourceSubscriptions[F]): SecuredStreamingMcpServer[F, S, E, P] = + copy(server = server.withSubscriptions(handler)) diff --git a/server/src/main/scala/chimp/server/ServerContext.scala b/server/src/main/scala/chimp/server/ServerContext.scala index 5ddd46e..ebcfa60 100644 --- a/server/src/main/scala/chimp/server/ServerContext.scala +++ b/server/src/main/scala/chimp/server/ServerContext.scala @@ -40,3 +40,19 @@ private[server] final class SinkStreamingServerContext[F[_]](sink: OutboundSink[ sink.send( JSONRPCMessage.Notification(method = "notifications/message", params = Some(LoggingMessageParams(level, data, logger).asJson)) ) + +/** A context which gives both the principal that the security logic of a [[SecuredMcpServer]] made from the request, and the streaming + * capabilities of a [[StreamingServerContext]]. + */ +trait SecuredStreamingServerContext[F[_], +P] extends SecuredServerContext[F, P] with StreamingServerContext[F] + +private[server] final class SinkSecuredStreamingServerContext[F[_], P]( + sink: OutboundSink[F], + progressToken: Option[ProgressToken], + val principal: P +)(using MonadError[F]) + extends SecuredStreamingServerContext[F, P]: + private val delegate = SinkStreamingServerContext[F](sink, progressToken) + def reportProgress(progress: Double, total: Option[Double] = None, message: Option[String] = None): F[Unit] = + delegate.reportProgress(progress, total, message) + def log(level: LoggingLevel, data: Json, logger: Option[String] = None): F[Unit] = delegate.log(level, data, logger) diff --git a/server/src/main/scala/chimp/server/Tool.scala b/server/src/main/scala/chimp/server/Tool.scala index 6ee06f8..ebdd65c 100644 --- a/server/src/main/scala/chimp/server/Tool.scala +++ b/server/src/main/scala/chimp/server/Tool.scala @@ -121,6 +121,22 @@ case class Tool[I, O]( ): ServerTool[I, O, F, StreamingServerContext[F]] = ServerTool(name, description, inputSchema, inputDecoder, outputSchema, annotations, logic) + /** Attaches effectful logic with access to the principal and the [[StreamingServerContext]]; usable only on a secured streaming + * server. + */ + def securedStreamingServerLogic[F[_], P]( + logic: (I, P, StreamingServerContext[F], Seq[Header]) => F[ToolResult[O]] + ): ServerTool[I, O, F, SecuredStreamingServerContext[F, P]] = + ServerTool( + name, + description, + inputSchema, + inputDecoder, + outputSchema, + annotations, + (input, context, headers) => logic(input, context.principal, context, headers) + ) + /** Attaches synchronous logic that also receives the request headers. */ def handleWithHeaders(logic: (I, Seq[Header]) => ToolResult[O]): ServerTool[I, O, Identity, ServerContext[Identity]] = ServerTool(name, description, inputSchema, inputDecoder, outputSchema, annotations, (i, _, headers) => logic(i, headers)) diff --git a/server/src/main/scala/chimp/server/transport/JsonRpcHttpResponses.scala b/server/src/main/scala/chimp/server/transport/JsonRpcHttpResponses.scala new file mode 100644 index 0000000..d00bd40 --- /dev/null +++ b/server/src/main/scala/chimp/server/transport/JsonRpcHttpResponses.scala @@ -0,0 +1,33 @@ +package chimp.server.transport + +import chimp.server.{McpResponse, OriginCheck, OutboundSink} +import io.circe.Json +import sttp.model.{Header, HeaderNames, StatusCode} +import sttp.monad.MonadError +import sttp.monad.syntax.* + +/** True if the request's `Host` and `Origin` headers pass the given check. Shared by every HTTP transport, secured or not, streaming or + * not. + */ +private[transport] def originAllowed(originCheck: OriginCheck, headers: Seq[Header]): Boolean = + val host = headers.find(_.name.equalsIgnoreCase(HeaderNames.Host)).map(_.value) + val origin = headers.find(_.name.equalsIgnoreCase(HeaderNames.Origin)).map(_.value) + originCheck.validate(host, origin) + +/** Runs `handle` if the origin check passes, otherwise gives a forbidden response, without running `handle`. */ +private[transport] def respondToJsonRpc[F[_]](originCheck: OriginCheck, headers: Seq[Header])(handle: => F[McpResponse])(using + m: MonadError[F] +): F[(StatusCode, Option[Json])] = + if !originAllowed(originCheck, headers) then m.unit((StatusCode.Forbidden, None)) + else handle.map(response => (response.statusCode, response.body)) + +/** The same as [[respondToJsonRpc]], but the response is an event stream produced by the given [[StreamingBackend]], rather than a single + * JSON body. + */ +private[transport] def respondWithEventStream[F[_], Caps]( + originCheck: OriginCheck, + headers: Seq[Header], + backend: StreamingBackend[F, Caps] +)(handle: OutboundSink[F] => F[Option[Json]])(using m: MonadError[F]): F[(StatusCode, backend.EventStream)] = + if !originAllowed(originCheck, headers) then m.unit((StatusCode.Forbidden, backend.emptyStream)) + else backend.eventStream(handle).map(events => (StatusCode.Ok, events)) diff --git a/server/src/main/scala/chimp/server/transport/SecuredServerHttpTransport.scala b/server/src/main/scala/chimp/server/transport/SecuredServerHttpTransport.scala index 0d43250..ebc452f 100644 --- a/server/src/main/scala/chimp/server/transport/SecuredServerHttpTransport.scala +++ b/server/src/main/scala/chimp/server/transport/SecuredServerHttpTransport.scala @@ -2,7 +2,7 @@ package chimp.server.transport import chimp.server.* import io.circe.Json -import sttp.model.{Header, HeaderNames, StatusCode} +import sttp.model.Header import sttp.monad.MonadError import sttp.monad.syntax.* import sttp.tapir.* @@ -34,12 +34,8 @@ final case class SecuredServerHttpTransport[F[_], S, E, P](path: List[String]): me => { (principal: P) => (input: (Seq[Header], Json)) => val (headers, json) = input given MonadError[F] = me - val host = headers.find(_.name.equalsIgnoreCase(HeaderNames.Host)).map(_.value) - val origin = headers.find(_.name.equalsIgnoreCase(HeaderNames.Origin)).map(_.value) - if !server.originCheck.validate(host, origin) then me.unit(Right((StatusCode.Forbidden, None))) - else - handler - .handleJsonRpc(json, headers, _ => SecuredServerContext[F, P](principal)) - .map(response => Right((response.statusCode, response.body))) + respondToJsonRpc(server.originCheck, headers) { + handler.handleJsonRpc(json, headers, _ => SecuredServerContext[F, P](principal)) + }.map(Right(_)) } ) diff --git a/server/src/main/scala/chimp/server/transport/SecuredServerStreamingHttpTransport.scala b/server/src/main/scala/chimp/server/transport/SecuredServerStreamingHttpTransport.scala new file mode 100644 index 0000000..46fe88f --- /dev/null +++ b/server/src/main/scala/chimp/server/transport/SecuredServerStreamingHttpTransport.scala @@ -0,0 +1,49 @@ +package chimp.server.transport + +import chimp.protocol.ProgressToken +import chimp.server.* +import io.circe.Json +import sttp.model.Header +import sttp.monad.MonadError +import sttp.monad.syntax.* +import sttp.tapir.* +import sttp.tapir.json.circe.* +import sttp.tapir.server.ServerEndpoint + +/** Implementation of bidirectional MCP server using Streamable HTTP, protected by the security logic of a [[SecuredMcpServer]]. The + * security logic runs before any MCP message is handled. If it gives a rejection, the server sends the error output and no tool logic + * runs. If it gives a principal, the principal goes to the tool logic, together with the [[StreamingServerContext]]. + * + * Wraps a [[StreamingBackend]] instead of extending it, so that the same backend instance - for example an + * `chimp.server.ox.OxServerHttpTransport` - serves both a plain [[StreamingMcpServer]] and a [[SecuredStreamingMcpServer]]. + * + * @param path + * The MCP endpoint path. + * @param backend + * The streaming machinery for the effect type `F` and the streaming capability `Caps`. + */ +final case class SecuredServerStreamingHttpTransport[F[_], Caps, S, E, P](path: List[String], backend: StreamingBackend[F, Caps]): + def serve(server: SecuredStreamingMcpServer[F, S, E, P]): ServerEndpoint[Caps, F] = + val handler = new McpHandler[F, SecuredStreamingServerContext[F, P]](server) + val mcpEndpoint = endpoint.post + .securityIn(server.securityInput) + .in(path.foldLeft(emptyInput)((inputSoFar, pathComponent) => inputSoFar / pathComponent)) + .in(extractFromRequest(_.headers)) + .in(jsonBody[Json]) + .errorOut(server.errorOutput) + .out(statusCode) + .out(backend.sseBody) + + ServerEndpoint( + mcpEndpoint, + server.securityLogic, + me => { (principal: P) => (input: (Seq[Header], Json)) => + val (headers, json) = input + given MonadError[F] = me + respondWithEventStream(server.originCheck, headers, backend) { sink => + val makeContext: Option[ProgressToken] => SecuredStreamingServerContext[F, P] = + token => SinkSecuredStreamingServerContext(sink, token, principal) + handler.handleJsonRpc(json, headers, makeContext).map(_.body) + }.map(Right(_)) + } + ) diff --git a/server/src/main/scala/chimp/server/transport/ServerHttpTransport.scala b/server/src/main/scala/chimp/server/transport/ServerHttpTransport.scala index 1b91bda..44be5a5 100644 --- a/server/src/main/scala/chimp/server/transport/ServerHttpTransport.scala +++ b/server/src/main/scala/chimp/server/transport/ServerHttpTransport.scala @@ -2,7 +2,7 @@ package chimp.server.transport import chimp.server.* import io.circe.Json -import sttp.model.{Header, HeaderNames, StatusCode} +import sttp.model.Header import sttp.monad.MonadError import sttp.monad.syntax.* import sttp.tapir.* @@ -30,12 +30,6 @@ final case class ServerHttpTransport[F[_]](path: List[String]) extends ServerTra me => { (input: (Seq[Header], Json)) => val (headers, json) = input given MonadError[F] = me - val host = headers.find(_.name.equalsIgnoreCase(HeaderNames.Host)).map(_.value) - val origin = headers.find(_.name.equalsIgnoreCase(HeaderNames.Origin)).map(_.value) - if !server.originCheck.validate(host, origin) then me.unit(Right((StatusCode.Forbidden, None))) - else - handler - .handleJsonRpc(json, headers) - .map(response => Right((response.statusCode, response.body))) + respondToJsonRpc(server.originCheck, headers)(handler.handleJsonRpc(json, headers)).map(Right(_)) } ) diff --git a/server/src/main/scala/chimp/server/transport/ServerStreamingHttpTransport.scala b/server/src/main/scala/chimp/server/transport/ServerStreamingHttpTransport.scala index 4dae664..bea642e 100644 --- a/server/src/main/scala/chimp/server/transport/ServerStreamingHttpTransport.scala +++ b/server/src/main/scala/chimp/server/transport/ServerStreamingHttpTransport.scala @@ -4,30 +4,38 @@ import chimp.protocol.ProgressToken import chimp.server.* import io.circe.Json import sttp.capabilities.Streams -import sttp.model.{Header, HeaderNames, StatusCode} +import sttp.model.Header import sttp.monad.MonadError import sttp.monad.syntax.* import sttp.tapir.* import sttp.tapir.json.circe.* import sttp.tapir.server.ServerEndpoint -/** Abstract base for bidirectional MCP server using Streamable HTTP. Responds to JSON-RPC messages from an MCP client with a - * Server-Sent-Event stream. Messages in the stream are interleaved with the final response on that stream. - * - * The extra type parameter `S` carries the streaming capability evidence required by the Tapir [[sttp.tapir.server.ServerEndpoint]] to - * produce asynchronous stream of Server-Sent Events as response. +/** The streaming machinery which a bidirectional MCP server transport needs from an effect backend: the streaming capability evidence, + * the codec for a Server-Sent-Event body, and a way to turn a [[chimp.server.OutboundSink]] into a stream of events. One instance is + * shared between the unsecured [[ServerStreamingHttpTransport]] and [[SecuredServerStreamingHttpTransport]]. * - * @param path - * The MCP endpoint path. + * @tparam Caps + * The streaming capability evidence required by the Tapir [[sttp.tapir.server.ServerEndpoint]] to produce an asynchronous stream of + * Server-Sent Events as response. */ -abstract class ServerStreamingHttpTransport[F[_], S](path: List[String]) extends StreamingServerTransport[F, ServerEndpoint[S, F]]: - val streams: Streams[S] +trait StreamingBackend[F[_], Caps]: + val streams: Streams[Caps] type EventStream - def sseBody: StreamBodyIO[streams.BinaryStream, EventStream, S] + def sseBody: StreamBodyIO[streams.BinaryStream, EventStream, Caps] def emptyStream: EventStream def eventStream(handle: OutboundSink[F] => F[Option[Json]]): F[EventStream] - final def serve(server: StreamingMcpServer[F]): ServerEndpoint[S, F] = +/** Abstract base for bidirectional MCP server using Streamable HTTP. Responds to JSON-RPC messages from an MCP client with a + * Server-Sent-Event stream. Messages in the stream are interleaved with the final response on that stream. + * + * @param path + * The MCP endpoint path. + */ +abstract class ServerStreamingHttpTransport[F[_], Caps](path: List[String]) + extends StreamingBackend[F, Caps] + with StreamingServerTransport[F, ServerEndpoint[Caps, F]]: + final def serve(server: StreamingMcpServer[F]): ServerEndpoint[Caps, F] = val handler = new McpHandler[F, StreamingServerContext[F]](server) val endpoint = infallibleEndpoint.post .in(path.foldLeft(emptyInput)((inputSoFar, pathComponent) => inputSoFar / pathComponent)) @@ -41,14 +49,10 @@ abstract class ServerStreamingHttpTransport[F[_], S](path: List[String]) extends me => { (input: (Seq[Header], Json)) => val (headers, json) = input given MonadError[F] = me - val host = headers.find(_.name.equalsIgnoreCase(HeaderNames.Host)).map(_.value) - val origin = headers.find(_.name.equalsIgnoreCase(HeaderNames.Origin)).map(_.value) - if !server.originCheck.validate(host, origin) then me.unit(Right((StatusCode.Forbidden, emptyStream))) - else - eventStream { sink => - val makeContext: Option[ProgressToken] => StreamingServerContext[F] = - token => SinkStreamingServerContext(sink, token) - handler.handleJsonRpc(json, headers, makeContext).map(_.body) - }.map(events => Right((StatusCode.Ok, events))) + respondWithEventStream(server.originCheck, headers, this) { sink => + val makeContext: Option[ProgressToken] => StreamingServerContext[F] = + token => SinkStreamingServerContext(sink, token) + handler.handleJsonRpc(json, headers, makeContext).map(_.body) + }.map(Right(_)) } ) diff --git a/server/src/test/scala/chimp/server/SecuredMcpServerStreamingTests.scala b/server/src/test/scala/chimp/server/SecuredMcpServerStreamingTests.scala new file mode 100644 index 0000000..6fcfb08 --- /dev/null +++ b/server/src/test/scala/chimp/server/SecuredMcpServerStreamingTests.scala @@ -0,0 +1,60 @@ +package chimp.server + +import chimp.client.BidirectionalMcpClient +import chimp.client.notifications.ServerNotification +import chimp.protocol.* +import io.circe.{Codec, Json} +import org.scalatest.Assertion +import org.scalatest.flatspec.AsyncFlatSpec +import org.scalatest.matchers.should.Matchers +import sttp.model.StatusCode +import sttp.monad.syntax.* +import sttp.tapir.* + +import java.util.concurrent.ConcurrentLinkedQueue +import scala.concurrent.Future +import scala.jdk.CollectionConverters.* + +trait SecuredMcpServerStreamingTests[F[_]] extends AsyncFlatSpec with Matchers: + this: ToFuture[F] => + + protected case class User(email: String) + + protected val validToken = "s3cret" + + protected def withSecuredStreamingServer(server: SecuredStreamingMcpServer[F, String, String, User])( + test: BidirectionalMcpClient[F] => F[Assertion] + ): Future[Assertion] + + private case class NoInput() derives Codec, Schema + + private def securityLogic(token: String): Either[String, User] = + if token == validToken then Right(User("employee@example.com")) else Left("Invalid token") + + protected def securedStreamingServer: SecuredStreamingMcpServer[F, String, String, User] = + McpServer[F]() + .serverSecurityLogicPure(auth.bearer[String](), statusCode(StatusCode.Unauthorized).and(stringBody))(securityLogic) + .withLoggingLevel(_ => monad.unit(())) + .streaming + .addStreamingTool( + tool("whoAmI") + .description("Logs the caller's email, then returns it") + .input[NoInput] + .securedStreamingServerLogic[F, User]: (_, user, ctx, _) => + ctx.log(LoggingLevel.Info, Json.fromString(user.email)).map(_ => ToolResult.text(user.email)) + ) + + "a secured streaming MCP server" should "give the principal to a streaming tool, and deliver its log notifications" in + withSecuredStreamingServer(securedStreamingServer): client => + val messages = ConcurrentLinkedQueue[Json]() + val listener: ServerNotification => F[Unit] = { + case ServerNotification.LoggingMessage(params) => messages.add(params.data); monad.unit(()) + case _ => monad.unit(()) + } + client + .onServerNotification(notification => listener(notification)) + .flatMap(_ => client.callTool("whoAmI", Json.obj())) + .flatMap: result => + waitUntil(messages.size >= 1).map: _ => + result.content shouldBe List(ToolContent.Text("text", "employee@example.com")) + messages.asScala.toList shouldBe List(Json.fromString("employee@example.com")) From ab58291a5a2a34bd1b83ddd85d241563fec27942 Mon Sep 17 00:00:00 2001 From: Maxime Perocheau Date: Wed, 9 Sep 2026 09:29:28 -0700 Subject: [PATCH 3/5] Fold StreamingBackend into ServerStreamingHttpTransport to not break bin-compat --- docs/server/transport.md | 2 +- .../main/scala/chimp/server/McpServer.scala | 4 +-- server/src/main/scala/chimp/server/Tool.scala | 3 +-- .../transport/JsonRpcHttpResponses.scala | 6 ++--- .../SecuredServerStreamingHttpTransport.scala | 12 ++++++--- .../ServerStreamingHttpTransport.scala | 26 +++++++------------ 6 files changed, 25 insertions(+), 28 deletions(-) diff --git a/docs/server/transport.md b/docs/server/transport.md index b7b191a..1bd27f2 100644 --- a/docs/server/transport.md +++ b/docs/server/transport.md @@ -184,4 +184,4 @@ object McpServerWithPrincipalAndStreaming: NettySyncServer().port(8080).addEndpoint(securedEndpoint).startAndWait() ``` -`SecuredServerStreamingHttpTransport` wraps a `StreamingBackend` rather than extending it, so the same effect-specific backend instance - `OxServerHttpTransport`, `PekkoServerHttpTransport`, `ZioServerHttpTransport` - serves both a plain `StreamingMcpServer` and a `SecuredStreamingMcpServer`. +`SecuredServerStreamingHttpTransport` takes an existing `ServerStreamingHttpTransport` as its source of streaming machinery, rather than extending it, so the same effect-specific backend instance - `OxServerHttpTransport`, `PekkoServerHttpTransport`, `ZioServerHttpTransport` - serves both a plain `StreamingMcpServer` and a `SecuredStreamingMcpServer`. diff --git a/server/src/main/scala/chimp/server/McpServer.scala b/server/src/main/scala/chimp/server/McpServer.scala index 3fe3ff0..49f380c 100644 --- a/server/src/main/scala/chimp/server/McpServer.scala +++ b/server/src/main/scala/chimp/server/McpServer.scala @@ -275,8 +275,8 @@ case class SecuredMcpServer[F[_], S, E, P]( def streaming: SecuredStreamingMcpServer[F, S, E, P] = SecuredStreamingMcpServer(this) -/** A [[SecuredMcpServer]] which also accepts streaming tools, which are given a [[SecuredStreamingServerContext]] combining the - * principal with the [[StreamingServerContext]]. Tools of the initial secured server, which need only the principal, are kept. +/** A [[SecuredMcpServer]] which also accepts streaming tools, which are given a [[SecuredStreamingServerContext]] combining the principal + * with the [[StreamingServerContext]]. Tools of the initial secured server, which need only the principal, are kept. */ case class SecuredStreamingMcpServer[F[_], S, E, P]( server: SecuredMcpServer[F, S, E, P], diff --git a/server/src/main/scala/chimp/server/Tool.scala b/server/src/main/scala/chimp/server/Tool.scala index ebdd65c..e6aecd3 100644 --- a/server/src/main/scala/chimp/server/Tool.scala +++ b/server/src/main/scala/chimp/server/Tool.scala @@ -121,8 +121,7 @@ case class Tool[I, O]( ): ServerTool[I, O, F, StreamingServerContext[F]] = ServerTool(name, description, inputSchema, inputDecoder, outputSchema, annotations, logic) - /** Attaches effectful logic with access to the principal and the [[StreamingServerContext]]; usable only on a secured streaming - * server. + /** Attaches effectful logic with access to the principal and the [[StreamingServerContext]]; usable only on a secured streaming server. */ def securedStreamingServerLogic[F[_], P]( logic: (I, P, StreamingServerContext[F], Seq[Header]) => F[ToolResult[O]] diff --git a/server/src/main/scala/chimp/server/transport/JsonRpcHttpResponses.scala b/server/src/main/scala/chimp/server/transport/JsonRpcHttpResponses.scala index d00bd40..3727dbe 100644 --- a/server/src/main/scala/chimp/server/transport/JsonRpcHttpResponses.scala +++ b/server/src/main/scala/chimp/server/transport/JsonRpcHttpResponses.scala @@ -21,13 +21,13 @@ private[transport] def respondToJsonRpc[F[_]](originCheck: OriginCheck, headers: if !originAllowed(originCheck, headers) then m.unit((StatusCode.Forbidden, None)) else handle.map(response => (response.statusCode, response.body)) -/** The same as [[respondToJsonRpc]], but the response is an event stream produced by the given [[StreamingBackend]], rather than a single - * JSON body. +/** The same as [[respondToJsonRpc]], but the response is an event stream produced by the given [[ServerStreamingHttpTransport]]'s streaming + * machinery, rather than a single JSON body. */ private[transport] def respondWithEventStream[F[_], Caps]( originCheck: OriginCheck, headers: Seq[Header], - backend: StreamingBackend[F, Caps] + backend: ServerStreamingHttpTransport[F, Caps] )(handle: OutboundSink[F] => F[Option[Json]])(using m: MonadError[F]): F[(StatusCode, backend.EventStream)] = if !originAllowed(originCheck, headers) then m.unit((StatusCode.Forbidden, backend.emptyStream)) else backend.eventStream(handle).map(events => (StatusCode.Ok, events)) diff --git a/server/src/main/scala/chimp/server/transport/SecuredServerStreamingHttpTransport.scala b/server/src/main/scala/chimp/server/transport/SecuredServerStreamingHttpTransport.scala index 46fe88f..4a6faa7 100644 --- a/server/src/main/scala/chimp/server/transport/SecuredServerStreamingHttpTransport.scala +++ b/server/src/main/scala/chimp/server/transport/SecuredServerStreamingHttpTransport.scala @@ -14,15 +14,19 @@ import sttp.tapir.server.ServerEndpoint * security logic runs before any MCP message is handled. If it gives a rejection, the server sends the error output and no tool logic * runs. If it gives a principal, the principal goes to the tool logic, together with the [[StreamingServerContext]]. * - * Wraps a [[StreamingBackend]] instead of extending it, so that the same backend instance - for example an - * `chimp.server.ox.OxServerHttpTransport` - serves both a plain [[StreamingMcpServer]] and a [[SecuredStreamingMcpServer]]. + * Takes an existing [[ServerStreamingHttpTransport]] as the source of the streaming machinery for the effect backend, rather than + * extending it, so that the same backend instance - for example an `chimp.server.ox.OxServerHttpTransport` - serves both a plain + * [[StreamingMcpServer]] and a [[SecuredStreamingMcpServer]]. Its own `path` plays no part here; only its streaming machinery is used. * * @param path * The MCP endpoint path. * @param backend - * The streaming machinery for the effect type `F` and the streaming capability `Caps`. + * Supplies the streaming machinery for the effect type `F` and the streaming capability `Caps`. */ -final case class SecuredServerStreamingHttpTransport[F[_], Caps, S, E, P](path: List[String], backend: StreamingBackend[F, Caps]): +final case class SecuredServerStreamingHttpTransport[F[_], Caps, S, E, P]( + path: List[String], + backend: ServerStreamingHttpTransport[F, Caps] +): def serve(server: SecuredStreamingMcpServer[F, S, E, P]): ServerEndpoint[Caps, F] = val handler = new McpHandler[F, SecuredStreamingServerContext[F, P]](server) val mcpEndpoint = endpoint.post diff --git a/server/src/main/scala/chimp/server/transport/ServerStreamingHttpTransport.scala b/server/src/main/scala/chimp/server/transport/ServerStreamingHttpTransport.scala index bea642e..8ae7503 100644 --- a/server/src/main/scala/chimp/server/transport/ServerStreamingHttpTransport.scala +++ b/server/src/main/scala/chimp/server/transport/ServerStreamingHttpTransport.scala @@ -11,30 +11,24 @@ import sttp.tapir.* import sttp.tapir.json.circe.* import sttp.tapir.server.ServerEndpoint -/** The streaming machinery which a bidirectional MCP server transport needs from an effect backend: the streaming capability evidence, - * the codec for a Server-Sent-Event body, and a way to turn a [[chimp.server.OutboundSink]] into a stream of events. One instance is - * shared between the unsecured [[ServerStreamingHttpTransport]] and [[SecuredServerStreamingHttpTransport]]. +/** Abstract base for bidirectional MCP server using Streamable HTTP. Responds to JSON-RPC messages from an MCP client with a + * Server-Sent-Event stream. Messages in the stream are interleaved with the final response on that stream. + * + * The extra type parameter `Caps` carries the streaming capability evidence required by the Tapir [[sttp.tapir.server.ServerEndpoint]] to + * produce asynchronous stream of Server-Sent Events as response. An instance also serves as the streaming machinery which + * [[SecuredServerStreamingHttpTransport]] needs from the same effect backend - its `path` plays no part in that, so the same instance, or + * another one for the same `F`/`Caps`, can back both a plain [[StreamingMcpServer]] and a [[SecuredStreamingMcpServer]]. * - * @tparam Caps - * The streaming capability evidence required by the Tapir [[sttp.tapir.server.ServerEndpoint]] to produce an asynchronous stream of - * Server-Sent Events as response. + * @param path + * The MCP endpoint path. */ -trait StreamingBackend[F[_], Caps]: +abstract class ServerStreamingHttpTransport[F[_], Caps](path: List[String]) extends StreamingServerTransport[F, ServerEndpoint[Caps, F]]: val streams: Streams[Caps] type EventStream def sseBody: StreamBodyIO[streams.BinaryStream, EventStream, Caps] def emptyStream: EventStream def eventStream(handle: OutboundSink[F] => F[Option[Json]]): F[EventStream] -/** Abstract base for bidirectional MCP server using Streamable HTTP. Responds to JSON-RPC messages from an MCP client with a - * Server-Sent-Event stream. Messages in the stream are interleaved with the final response on that stream. - * - * @param path - * The MCP endpoint path. - */ -abstract class ServerStreamingHttpTransport[F[_], Caps](path: List[String]) - extends StreamingBackend[F, Caps] - with StreamingServerTransport[F, ServerEndpoint[Caps, F]]: final def serve(server: StreamingMcpServer[F]): ServerEndpoint[Caps, F] = val handler = new McpHandler[F, StreamingServerContext[F]](server) val endpoint = infallibleEndpoint.post From 000fbc8d976e33ba1e92bb29950cc19a10f5e2c8 Mon Sep 17 00:00:00 2001 From: Maxime Perocheau Date: Thu, 10 Sep 2026 14:58:23 -0700 Subject: [PATCH 4/5] Add more tests + doc --- docs/server/tools.md | 3 +- .../chimp/server/ox/OxMcpServerHttpSpec.scala | 5 +- .../server/pekko/PekkoMcpServerHttpSpec.scala | 5 +- .../server/zio/ZioMcpServerHttpSpec.scala | 5 +- .../server/SecuredHttpMcpServerSpec.scala | 20 ++++++ .../SecuredMcpServerStreamingTests.scala | 72 ++++++++++++------- 6 files changed, 79 insertions(+), 31 deletions(-) diff --git a/docs/server/tools.md b/docs/server/tools.md index 9ac1ee1..656d7dc 100644 --- a/docs/server/tools.md +++ b/docs/server/tools.md @@ -8,8 +8,9 @@ - `handle` — synchronous logic from input to `ToolResult`. - `handleWithHeaders` — synchronous logic that also receives the request headers. - `serverLogic` — effectful logic, with the request headers. + - `handleSecured` (synchronous) or `securedServerLogic` (effectful) — logic that also receives the principal made by the server's security logic; usable only on a secured server — see [transport security](transport.md). - A tool that pushes to the client while running (progress, logging) instead uses `streamingServerLogic` — see [server capabilities](capabilities.md). + A tool that pushes to the client while running (progress, logging) instead uses `streamingServerLogic` — see [server capabilities](capabilities.md). A secured, streaming tool uses `securedStreamingServerLogic`, which gives both the principal and the streaming context — see [transport security](transport.md). - Assemble tools into an `McpServer` and call `.endpoint(path)` to create a Tapir endpoint. ```scala mdoc:compile-only diff --git a/server-streaming/server-ox/src/test/scala/chimp/server/ox/OxMcpServerHttpSpec.scala b/server-streaming/server-ox/src/test/scala/chimp/server/ox/OxMcpServerHttpSpec.scala index 2f946da..c1a0286 100644 --- a/server-streaming/server-ox/src/test/scala/chimp/server/ox/OxMcpServerHttpSpec.scala +++ b/server-streaming/server-ox/src/test/scala/chimp/server/ox/OxMcpServerHttpSpec.scala @@ -57,7 +57,8 @@ class OxMcpServerHttpSpec finally binding.stop() override protected def withSecuredStreamingServer( - server: SecuredStreamingMcpServer[Identity, String, String, User] + server: SecuredStreamingMcpServer[Identity, String, String, User], + token: String )(test: BidirectionalMcpClient[Identity] => Identity[Assertion]): Future[Assertion] = toFuture: supervised: @@ -72,7 +73,7 @@ class OxMcpServerHttpSpec uri"http://localhost:${binding.port}/mcp", ProtocolVersion.Latest, ClientTransport.defaultTimeout, - headers = List(Header.authorization("Bearer", validToken)) + headers = List(Header.authorization("Bearer", token)) ) try test(McpClient.bidirectional(transport, clientInfo)) finally transport.close() diff --git a/server-streaming/server-pekko/src/test/scala/chimp/server/pekko/PekkoMcpServerHttpSpec.scala b/server-streaming/server-pekko/src/test/scala/chimp/server/pekko/PekkoMcpServerHttpSpec.scala index 32895d9..2626ab8 100644 --- a/server-streaming/server-pekko/src/test/scala/chimp/server/pekko/PekkoMcpServerHttpSpec.scala +++ b/server-streaming/server-pekko/src/test/scala/chimp/server/pekko/PekkoMcpServerHttpSpec.scala @@ -54,7 +54,8 @@ class PekkoMcpServerHttpSpec .transform(_ => result) override protected def withSecuredStreamingServer( - server: SecuredStreamingMcpServer[Future, String, String, User] + server: SecuredStreamingMcpServer[Future, String, String, User], + token: String )(test: BidirectionalMcpClient[Future] => Future[Assertion]): Future[Assertion] = given ExecutionContext = actorSystem.dispatcher val endpoint = SecuredServerStreamingHttpTransport(List("mcp"), PekkoServerHttpTransport(List("mcp"))).serve(server) @@ -66,7 +67,7 @@ class PekkoMcpServerHttpSpec val transport = PekkoClientHttpTransport( backend, uri"http://localhost:${binding.localAddress.getPort}/mcp", - headers = List(Header.authorization("Bearer", validToken)) + headers = List(Header.authorization("Bearer", token)) ) McpClient .bidirectional(transport, clientInfo) diff --git a/server-streaming/server-zio/src/test/scala/chimp/server/zio/ZioMcpServerHttpSpec.scala b/server-streaming/server-zio/src/test/scala/chimp/server/zio/ZioMcpServerHttpSpec.scala index f4da5c0..8918783 100644 --- a/server-streaming/server-zio/src/test/scala/chimp/server/zio/ZioMcpServerHttpSpec.scala +++ b/server-streaming/server-zio/src/test/scala/chimp/server/zio/ZioMcpServerHttpSpec.scala @@ -50,7 +50,8 @@ class ZioMcpServerHttpSpec yield result).provideSome[Scope](Server.defaultWithPort(0)) override protected def withSecuredStreamingServer( - server: SecuredStreamingMcpServer[Task, String, String, User] + server: SecuredStreamingMcpServer[Task, String, String, User], + token: String )(test: BidirectionalMcpClient[Task] => Task[Assertion]): Future[Assertion] = toFuture: val routes = @@ -65,7 +66,7 @@ class ZioMcpServerHttpSpec uri"http://localhost:$port/mcp", ProtocolVersion.Latest, ClientTransport.defaultTimeout, - headers = List(Header.authorization("Bearer", validToken)) + headers = List(Header.authorization("Bearer", token)) ) .flatMap(transport => McpClient.bidirectional(transport, clientInfo)) .flatMap(client => test(client)) diff --git a/server/src/test/scala/chimp/server/SecuredHttpMcpServerSpec.scala b/server/src/test/scala/chimp/server/SecuredHttpMcpServerSpec.scala index 3154823..7afa100 100644 --- a/server/src/test/scala/chimp/server/SecuredHttpMcpServerSpec.scala +++ b/server/src/test/scala/chimp/server/SecuredHttpMcpServerSpec.scala @@ -39,11 +39,18 @@ class SecuredHttpMcpServerSpec extends AnyFlatSpec with Matchers: private def securityLogic(token: String): Either[String, User] = if token == validToken then Right(User("employee@example.com")) else Left("Invalid token") + private def securityLogicEffectful(token: String): Identity[Either[String, User]] = securityLogic(token) + private val securedServer = McpServer[Identity]() .addTool(echoTool) .serverSecurityLogicPure(auth.bearer[String](), statusCode(StatusCode.Unauthorized).and(stringBody))(securityLogic) .addTool(whoAmITool) + private val securedServerWithEffectfulSecurityLogic = McpServer[Identity]() + .addTool(echoTool) + .serverSecurityLogic(auth.bearer[String](), statusCode(StatusCode.Unauthorized).and(stringBody))(securityLogicEffectful) + .addTool(whoAmITool) + private val serverConfiguredAfterSecurity = McpServer[Identity]() .serverSecurityLogicPure(auth.bearer[String](), statusCode(StatusCode.Unauthorized).and(stringBody))(securityLogic) .name("secured-server") @@ -68,6 +75,7 @@ class SecuredHttpMcpServerSpec extends AnyFlatSpec with Matchers: private val securedEndpoint = securedServer.endpoint(List("mcp")) private val endpointConfiguredAfterSecurity = serverConfiguredAfterSecurity.endpoint(List("mcp")) + private val endpointWithEffectfulSecurityLogic = securedServerWithEffectfulSecurityLogic.endpoint(List("mcp")) "a secured MCP server" should "give the principal to the tool logic" in withServer(securedEndpoint): (port, backend) => withClient(port, backend, validToken): client => @@ -75,6 +83,18 @@ class SecuredHttpMcpServerSpec extends AnyFlatSpec with Matchers: result.isError shouldBe false result.content shouldBe List(ToolContent.Text("text", "hi employee@example.com")) + it should "give the principal to the tool logic when the security logic is effectful" in + withServer(endpointWithEffectfulSecurityLogic): (port, backend) => + withClient(port, backend, validToken): client => + val result = client.callTool("whoAmI", Json.obj("message" -> Json.fromString("hi"))) + result.isError shouldBe false + result.content shouldBe List(ToolContent.Text("text", "hi employee@example.com")) + + it should "reject an invalid security input with the error output when the security logic is effectful" in + withServer(endpointWithEffectfulSecurityLogic): (port, backend) => + val exception = intercept[McpAuthorizationException](withClient(port, backend, "wrong")(_ => ())) + exception.statusCode shouldBe StatusCode.Unauthorized.code + it should "also serve the tools which do not need the principal" in withServer(securedEndpoint): (port, backend) => withClient(port, backend, validToken): client => client.listTools().tools.map(_.name) should contain allOf ("echo", "whoAmI") diff --git a/server/src/test/scala/chimp/server/SecuredMcpServerStreamingTests.scala b/server/src/test/scala/chimp/server/SecuredMcpServerStreamingTests.scala index 6fcfb08..ece3242 100644 --- a/server/src/test/scala/chimp/server/SecuredMcpServerStreamingTests.scala +++ b/server/src/test/scala/chimp/server/SecuredMcpServerStreamingTests.scala @@ -1,10 +1,10 @@ package chimp.server -import chimp.client.BidirectionalMcpClient +import chimp.client.{BidirectionalMcpClient, McpAuthorizationException} import chimp.client.notifications.ServerNotification import chimp.protocol.* import io.circe.{Codec, Json} -import org.scalatest.Assertion +import org.scalatest.{Assertion, RecoverMethods} import org.scalatest.flatspec.AsyncFlatSpec import org.scalatest.matchers.should.Matchers import sttp.model.StatusCode @@ -15,14 +15,17 @@ import java.util.concurrent.ConcurrentLinkedQueue import scala.concurrent.Future import scala.jdk.CollectionConverters.* -trait SecuredMcpServerStreamingTests[F[_]] extends AsyncFlatSpec with Matchers: +trait SecuredMcpServerStreamingTests[F[_]] extends AsyncFlatSpec with Matchers with RecoverMethods: this: ToFuture[F] => protected case class User(email: String) protected val validToken = "s3cret" - protected def withSecuredStreamingServer(server: SecuredStreamingMcpServer[F, String, String, User])( + /** @param token + * The bearer token the client authenticates with; defaults to [[validToken]] so most tests need not pass it. + */ + protected def withSecuredStreamingServer(server: SecuredStreamingMcpServer[F, String, String, User], token: String = validToken)( test: BidirectionalMcpClient[F] => F[Assertion] ): Future[Assertion] @@ -31,30 +34,51 @@ trait SecuredMcpServerStreamingTests[F[_]] extends AsyncFlatSpec with Matchers: private def securityLogic(token: String): Either[String, User] = if token == validToken then Right(User("employee@example.com")) else Left("Invalid token") + private def securityLogicEffectful(token: String): F[Either[String, User]] = + monad.unit(securityLogic(token)) + + private def whoAmITool: ServerTool[NoInput, NoStructuredOutput, F, SecuredStreamingServerContext[F, User]] = + tool("whoAmI") + .description("Logs the caller's email, then returns it") + .input[NoInput] + .securedStreamingServerLogic[F, User]: (_, user, ctx, _) => + ctx.log(LoggingLevel.Info, Json.fromString(user.email)).map(_ => ToolResult.text(user.email)) + protected def securedStreamingServer: SecuredStreamingMcpServer[F, String, String, User] = McpServer[F]() .serverSecurityLogicPure(auth.bearer[String](), statusCode(StatusCode.Unauthorized).and(stringBody))(securityLogic) .withLoggingLevel(_ => monad.unit(())) .streaming - .addStreamingTool( - tool("whoAmI") - .description("Logs the caller's email, then returns it") - .input[NoInput] - .securedStreamingServerLogic[F, User]: (_, user, ctx, _) => - ctx.log(LoggingLevel.Info, Json.fromString(user.email)).map(_ => ToolResult.text(user.email)) - ) + .addStreamingTool(whoAmITool) + + protected def securedStreamingServerWithEffectfulSecurityLogic: SecuredStreamingMcpServer[F, String, String, User] = + McpServer[F]() + .serverSecurityLogic(auth.bearer[String](), statusCode(StatusCode.Unauthorized).and(stringBody))(securityLogicEffectful) + .withLoggingLevel(_ => monad.unit(())) + .streaming + .addStreamingTool(whoAmITool) + + private def assertWhoAmIDeliversPrincipal(client: BidirectionalMcpClient[F]): F[Assertion] = + val messages = ConcurrentLinkedQueue[Json]() + val listener: ServerNotification => F[Unit] = { + case ServerNotification.LoggingMessage(params) => messages.add(params.data); monad.unit(()) + case _ => monad.unit(()) + } + client + .onServerNotification(notification => listener(notification)) + .flatMap(_ => client.callTool("whoAmI", Json.obj())) + .flatMap: result => + waitUntil(messages.size >= 1).map: _ => + result.content shouldBe List(ToolContent.Text("text", "employee@example.com")) + messages.asScala.toList shouldBe List(Json.fromString("employee@example.com")) "a secured streaming MCP server" should "give the principal to a streaming tool, and deliver its log notifications" in - withSecuredStreamingServer(securedStreamingServer): client => - val messages = ConcurrentLinkedQueue[Json]() - val listener: ServerNotification => F[Unit] = { - case ServerNotification.LoggingMessage(params) => messages.add(params.data); monad.unit(()) - case _ => monad.unit(()) - } - client - .onServerNotification(notification => listener(notification)) - .flatMap(_ => client.callTool("whoAmI", Json.obj())) - .flatMap: result => - waitUntil(messages.size >= 1).map: _ => - result.content shouldBe List(ToolContent.Text("text", "employee@example.com")) - messages.asScala.toList shouldBe List(Json.fromString("employee@example.com")) + withSecuredStreamingServer(securedStreamingServer)(assertWhoAmIDeliversPrincipal) + + it should "give the principal to a streaming tool when the security logic is effectful" in + withSecuredStreamingServer(securedStreamingServerWithEffectfulSecurityLogic)(assertWhoAmIDeliversPrincipal) + + it should "reject an invalid security input with the error output, before any tool logic runs" in + recoverToExceptionIf[McpAuthorizationException] { + Future(withSecuredStreamingServer(securedStreamingServer, token = "wrong")(_ => monad.unit(succeed))).flatten + }.map(_.statusCode shouldBe StatusCode.Unauthorized.code) From 1591a2edfc28e9d1d757eff3fb674fa7355545ca Mon Sep 17 00:00:00 2001 From: Maxime Perocheau Date: Mon, 14 Sep 2026 13:49:57 -0700 Subject: [PATCH 5/5] Address review comments --- docs/server/transport.md | 2 +- .../chimp/server/ox/OxMcpServerHttpSpec.scala | 79 ++++++----- .../server/pekko/PekkoMcpServerHttpSpec.scala | 52 +++---- .../server/zio/ZioMcpServerHttpSpec.scala | 43 +++--- .../transport/JsonRpcHttpResponses.scala | 7 - .../server/SecuredHttpMcpServerSpec.scala | 129 ------------------ .../SecuredMcpServerStreamingTests.scala | 7 +- .../chimp/server/SecuredMcpServerTests.scala | 120 ++++++++++++++++ 8 files changed, 222 insertions(+), 217 deletions(-) delete mode 100644 server/src/test/scala/chimp/server/SecuredHttpMcpServerSpec.scala create mode 100644 server/src/test/scala/chimp/server/SecuredMcpServerTests.scala diff --git a/docs/server/transport.md b/docs/server/transport.md index 1bd27f2..c63f21a 100644 --- a/docs/server/transport.md +++ b/docs/server/transport.md @@ -93,7 +93,7 @@ import sttp.model.StatusCode import sttp.tapir.* import sttp.tapir.server.netty.sync.NettySyncServer -object McpServerWithPrependedSecurity: +object SecuredMcpServer: def main(args: Array[String]): Unit = val adder = tool("echo").input[String].handle(echo => ToolResult.text(echo)) val mcpEndpoint = McpServer(tools = List(adder)).endpoint(List("mcp")) diff --git a/server-streaming/server-ox/src/test/scala/chimp/server/ox/OxMcpServerHttpSpec.scala b/server-streaming/server-ox/src/test/scala/chimp/server/ox/OxMcpServerHttpSpec.scala index c1a0286..724be68 100644 --- a/server-streaming/server-ox/src/test/scala/chimp/server/ox/OxMcpServerHttpSpec.scala +++ b/server-streaming/server-ox/src/test/scala/chimp/server/ox/OxMcpServerHttpSpec.scala @@ -1,6 +1,6 @@ package chimp.server.ox -import chimp.client.transport.ClientTransport +import chimp.client.transport.{ClientHttpTransport, ClientTransport} import chimp.client.transport.ox.OxClientHttpTransport import chimp.client.{BidirectionalMcpClient, McpClient} import chimp.protocol.{Implementation, ProtocolVersion} @@ -9,24 +9,29 @@ import chimp.server.{ McpServer, McpServerStreamingTests, McpServerTests, + SecuredMcpServer, SecuredMcpServerStreamingTests, + SecuredMcpServerTests, SecuredStreamingMcpServer, StreamingMcpServer, SyncToFuture } import org.scalatest.Assertion -import ox.supervised -import sttp.client4.DefaultSyncBackend +import ox.{supervised, Ox} +import sttp.capabilities.WebSockets +import sttp.client4.{DefaultSyncBackend, SyncBackend} import sttp.model.Header import sttp.model.Uri.UriContext import sttp.shared.Identity -import sttp.tapir.server.netty.sync.NettySyncServer +import sttp.tapir.server.netty.sync.{NettySyncServer, OxStreams} +import sttp.tapir.server.ServerEndpoint import scala.concurrent.Future class OxMcpServerHttpSpec extends McpServerTests[Identity] with McpServerStreamingTests[Identity] + with SecuredMcpServerTests[Identity] with SecuredMcpServerStreamingTests[Identity] with SyncToFuture: private val clientInfo = Implementation("chimp-server-test", "0.0.1") @@ -37,45 +42,51 @@ class OxMcpServerHttpSpec override protected def withStreamingServer( server: StreamingMcpServer[Identity] )(test: BidirectionalMcpClient[Identity] => Identity[Assertion]): Future[Assertion] = - toFuture: - supervised: - val endpoint = OxServerHttpTransport(List("mcp")).serve(server) - val binding = NettySyncServer().port(0).addEndpoint(endpoint).start() - try - val backend = DefaultSyncBackend() - try - val transport = - OxClientHttpTransport( - backend, - uri"http://localhost:${binding.port}/mcp", - ProtocolVersion.Latest, - ClientTransport.defaultTimeout - ) - try test(McpClient.bidirectional(transport, clientInfo)) - finally transport.close() - finally backend.close() - finally binding.stop() + withHttpServer(OxServerHttpTransport(List("mcp")).serve(server)): (port, backend, ox) => + given Ox = ox + val transport = + OxClientHttpTransport(backend, uri"http://localhost:$port/mcp", ProtocolVersion.Latest, ClientTransport.defaultTimeout) + try test(McpClient.bidirectional(transport, clientInfo)) + finally transport.close() + + override protected def withSecuredServer( + server: SecuredMcpServer[Identity, String, String, User], + token: String + )(test: McpClient[Identity] => Identity[Assertion]): Future[Assertion] = + withHttpServer(server.endpoint(List("mcp"))): (port, backend, _) => + val transport = ClientHttpTransport[Identity]( + backend, + uri"http://localhost:$port/mcp", + headers = List(Header.authorization("Bearer", token)) + ) + try test(McpClient(transport, clientInfo)) + finally transport.close() override protected def withSecuredStreamingServer( server: SecuredStreamingMcpServer[Identity, String, String, User], token: String )(test: BidirectionalMcpClient[Identity] => Identity[Assertion]): Future[Assertion] = + val endpoint = SecuredServerStreamingHttpTransport(List("mcp"), OxServerHttpTransport(List("mcp"))).serve(server) + withHttpServer(endpoint): (port, backend, ox) => + given Ox = ox + val transport = OxClientHttpTransport( + backend, + uri"http://localhost:$port/mcp", + ProtocolVersion.Latest, + ClientTransport.defaultTimeout, + headers = List(Header.authorization("Bearer", token)) + ) + try test(McpClient.bidirectional(transport, clientInfo)) + finally transport.close() + + private def withHttpServer( + endpoint: ServerEndpoint[OxStreams & WebSockets, Identity] + )(test: (Int, SyncBackend, Ox) => Assertion): Future[Assertion] = toFuture: supervised: - val endpoint = SecuredServerStreamingHttpTransport(List("mcp"), OxServerHttpTransport(List("mcp"))).serve(server) val binding = NettySyncServer().port(0).addEndpoint(endpoint).start() try val backend = DefaultSyncBackend() - try - val transport = - OxClientHttpTransport( - backend, - uri"http://localhost:${binding.port}/mcp", - ProtocolVersion.Latest, - ClientTransport.defaultTimeout, - headers = List(Header.authorization("Bearer", token)) - ) - try test(McpClient.bidirectional(transport, clientInfo)) - finally transport.close() + try test(binding.port, backend, summon[Ox]) finally backend.close() finally binding.stop() diff --git a/server-streaming/server-pekko/src/test/scala/chimp/server/pekko/PekkoMcpServerHttpSpec.scala b/server-streaming/server-pekko/src/test/scala/chimp/server/pekko/PekkoMcpServerHttpSpec.scala index 2626ab8..f3c0944 100644 --- a/server-streaming/server-pekko/src/test/scala/chimp/server/pekko/PekkoMcpServerHttpSpec.scala +++ b/server-streaming/server-pekko/src/test/scala/chimp/server/pekko/PekkoMcpServerHttpSpec.scala @@ -8,16 +8,21 @@ import chimp.server.{ McpServer, McpServerStreamingTests, McpServerTests, + SecuredMcpServer, SecuredMcpServerStreamingTests, + SecuredMcpServerTests, SecuredStreamingMcpServer, StreamingMcpServer } import org.apache.pekko.http.scaladsl.Http import org.scalatest.Assertion +import sttp.capabilities.WebSockets +import sttp.capabilities.pekko.PekkoStreams import sttp.client4.pekkohttp.PekkoHttpBackend import sttp.model.Header import sttp.model.Uri.UriContext import sttp.tapir.server.pekkohttp.PekkoHttpServerInterpreter +import sttp.tapir.server.ServerEndpoint import scala.concurrent.duration.DurationInt import scala.concurrent.{ExecutionContext, Future} @@ -25,6 +30,7 @@ import scala.concurrent.{ExecutionContext, Future} class PekkoMcpServerHttpSpec extends McpServerTests[Future] with McpServerStreamingTests[Future] + with SecuredMcpServerTests[Future] with SecuredMcpServerStreamingTests[Future] with PekkoToFuture: private val clientInfo = Implementation("chimp-server-test", "0.0.1") @@ -35,43 +41,37 @@ class PekkoMcpServerHttpSpec override protected def withStreamingServer( server: StreamingMcpServer[Future] )(test: BidirectionalMcpClient[Future] => Future[Assertion]): Future[Assertion] = - given ExecutionContext = actorSystem.dispatcher - val endpoint = PekkoServerHttpTransport(List("mcp")).serve(server) - Http() - .newServerAt("localhost", 0) - .bind(PekkoHttpServerInterpreter().toRoute(endpoint)) - .flatMap: binding => - val backend = PekkoHttpBackend.usingActorSystem(actorSystem) - val transport = PekkoClientHttpTransport(backend, uri"http://localhost:${binding.localAddress.getPort}/mcp") - McpClient - .bidirectional(transport, clientInfo) - .flatMap(test) - .transformWith: result => - transport - .close() - .transformWith(_ => backend.close()) - .transformWith(_ => binding.terminate(5.seconds)) - .transform(_ => result) + withHttpServer(PekkoServerHttpTransport(List("mcp")).serve(server), Nil): transport => + McpClient.bidirectional(transport, clientInfo).flatMap(test) + + override protected def withSecuredServer( + server: SecuredMcpServer[Future, String, String, User], + token: String + )(test: McpClient[Future] => Future[Assertion]): Future[Assertion] = + withHttpServer(server.endpoint(List("mcp")), List(Header.authorization("Bearer", token))): transport => + McpClient(transport, clientInfo).flatMap(test) override protected def withSecuredStreamingServer( server: SecuredStreamingMcpServer[Future, String, String, User], token: String )(test: BidirectionalMcpClient[Future] => Future[Assertion]): Future[Assertion] = - given ExecutionContext = actorSystem.dispatcher val endpoint = SecuredServerStreamingHttpTransport(List("mcp"), PekkoServerHttpTransport(List("mcp"))).serve(server) + withHttpServer(endpoint, List(Header.authorization("Bearer", token))): transport => + McpClient.bidirectional(transport, clientInfo).flatMap(test) + + private def withHttpServer( + endpoint: ServerEndpoint[PekkoStreams & WebSockets, Future], + headers: List[Header] + )(test: PekkoClientHttpTransport => Future[Assertion]): Future[Assertion] = + given ExecutionContext = actorSystem.dispatcher Http() .newServerAt("localhost", 0) .bind(PekkoHttpServerInterpreter().toRoute(endpoint)) .flatMap: binding => val backend = PekkoHttpBackend.usingActorSystem(actorSystem) - val transport = PekkoClientHttpTransport( - backend, - uri"http://localhost:${binding.localAddress.getPort}/mcp", - headers = List(Header.authorization("Bearer", token)) - ) - McpClient - .bidirectional(transport, clientInfo) - .flatMap(test) + val transport = + PekkoClientHttpTransport(backend, uri"http://localhost:${binding.localAddress.getPort}/mcp", headers = headers) + test(transport) .transformWith: result => transport .close() diff --git a/server-streaming/server-zio/src/test/scala/chimp/server/zio/ZioMcpServerHttpSpec.scala b/server-streaming/server-zio/src/test/scala/chimp/server/zio/ZioMcpServerHttpSpec.scala index 8918783..8bd995a 100644 --- a/server-streaming/server-zio/src/test/scala/chimp/server/zio/ZioMcpServerHttpSpec.scala +++ b/server-streaming/server-zio/src/test/scala/chimp/server/zio/ZioMcpServerHttpSpec.scala @@ -9,15 +9,20 @@ import chimp.server.{ McpServer, McpServerStreamingTests, McpServerTests, + SecuredMcpServer, SecuredMcpServerStreamingTests, + SecuredMcpServerTests, SecuredStreamingMcpServer, StreamingMcpServer } import org.scalatest.Assertion +import sttp.capabilities.WebSockets +import sttp.capabilities.zio.ZioStreams import sttp.client4.* import sttp.client4.httpclient.zio.HttpClientZioBackend import sttp.model.Header import sttp.tapir.server.ziohttp.ZioHttpInterpreter +import sttp.tapir.server.ServerEndpoint import zio.http.Server import zio.{Scope, Task, ZIO} @@ -26,6 +31,7 @@ import scala.concurrent.Future class ZioMcpServerHttpSpec extends McpServerTests[Task] with McpServerStreamingTests[Task] + with SecuredMcpServerTests[Task] with SecuredMcpServerStreamingTests[Task] with ZioToFuture: private val clientInfo = Implementation("chimp-server-test", "0.0.1") @@ -36,26 +42,30 @@ class ZioMcpServerHttpSpec override protected def withStreamingServer( server: StreamingMcpServer[Task] )(test: BidirectionalMcpClient[Task] => Task[Assertion]): Future[Assertion] = - toFuture: - val routes = ZioHttpInterpreter().toHttp(ZioServerHttpTransport(List("mcp")).serve(server)) - ZIO.scoped: - (for - port <- Server.install(routes) - result <- HttpClientZioBackend().flatMap: backend => - ZioClientHttpTransport - .scoped(backend, uri"http://localhost:$port/mcp", ProtocolVersion.Latest, ClientTransport.defaultTimeout) - .flatMap(transport => McpClient.bidirectional(transport, clientInfo)) - .flatMap(client => test(client)) - .ensuring(backend.close().ignore) - yield result).provideSome[Scope](Server.defaultWithPort(0)) + withHttpServer(ZioServerHttpTransport(List("mcp")).serve(server), Nil): transport => + McpClient.bidirectional(transport, clientInfo).flatMap(test) + + override protected def withSecuredServer( + server: SecuredMcpServer[Task, String, String, User], + token: String + )(test: McpClient[Task] => Task[Assertion]): Future[Assertion] = + withHttpServer(server.endpoint(List("mcp")), List(Header.authorization("Bearer", token))): transport => + McpClient(transport, clientInfo).flatMap(test) override protected def withSecuredStreamingServer( server: SecuredStreamingMcpServer[Task, String, String, User], token: String )(test: BidirectionalMcpClient[Task] => Task[Assertion]): Future[Assertion] = + val endpoint = SecuredServerStreamingHttpTransport(List("mcp"), ZioServerHttpTransport(List("mcp"))).serve(server) + withHttpServer(endpoint, List(Header.authorization("Bearer", token))): transport => + McpClient.bidirectional(transport, clientInfo).flatMap(test) + + private def withHttpServer( + endpoint: ServerEndpoint[ZioStreams & WebSockets, Task], + headers: List[Header] + )(test: ZioClientHttpTransport => Task[Assertion]): Future[Assertion] = toFuture: - val routes = - ZioHttpInterpreter().toHttp(SecuredServerStreamingHttpTransport(List("mcp"), ZioServerHttpTransport(List("mcp"))).serve(server)) + val routes = ZioHttpInterpreter().toHttp(endpoint) ZIO.scoped: (for port <- Server.install(routes) @@ -66,9 +76,8 @@ class ZioMcpServerHttpSpec uri"http://localhost:$port/mcp", ProtocolVersion.Latest, ClientTransport.defaultTimeout, - headers = List(Header.authorization("Bearer", token)) + headers = headers ) - .flatMap(transport => McpClient.bidirectional(transport, clientInfo)) - .flatMap(client => test(client)) + .flatMap(test) .ensuring(backend.close().ignore) yield result).provideSome[Scope](Server.defaultWithPort(0)) diff --git a/server/src/main/scala/chimp/server/transport/JsonRpcHttpResponses.scala b/server/src/main/scala/chimp/server/transport/JsonRpcHttpResponses.scala index 3727dbe..b89b9ec 100644 --- a/server/src/main/scala/chimp/server/transport/JsonRpcHttpResponses.scala +++ b/server/src/main/scala/chimp/server/transport/JsonRpcHttpResponses.scala @@ -6,24 +6,17 @@ import sttp.model.{Header, HeaderNames, StatusCode} import sttp.monad.MonadError import sttp.monad.syntax.* -/** True if the request's `Host` and `Origin` headers pass the given check. Shared by every HTTP transport, secured or not, streaming or - * not. - */ private[transport] def originAllowed(originCheck: OriginCheck, headers: Seq[Header]): Boolean = val host = headers.find(_.name.equalsIgnoreCase(HeaderNames.Host)).map(_.value) val origin = headers.find(_.name.equalsIgnoreCase(HeaderNames.Origin)).map(_.value) originCheck.validate(host, origin) -/** Runs `handle` if the origin check passes, otherwise gives a forbidden response, without running `handle`. */ private[transport] def respondToJsonRpc[F[_]](originCheck: OriginCheck, headers: Seq[Header])(handle: => F[McpResponse])(using m: MonadError[F] ): F[(StatusCode, Option[Json])] = if !originAllowed(originCheck, headers) then m.unit((StatusCode.Forbidden, None)) else handle.map(response => (response.statusCode, response.body)) -/** The same as [[respondToJsonRpc]], but the response is an event stream produced by the given [[ServerStreamingHttpTransport]]'s streaming - * machinery, rather than a single JSON body. - */ private[transport] def respondWithEventStream[F[_], Caps]( originCheck: OriginCheck, headers: Seq[Header], diff --git a/server/src/test/scala/chimp/server/SecuredHttpMcpServerSpec.scala b/server/src/test/scala/chimp/server/SecuredHttpMcpServerSpec.scala deleted file mode 100644 index 7afa100..0000000 --- a/server/src/test/scala/chimp/server/SecuredHttpMcpServerSpec.scala +++ /dev/null @@ -1,129 +0,0 @@ -package chimp.server - -import chimp.client.{McpAuthorizationException, McpClient} -import chimp.client.transport.ClientHttpTransport -import chimp.protocol.{Implementation, ResourceContents, ToolContent} -import io.circe.{Codec, Json} -import org.scalatest.flatspec.AnyFlatSpec -import org.scalatest.matchers.should.Matchers -import ox.supervised -import sttp.client4.* -import sttp.model.{Header, StatusCode} -import sttp.shared.Identity -import sttp.tapir.* -import sttp.tapir.server.ServerEndpoint -import sttp.tapir.server.netty.sync.NettySyncServer - -class SecuredHttpMcpServerSpec extends AnyFlatSpec with Matchers: - private case class EchoInput(message: String) derives Codec, Schema - private case class User(email: String) - - private val clientInfo = Implementation("chimp-server-test", "0.0.1") - private val validToken = "s3cret" - - private val echoTool = tool("echo") - .description("Echoes a message.") - .input[EchoInput] - .handle(in => ToolResult.text(in.message)) - - private val whoAmITool = tool("whoAmI") - .description("Echoes a message and the caller's email.") - .input[EchoInput] - .handleSecured[User]((in, user) => ToolResult.text(s"${in.message} ${user.email}")) - - private val greetingResource = resource("test://greeting") - .name("greeting") - .mimeType("text/plain") - .handle(() => Right(List(ResourceContents.Text(uri = "test://greeting", text = "hello", mimeType = Some("text/plain"))))) - - private def securityLogic(token: String): Either[String, User] = - if token == validToken then Right(User("employee@example.com")) else Left("Invalid token") - - private def securityLogicEffectful(token: String): Identity[Either[String, User]] = securityLogic(token) - - private val securedServer = McpServer[Identity]() - .addTool(echoTool) - .serverSecurityLogicPure(auth.bearer[String](), statusCode(StatusCode.Unauthorized).and(stringBody))(securityLogic) - .addTool(whoAmITool) - - private val securedServerWithEffectfulSecurityLogic = McpServer[Identity]() - .addTool(echoTool) - .serverSecurityLogic(auth.bearer[String](), statusCode(StatusCode.Unauthorized).and(stringBody))(securityLogicEffectful) - .addTool(whoAmITool) - - private val serverConfiguredAfterSecurity = McpServer[Identity]() - .serverSecurityLogicPure(auth.bearer[String](), statusCode(StatusCode.Unauthorized).and(stringBody))(securityLogic) - .name("secured-server") - .version("2.0.0") - .addResource(greetingResource) - .addTool(whoAmITool) - - private def withServer[T](mcpEndpoint: ServerEndpoint[Any, Identity])(test: (Int, SyncBackend) => T): T = - supervised: - val binding = NettySyncServer().port(0).addEndpoint(mcpEndpoint).start() - try - val backend = DefaultSyncBackend() - try test(binding.port, backend) - finally backend.close() - finally binding.stop() - - private def withClient[T](port: Int, backend: SyncBackend, token: String)(test: McpClient[Identity] => T): T = - val transport = - ClientHttpTransport[Identity](backend, uri"http://localhost:$port/mcp", headers = List(Header.authorization("Bearer", token))) - try test(McpClient(transport, clientInfo)) - finally transport.close() - - private val securedEndpoint = securedServer.endpoint(List("mcp")) - private val endpointConfiguredAfterSecurity = serverConfiguredAfterSecurity.endpoint(List("mcp")) - private val endpointWithEffectfulSecurityLogic = securedServerWithEffectfulSecurityLogic.endpoint(List("mcp")) - - "a secured MCP server" should "give the principal to the tool logic" in withServer(securedEndpoint): (port, backend) => - withClient(port, backend, validToken): client => - val result = client.callTool("whoAmI", Json.obj("message" -> Json.fromString("hi"))) - result.isError shouldBe false - result.content shouldBe List(ToolContent.Text("text", "hi employee@example.com")) - - it should "give the principal to the tool logic when the security logic is effectful" in - withServer(endpointWithEffectfulSecurityLogic): (port, backend) => - withClient(port, backend, validToken): client => - val result = client.callTool("whoAmI", Json.obj("message" -> Json.fromString("hi"))) - result.isError shouldBe false - result.content shouldBe List(ToolContent.Text("text", "hi employee@example.com")) - - it should "reject an invalid security input with the error output when the security logic is effectful" in - withServer(endpointWithEffectfulSecurityLogic): (port, backend) => - val exception = intercept[McpAuthorizationException](withClient(port, backend, "wrong")(_ => ())) - exception.statusCode shouldBe StatusCode.Unauthorized.code - - it should "also serve the tools which do not need the principal" in withServer(securedEndpoint): (port, backend) => - withClient(port, backend, validToken): client => - client.listTools().tools.map(_.name) should contain allOf ("echo", "whoAmI") - client.callTool("echo", Json.obj("message" -> Json.fromString("hi"))).content shouldBe List(ToolContent.Text("text", "hi")) - - it should "use the identity which the builders set after the security logic" in - withServer(endpointConfiguredAfterSecurity): (port, backend) => - withClient(port, backend, validToken): client => - client.serverInfo shouldBe Implementation("secured-server", "2.0.0") - - it should "serve a resource which was added after the security logic" in - withServer(endpointConfiguredAfterSecurity): (port, backend) => - withClient(port, backend, validToken): client => - client.serverCapabilities.resources shouldBe defined - client.listResources().resources.map(_.uri) shouldBe List("test://greeting") - client.readResource("test://greeting").contents.head match - case ResourceContents.Text(_, text, _, _) => text shouldBe "hello" - case other => fail(s"expected text contents, got $other") - - it should "reject an invalid security input with the error output, before any tool logic runs" in - withServer(securedEndpoint): (port, backend) => - val exception = intercept[McpAuthorizationException](withClient(port, backend, "wrong")(_ => ())) - exception.statusCode shouldBe StatusCode.Unauthorized.code - - it should "reject a request with no security input with the error output" in withServer(securedEndpoint): (port, backend) => - val response = basicRequest - .post(uri"http://localhost:$port/mcp") - .header("Content-Type", "application/json") - .body("""{"jsonrpc":"2.0","id":"1","method":"tools/call","params":{"name":"whoAmI","arguments":{"message":"hi"}}}""") - .send(backend) - response.code shouldBe StatusCode.Unauthorized - response.body shouldBe Left("Invalid value for: header Authorization (missing)") diff --git a/server/src/test/scala/chimp/server/SecuredMcpServerStreamingTests.scala b/server/src/test/scala/chimp/server/SecuredMcpServerStreamingTests.scala index ece3242..6dc12c1 100644 --- a/server/src/test/scala/chimp/server/SecuredMcpServerStreamingTests.scala +++ b/server/src/test/scala/chimp/server/SecuredMcpServerStreamingTests.scala @@ -15,13 +15,14 @@ import java.util.concurrent.ConcurrentLinkedQueue import scala.concurrent.Future import scala.jdk.CollectionConverters.* -trait SecuredMcpServerStreamingTests[F[_]] extends AsyncFlatSpec with Matchers with RecoverMethods: - this: ToFuture[F] => - +trait SecuredMcpServerTestFixtures[F[_]]: protected case class User(email: String) protected val validToken = "s3cret" +trait SecuredMcpServerStreamingTests[F[_]] extends AsyncFlatSpec with Matchers with RecoverMethods with SecuredMcpServerTestFixtures[F]: + this: ToFuture[F] => + /** @param token * The bearer token the client authenticates with; defaults to [[validToken]] so most tests need not pass it. */ diff --git a/server/src/test/scala/chimp/server/SecuredMcpServerTests.scala b/server/src/test/scala/chimp/server/SecuredMcpServerTests.scala new file mode 100644 index 0000000..3a4208e --- /dev/null +++ b/server/src/test/scala/chimp/server/SecuredMcpServerTests.scala @@ -0,0 +1,120 @@ +package chimp.server + +import chimp.client.{McpAuthorizationException, McpClient} +import chimp.protocol.{Implementation, ResourceContents, ToolContent} +import io.circe.{Codec, Json} +import org.scalatest.{Assertion, RecoverMethods} +import org.scalatest.flatspec.AsyncFlatSpec +import org.scalatest.matchers.should.Matchers +import sttp.monad.syntax.* +import sttp.model.StatusCode +import sttp.tapir.* + +import scala.concurrent.Future + +trait SecuredMcpServerTests[F[_]] extends AsyncFlatSpec with Matchers with RecoverMethods with SecuredMcpServerTestFixtures[F]: + this: ToFuture[F] => + + protected def withSecuredServer(server: SecuredMcpServer[F, String, String, User], token: String = validToken)( + test: McpClient[F] => F[Assertion] + ): Future[Assertion] + + private case class EchoInput(message: String) derives Codec, Schema + + private def securityLogic(token: String): Either[String, User] = + if token == validToken then Right(User("employee@example.com")) else Left("Invalid token") + + private def securityLogicEffectful(token: String): F[Either[String, User]] = + monad.unit(securityLogic(token)) + + private def echoTool: ServerTool[EchoInput, NoStructuredOutput, F, ServerContext[F]] = + tool("echo") + .description("Echoes a message.") + .input[EchoInput] + .serverLogic[F]((in, _) => monad.unit(ToolResult.text(in.message))) + + private def whoAmITool: ServerTool[EchoInput, NoStructuredOutput, F, SecuredServerContext[F, User]] = + tool("whoAmI") + .description("Echoes a message and the caller's email.") + .input[EchoInput] + .securedServerLogic[F, User]((in, user, _) => monad.unit(ToolResult.text(s"${in.message} ${user.email}"))) + + private def greetingResource: ServerResource[F] = + resource("test://greeting") + .name("greeting") + .mimeType("text/plain") + .serverLogic[F](_ => + monad.unit(Right(List(ResourceContents.Text(uri = "test://greeting", text = "hello", mimeType = Some("text/plain"))))) + ) + + private def securedServer: SecuredMcpServer[F, String, String, User] = + McpServer[F]() + .addTool(echoTool) + .serverSecurityLogicPure(auth.bearer[String](), statusCode(StatusCode.Unauthorized).and(stringBody))(securityLogic) + .addTool(whoAmITool) + + private def securedServerWithEffectfulSecurityLogic: SecuredMcpServer[F, String, String, User] = + McpServer[F]() + .addTool(echoTool) + .serverSecurityLogic(auth.bearer[String](), statusCode(StatusCode.Unauthorized).and(stringBody))(securityLogicEffectful) + .addTool(whoAmITool) + + private def serverConfiguredAfterSecurity: SecuredMcpServer[F, String, String, User] = + McpServer[F]() + .serverSecurityLogicPure(auth.bearer[String](), statusCode(StatusCode.Unauthorized).and(stringBody))(securityLogic) + .name("secured-server") + .version("2.0.0") + .addResource(greetingResource) + .addTool(whoAmITool) + + private def assertWhoAmIDeliversPrincipal(client: McpClient[F]): F[Assertion] = + client + .callTool("whoAmI", Json.obj("message" -> Json.fromString("hi"))) + .map: result => + result.isError shouldBe false + result.content shouldBe List(ToolContent.Text("text", "hi employee@example.com")) + + "a secured MCP server" should "give the principal to the tool logic" in + withSecuredServer(securedServer)(assertWhoAmIDeliversPrincipal) + + it should "give the principal to the tool logic when the security logic is effectful" in + withSecuredServer(securedServerWithEffectfulSecurityLogic)(assertWhoAmIDeliversPrincipal) + + it should "also serve tools which do not need the principal" in + withSecuredServer(securedServer): client => + client + .listTools() + .flatMap: tools => + tools.tools.map(_.name) should contain allOf ("echo", "whoAmI") + client + .callTool("echo", Json.obj("message" -> Json.fromString("hi"))) + .map: result => + result.content shouldBe List(ToolContent.Text("text", "hi")) + + it should "use the identity which the builders set after the security logic" in + withSecuredServer(serverConfiguredAfterSecurity): client => + monad.unit(client.serverInfo shouldBe Implementation("secured-server", "2.0.0")) + + it should "serve a resource which was added after the security logic" in + withSecuredServer(serverConfiguredAfterSecurity): client => + client.serverCapabilities.resources shouldBe defined + client + .listResources() + .flatMap: listed => + listed.resources.map(_.uri) shouldBe List("test://greeting") + client + .readResource("test://greeting") + .map: result => + result.contents.head match + case ResourceContents.Text(_, text, _, _) => text shouldBe "hello" + case other => fail(s"expected text contents, got $other") + + it should "reject an invalid security input with HTTP 401 before any tool logic runs" in + recoverToExceptionIf[McpAuthorizationException] { + Future(withSecuredServer(securedServer, token = "wrong")(_ => monad.unit(succeed))).flatten + }.map(_.statusCode shouldBe StatusCode.Unauthorized.code) + + it should "reject invalid security input with HTTP 401 when the security logic is effectful" in + recoverToExceptionIf[McpAuthorizationException] { + Future(withSecuredServer(securedServerWithEffectfulSecurityLogic, token = "wrong")(_ => monad.unit(succeed))).flatten + }.map(_.statusCode shouldBe StatusCode.Unauthorized.code)