From 4b5e95b6fb67802174a64c33b85dc41dee37c309 Mon Sep 17 00:00:00 2001 From: Maxime Perocheau Date: Wed, 16 Sep 2026 10:55:47 -0700 Subject: [PATCH 1/3] Add secure prompts and resources --- docs/server/prompts.md | 1 + docs/server/resources.md | 1 + docs/server/transport.md | 14 +- .../main/scala/chimp/server/McpHandler.scala | 67 ++++---- .../main/scala/chimp/server/McpServer.scala | 153 +++++++++++++++++- .../src/main/scala/chimp/server/Prompt.scala | 16 ++ .../main/scala/chimp/server/Resource.scala | 35 ++++ .../SecuredMcpServerStreamingTests.scala | 20 +++ .../chimp/server/SecuredMcpServerTests.scala | 78 ++++++++- 9 files changed, 342 insertions(+), 43 deletions(-) diff --git a/docs/server/prompts.md b/docs/server/prompts.md index 34128d37..0e818c17 100644 --- a/docs/server/prompts.md +++ b/docs/server/prompts.md @@ -6,6 +6,7 @@ - `handle` — synchronous logic from the argument values to `GetPromptResult`. - `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). - Register prompts with `.addPrompt` / `.addPrompts`. ```scala mdoc:compile-only diff --git a/docs/server/resources.md b/docs/server/resources.md index 7bbe3e8c..86f823f5 100644 --- a/docs/server/resources.md +++ b/docs/server/resources.md @@ -6,6 +6,7 @@ - `handle` — synchronous read logic. - `handleWithHeaders` — synchronous read logic that also receives the request headers. - `serverLogic` — effectful read 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). - Register with `.addResource` / `.addResourceTemplate`. Subscriptions are wired with `.withSubscriptions`. ```scala mdoc:compile-only diff --git a/docs/server/transport.md b/docs/server/transport.md index c63f21a9..d4b7dd65 100644 --- a/docs/server/transport.md +++ b/docs/server/transport.md @@ -106,17 +106,18 @@ object SecuredMcpServer: NettySyncServer().port(8080).addEndpoint(securedEndpoint).startAndWait() ``` -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. +The result of the security logic of `prependSecurity` does not reach the tool, prompt, or resource logic. Use `prependSecurity` if they do not need data from the caller. If they need such data, read the request headers with `handleWithHeaders` (or `serverLogic` with headers), or give them 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 +### Giving the security result to tools, prompts, and resources -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. +To validate the caller one time and give the result to the tool, prompt, or resource 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: +The server gives the principal to the logic of each tool, prompt, or resource which you add to it. Define such handlers with `handleSecured`, or with `securedServerLogic` if the logic needs an effect. Handlers 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.protocol.ResourceContents import chimp.server.* import sttp.model.StatusCode import sttp.shared.Identity @@ -129,6 +130,8 @@ 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 profile = resource("user://profile") + .handleSecured[User](user => Right(List(ResourceContents.Text(uri = "user://profile", text = user.email)))) val securedEndpoint = McpServer[Identity]() .serverSecurityLogicPure( @@ -137,12 +140,13 @@ object McpServerWithPrincipal: )(token => if token == "s3cret" then Right(User("employee@example.com")) else Left("Invalid token")) .name("my-mcp-server") .addTools(echo, whoAmI) + .addResource(profile) .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. +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, prompt, or resource 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 diff --git a/server/src/main/scala/chimp/server/McpHandler.scala b/server/src/main/scala/chimp/server/McpHandler.scala index 585dcee2..d4121445 100644 --- a/server/src/main/scala/chimp/server/McpHandler.scala +++ b/server/src/main/scala/chimp/server/McpHandler.scala @@ -29,9 +29,8 @@ enum McpResponse: private[server] class McpHandler[F[_], C <: ServerContext[F]](server: McpServerDef[F, C]): private val logger = LoggerFactory.getLogger(classOf[McpHandler[?, ?]]) private val toolsByName = server.tools.map(tool => tool.name -> tool).toMap - private val promptsByName = server.prompts.map(prompt => prompt.definition.name -> prompt).toMap - private val resourcesByUri = server.resources.map(resource => resource.definition.uri -> resource).toMap - private val hasResources = server.resources.nonEmpty || server.resourceTemplates.nonEmpty + private val hasResources = server.resourceDefinitions.nonEmpty || server.resourceTemplateDefinitions.nonEmpty + private val hasPrompts = server.promptDefinitions.nonEmpty private val toolDefinitions = server.tools.map(toolToDefinition) private def toJsonSchema(toolSchema: ToolSchema): Json = toolSchema match @@ -83,21 +82,21 @@ private[server] class McpHandler[F[_], C <: ServerContext[F]](server: McpServerD case "tools/call" => handleToolsCall(params, id, headers, makeContext).map(jsonResponse) case "resources/list" if hasResources => - jsonResponse(JSONRPCMessage.Response(id = id, result = ListResourcesResult(server.resources.map(_.definition)).asJson)).unit + jsonResponse(JSONRPCMessage.Response(id = id, result = ListResourcesResult(server.resourceDefinitions).asJson)).unit case "resources/templates/list" if hasResources => jsonResponse( - JSONRPCMessage.Response(id = id, result = ListResourceTemplatesResult(server.resourceTemplates.map(_.definition)).asJson) + JSONRPCMessage.Response(id = id, result = ListResourceTemplatesResult(server.resourceTemplateDefinitions).asJson) ).unit case "resources/read" if hasResources => - handleResourcesRead(params, id, headers).map(jsonResponse) + handleResourcesRead(params, id, headers, makeContext).map(jsonResponse) case "resources/subscribe" if server.subscriptions.isDefined => handleSubscribe(params, id, subscribe = true).map(jsonResponse) case "resources/unsubscribe" if server.subscriptions.isDefined => handleSubscribe(params, id, subscribe = false).map(jsonResponse) - case "prompts/list" if server.prompts.nonEmpty => - jsonResponse(JSONRPCMessage.Response(id = id, result = ListPromptsResult(server.prompts.map(_.definition)).asJson)).unit - case "prompts/get" if server.prompts.nonEmpty => - handlePromptsGet(params, id, headers).map(jsonResponse) + case "prompts/list" if hasPrompts => + jsonResponse(JSONRPCMessage.Response(id = id, result = ListPromptsResult(server.promptDefinitions).asJson)).unit + case "prompts/get" if hasPrompts => + handlePromptsGet(params, id, headers, makeContext).map(jsonResponse) case "completion/complete" if server.completion.isDefined => handleComplete(params, id).map(jsonResponse) case "logging/setLevel" if server.loggingLevel.isDefined => @@ -123,7 +122,7 @@ private[server] class McpHandler[F[_], C <: ServerContext[F]](server: McpServerD val capabilities = ServerCapabilities( logging = Option.when(server.loggingLevel.isDefined)(Json.obj()), completions = Option.when(server.completion.isDefined)(Json.obj()), - prompts = Option.when(server.prompts.nonEmpty)(ServerPromptsCapability(listChanged = Some(false))), + prompts = Option.when(hasPrompts)(ServerPromptsCapability(listChanged = Some(false))), resources = Option.when(hasResources)(ServerResourcesCapability(subscribe = Some(server.subscriptions.isDefined), listChanged = Some(false))), tools = Option.when(server.tools.nonEmpty)(ServerToolsCapability(listChanged = Some(false))) @@ -178,23 +177,22 @@ private[server] class McpHandler[F[_], C <: ServerContext[F]](server: McpServerD ).asJson ) - private def handleResourcesRead(params: Option[Json], id: RequestId, headers: Seq[Header])(using MonadError[F]): F[JSONRPCMessage] = + private def handleResourcesRead( + params: Option[Json], + id: RequestId, + headers: Seq[Header], + makeContext: Option[ProgressToken] => C + )(using MonadError[F]): F[JSONRPCMessage] = decodeParams[ReadResourceParams](params, id): params => - resourcesByUri.get(params.uri) match - case Some(resource) => resource.read(headers).map(resourceReadResponse(id, params.uri)) - case None => - val templateMatch = server.resourceTemplates.iterator - .map(template => (template, template.matcher.matchUri(params.uri))) - .collectFirst { case (template, Some(vars)) => (template, vars) } - templateMatch match - case Some((template, vars)) => template.read(vars, params.uri, headers).map(resourceReadResponse(id, params.uri)) - case None => - protocolError( - id, - JSONRPCErrorCodes.ResourceNotFound.code, - s"Resource not found: ${params.uri}", - Some(Json.obj("uri" -> Json.fromString(params.uri))) - ).unit + server.invokeResource(params.uri, makeContext(None), headers) match + case Some(result) => result.map(resourceReadResponse(id, params.uri)) + case None => + protocolError( + id, + JSONRPCErrorCodes.ResourceNotFound.code, + s"Resource not found: ${params.uri}", + Some(Json.obj("uri" -> Json.fromString(params.uri))) + ).unit private def decodeParams[P: Decoder](params: Option[Json], id: RequestId)(f: P => F[JSONRPCMessage])(using MonadError[F] @@ -221,13 +219,16 @@ private[server] class McpHandler[F[_], C <: ServerContext[F]](server: McpServerD private def emptyResult(id: RequestId): JSONRPCMessage = JSONRPCMessage.Response(id = id, result = Json.obj()) - private def handlePromptsGet(params: Option[Json], id: RequestId, headers: Seq[Header])(using MonadError[F]): F[JSONRPCMessage] = + private def handlePromptsGet( + params: Option[Json], + id: RequestId, + headers: Seq[Header], + makeContext: Option[ProgressToken] => C + )(using MonadError[F]): F[JSONRPCMessage] = decodeParams[GetPromptParams](params, id): params => - promptsByName.get(params.name) match - case Some(prompt) => - prompt - .logic(params.arguments.getOrElse(Map.empty), headers) - .map(result => JSONRPCMessage.Response(id = id, result = result.asJson)) + server.invokePrompt(params.name, params.arguments.getOrElse(Map.empty), makeContext(None), headers) match + case Some(result) => + result.map(promptResult => JSONRPCMessage.Response(id = id, result = promptResult.asJson)) case None => protocolError(id, JSONRPCErrorCodes.InvalidParams.code, s"Unknown prompt: ${params.name}").unit private def handleComplete(params: Option[Json], id: RequestId)(using MonadError[F]): F[JSONRPCMessage] = diff --git a/server/src/main/scala/chimp/server/McpServer.scala b/server/src/main/scala/chimp/server/McpServer.scala index 49f380c8..fc0ffcde 100644 --- a/server/src/main/scala/chimp/server/McpServer.scala +++ b/server/src/main/scala/chimp/server/McpServer.scala @@ -2,10 +2,13 @@ package chimp.server import chimp.protocol.* import chimp.server.transport.{SecuredServerHttpTransport, ServerHttpTransport} +import sttp.model.Header import sttp.monad.MonadError import sttp.tapir.server.ServerEndpoint import sttp.tapir.{EndpointInput, EndpointOutput} +import scala.annotation.{targetName, unused} + type CompletionHandler[F[_]] = (CompleteRef, CompleteArgument, Option[CompleteContext]) => F[Completion] type SetLoggingLevelHandler[F[_]] = LoggingLevel => F[Unit] @@ -29,6 +32,41 @@ sealed trait McpServerDef[F[_], C <: ServerContext[F]]: def loggingLevel: Option[SetLoggingLevelHandler[F]] def subscriptions: Option[ResourceSubscriptions[F]] + private[server] lazy val promptsByName: Map[String, ServerPrompt[F]] = + prompts.map(prompt => prompt.definition.name -> prompt).toMap + + private[server] lazy val resourcesByUri: Map[String, ServerResource[F]] = + resources.map(resource => resource.definition.uri -> resource).toMap + + private[server] def promptDefinitions: List[Prompt] = prompts.map(_.definition) + + private[server] def invokePrompt( + name: String, + args: Map[String, String], + @unused context: C, + headers: Seq[Header] + ): Option[F[GetPromptResult]] = + promptsByName.get(name).map(_.logic(args, headers)) + + private[server] def resourceDefinitions: List[Resource] = resources.map(_.definition) + + private[server] def resourceTemplateDefinitions: List[ResourceTemplate] = resourceTemplates.map(_.definition) + + private[server] def invokeResource( + uri: String, + @unused context: C, + headers: Seq[Header] + ): Option[F[Either[ResourceError, List[ResourceContents]]]] = + resourcesByUri + .get(uri) + .map(_.read(headers)) + .orElse(matchResourceTemplate(uri).map((template, vars) => template.read(vars, uri, headers))) + + private[server] def matchResourceTemplate(uri: String): Option[(ServerResourceTemplate[F], Map[String, String])] = + resourceTemplates.iterator + .map(template => (template, template.matcher.matchUri(uri))) + .collectFirst { case (template, Some(vars)) => (template, vars) } + case class McpServer[F[_]]( name: String = "Chimp MCP server", version: String = "1.0.0", @@ -94,7 +132,7 @@ 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. + * principal. The principal is given to the logic of the tools, prompts, and resources which are added to the returned server. */ def serverSecurityLogic[S, E, P](securityInput: EndpointInput[S], errorOutput: EndpointOutput[E])( logic: S => F[Either[E, P]] @@ -192,8 +230,8 @@ case class 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. + * principal, is given to the logic of the tools, prompts, and resources which are added to this server. Tools, prompts, and resources 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. @@ -207,7 +245,10 @@ case class SecuredMcpServer[F[_], S, E, P]( securityInput: EndpointInput[S], errorOutput: EndpointOutput[E], securityLogic: MonadError[F] => S => F[Either[E, P]], - securedTools: List[ServerTool[?, ?, F, SecuredServerContext[F, P]]] = Nil + securedTools: List[ServerTool[?, ?, F, SecuredServerContext[F, P]]] = Nil, + securedPrompts: List[SecuredServerPrompt[F, P]] = Nil, + securedResources: List[SecuredServerResource[F, P]] = Nil, + securedResourceTemplates: List[SecuredServerResourceTemplate[F, P]] = Nil ) extends McpServerDef[F, SecuredServerContext[F, P]]: def name: String = server.name def version: String = server.version @@ -223,6 +264,47 @@ case class SecuredMcpServer[F[_], S, E, P]( def tools: List[ServerTool[?, ?, F, SecuredServerContext[F, P]]] = server.tools ++ securedTools + private lazy val securedPromptsByName: Map[String, SecuredServerPrompt[F, P]] = + securedPrompts.map(prompt => prompt.definition.name -> prompt).toMap + + private lazy val securedResourcesByUri: Map[String, SecuredServerResource[F, P]] = + securedResources.map(resource => resource.definition.uri -> resource).toMap + + override private[server] def promptDefinitions: List[Prompt] = + prompts.map(_.definition) ++ securedPrompts.map(_.definition) + + override private[server] def invokePrompt( + name: String, + args: Map[String, String], + context: SecuredServerContext[F, P], + headers: Seq[Header] + ): Option[F[GetPromptResult]] = + promptsByName + .get(name) + .map(_.logic(args, headers)) + .orElse(securedPromptsByName.get(name).map(_.logic(args, context.principal, headers))) + + override private[server] def resourceDefinitions: List[Resource] = + resources.map(_.definition) ++ securedResources.map(_.definition) + + override private[server] def resourceTemplateDefinitions: List[ResourceTemplate] = + resourceTemplates.map(_.definition) ++ securedResourceTemplates.map(_.definition) + + override private[server] def invokeResource( + uri: String, + context: SecuredServerContext[F, P], + headers: Seq[Header] + ): Option[F[Either[ResourceError, List[ResourceContents]]]] = + resourcesByUri + .get(uri) + .map(_.read(headers)) + .orElse(securedResourcesByUri.get(uri).map(_.read(context.principal, headers))) + .orElse(matchResourceTemplate(uri).map((template, vars) => template.read(vars, uri, headers))) + .orElse: + securedResourceTemplates.iterator + .map(template => (template, template.matcher.matchUri(uri))) + .collectFirst { case (template, Some(vars)) => template.read(vars, uri, context.principal, headers) } + def name(value: String): SecuredMcpServer[F, S, E, P] = copy(server = server.name(value)) @@ -247,21 +329,42 @@ case class SecuredMcpServer[F[_], S, E, P]( def addPrompt(prompt: ServerPrompt[F]): SecuredMcpServer[F, S, E, P] = copy(server = server.addPrompt(prompt)) + def addPrompt(prompt: SecuredServerPrompt[F, P]): SecuredMcpServer[F, S, E, P] = + copy(securedPrompts = securedPrompts :+ prompt) + def addPrompts(prompts: ServerPrompt[F]*): SecuredMcpServer[F, S, E, P] = copy(server = server.addPrompts(prompts*)) + @targetName("addSecuredPrompts") + def addPrompts(prompts: SecuredServerPrompt[F, P]*): SecuredMcpServer[F, S, E, P] = + copy(securedPrompts = securedPrompts ++ prompts) + def addResource(resource: ServerResource[F]): SecuredMcpServer[F, S, E, P] = copy(server = server.addResource(resource)) + def addResource(resource: SecuredServerResource[F, P]): SecuredMcpServer[F, S, E, P] = + copy(securedResources = securedResources :+ resource) + def addResources(resources: ServerResource[F]*): SecuredMcpServer[F, S, E, P] = copy(server = server.addResources(resources*)) + @targetName("addSecuredResources") + def addResources(resources: SecuredServerResource[F, P]*): SecuredMcpServer[F, S, E, P] = + copy(securedResources = securedResources ++ resources) + def addResourceTemplate(resourceTemplate: ServerResourceTemplate[F]): SecuredMcpServer[F, S, E, P] = copy(server = server.addResourceTemplate(resourceTemplate)) + def addResourceTemplate(resourceTemplate: SecuredServerResourceTemplate[F, P]): SecuredMcpServer[F, S, E, P] = + copy(securedResourceTemplates = securedResourceTemplates :+ resourceTemplate) + def addResourceTemplates(resourceTemplates: ServerResourceTemplate[F]*): SecuredMcpServer[F, S, E, P] = copy(server = server.addResourceTemplates(resourceTemplates*)) + @targetName("addSecuredResourceTemplates") + def addResourceTemplates(resourceTemplates: SecuredServerResourceTemplate[F, P]*): SecuredMcpServer[F, S, E, P] = + copy(securedResourceTemplates = securedResourceTemplates ++ resourceTemplates) + def withCompletion(handler: CompletionHandler[F]): SecuredMcpServer[F, S, E, P] = copy(server = server.withCompletion(handler)) @@ -300,6 +403,27 @@ case class SecuredStreamingMcpServer[F[_], S, E, P]( def tools: List[ServerTool[?, ?, F, SecuredStreamingServerContext[F, P]]] = server.tools ++ streamingTools + override private[server] def promptDefinitions: List[Prompt] = server.promptDefinitions + + override private[server] def invokePrompt( + name: String, + args: Map[String, String], + context: SecuredStreamingServerContext[F, P], + headers: Seq[Header] + ): Option[F[GetPromptResult]] = + server.invokePrompt(name, args, context, headers) + + override private[server] def resourceDefinitions: List[Resource] = server.resourceDefinitions + + override private[server] def resourceTemplateDefinitions: List[ResourceTemplate] = server.resourceTemplateDefinitions + + override private[server] def invokeResource( + uri: String, + context: SecuredStreamingServerContext[F, P], + headers: Seq[Header] + ): Option[F[Either[ResourceError, List[ResourceContents]]]] = + server.invokeResource(uri, context, headers) + def name(value: String): SecuredStreamingMcpServer[F, S, E, P] = copy(server = server.name(value)) @@ -330,21 +454,42 @@ case class SecuredStreamingMcpServer[F[_], S, E, P]( def addPrompt(prompt: ServerPrompt[F]): SecuredStreamingMcpServer[F, S, E, P] = copy(server = server.addPrompt(prompt)) + def addPrompt(prompt: SecuredServerPrompt[F, P]): 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*)) + @targetName("addSecuredPrompts") + def addPrompts(prompts: SecuredServerPrompt[F, P]*): 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 addResource(resource: SecuredServerResource[F, P]): 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*)) + @targetName("addSecuredResources") + def addResources(resources: SecuredServerResource[F, P]*): 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 addResourceTemplate(resourceTemplate: SecuredServerResourceTemplate[F, P]): 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*)) + @targetName("addSecuredResourceTemplates") + def addResourceTemplates(resourceTemplates: SecuredServerResourceTemplate[F, P]*): 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)) diff --git a/server/src/main/scala/chimp/server/Prompt.scala b/server/src/main/scala/chimp/server/Prompt.scala index 6e420819..bacc1dfb 100644 --- a/server/src/main/scala/chimp/server/Prompt.scala +++ b/server/src/main/scala/chimp/server/Prompt.scala @@ -32,6 +32,12 @@ case class PartialPrompt( def serverLogic[F[_]](logic: (Map[String, String], Seq[Header]) => F[GetPromptResult]): ServerPrompt[F] = ServerPrompt(definition, logic) + /** Attaches effectful logic, with access to the principal; usable only on a [[SecuredMcpServer]]. */ + def securedServerLogic[F[_], P]( + logic: (Map[String, String], P, Seq[Header]) => F[GetPromptResult] + ): SecuredServerPrompt[F, P] = + SecuredServerPrompt(definition, logic) + /** Attaches synchronous logic that also receives the request headers. */ def handleWithHeaders(logic: (Map[String, String], Seq[Header]) => GetPromptResult): ServerPrompt[Identity] = ServerPrompt(definition, logic) @@ -40,6 +46,10 @@ case class PartialPrompt( def handle(logic: Map[String, String] => GetPromptResult): ServerPrompt[Identity] = handleWithHeaders((args, _) => logic(args)) + /** Attaches synchronous logic over the supplied argument values and the principal; usable only on a [[SecuredMcpServer]]. */ + def handleSecured[P](logic: (Map[String, String], P) => GetPromptResult): SecuredServerPrompt[Identity, P] = + securedServerLogic[Identity, P]((args, principal, _) => logic(args, principal)) + private def definition: Prompt = Prompt(name, title, description, Option.when(arguments.nonEmpty)(arguments)) @@ -47,3 +57,9 @@ end PartialPrompt /** A fully-defined prompt: its metadata plus the logic producing its messages. */ case class ServerPrompt[F[_]](definition: Prompt, logic: (Map[String, String], Seq[Header]) => F[GetPromptResult]) + +/** A prompt whose logic also receives the principal made by a [[SecuredMcpServer]]. */ +case class SecuredServerPrompt[F[_], P]( + definition: Prompt, + logic: (Map[String, String], P, Seq[Header]) => F[GetPromptResult] +) diff --git a/server/src/main/scala/chimp/server/Resource.scala b/server/src/main/scala/chimp/server/Resource.scala index 60ec2d29..56901906 100644 --- a/server/src/main/scala/chimp/server/Resource.scala +++ b/server/src/main/scala/chimp/server/Resource.scala @@ -44,6 +44,12 @@ case class PartialResource( def serverLogic[F[_]](logic: Seq[Header] => F[Either[ResourceError, List[ResourceContents]]]): ServerResource[F] = ServerResource(definition, logic) + /** Attaches effectful logic, with access to the principal; usable only on a [[SecuredMcpServer]]. */ + def securedServerLogic[F[_], P]( + logic: (P, Seq[Header]) => F[Either[ResourceError, List[ResourceContents]]] + ): SecuredServerResource[F, P] = + SecuredServerResource(definition, logic) + /** Attaches synchronous logic that also receives the request headers. */ def handleWithHeaders(logic: Seq[Header] => Either[ResourceError, List[ResourceContents]]): ServerResource[Identity] = ServerResource(definition, logic) @@ -52,6 +58,10 @@ case class PartialResource( def handle(logic: () => Either[ResourceError, List[ResourceContents]]): ServerResource[Identity] = handleWithHeaders(_ => logic()) + /** Attaches synchronous logic over the principal; usable only on a [[SecuredMcpServer]]. */ + def handleSecured[P](logic: P => Either[ResourceError, List[ResourceContents]]): SecuredServerResource[Identity, P] = + securedServerLogic[Identity, P]((principal, _) => logic(principal)) + private def definition: Resource = Resource(uri, name.getOrElse(uri), title, description, mimeType, size) end PartialResource @@ -59,6 +69,12 @@ end PartialResource /** A fully-defined resource: its metadata plus the logic reading its contents. */ case class ServerResource[F[_]](definition: Resource, read: Seq[Header] => F[Either[ResourceError, List[ResourceContents]]]) +/** A resource whose read logic also receives the principal made by a [[SecuredMcpServer]]. */ +case class SecuredServerResource[F[_], P]( + definition: Resource, + read: (P, Seq[Header]) => F[Either[ResourceError, List[ResourceContents]]] +) + /** A resource template being defined, before its read logic is attached. */ case class PartialResourceTemplate( uriTemplate: String, @@ -85,6 +101,12 @@ case class PartialResourceTemplate( ): ServerResourceTemplate[F] = ServerResourceTemplate(definition, UriTemplate.compile(uriTemplate), logic) + /** Attaches effectful logic, with access to the principal; usable only on a [[SecuredMcpServer]]. */ + def securedServerLogic[F[_], P]( + logic: (Map[String, String], String, P, Seq[Header]) => F[Either[ResourceError, List[ResourceContents]]] + ): SecuredServerResourceTemplate[F, P] = + SecuredServerResourceTemplate(definition, UriTemplate.compile(uriTemplate), logic) + /** Attaches synchronous logic that also receives the request headers. */ def handleWithHeaders( logic: (Map[String, String], String, Seq[Header]) => Either[ResourceError, List[ResourceContents]] @@ -95,6 +117,12 @@ case class PartialResourceTemplate( def handle(logic: (Map[String, String], String) => Either[ResourceError, List[ResourceContents]]): ServerResourceTemplate[Identity] = handleWithHeaders((vars, uri, _) => logic(vars, uri)) + /** Attaches synchronous logic over the extracted variables, the full URI, and the principal; usable only on a [[SecuredMcpServer]]. */ + def handleSecured[P]( + logic: (Map[String, String], String, P) => Either[ResourceError, List[ResourceContents]] + ): SecuredServerResourceTemplate[Identity, P] = + securedServerLogic[Identity, P]((vars, uri, principal, _) => logic(vars, uri, principal)) + private def definition: ResourceTemplate = ResourceTemplate(uriTemplate, name.getOrElse(uriTemplate), title, description, mimeType) @@ -107,6 +135,13 @@ case class ServerResourceTemplate[F[_]]( read: (Map[String, String], String, Seq[Header]) => F[Either[ResourceError, List[ResourceContents]]] ) +/** A resource template whose read logic also receives the principal made by a [[SecuredMcpServer]]. */ +case class SecuredServerResourceTemplate[F[_], P]( + definition: ResourceTemplate, + matcher: UriTemplate, + read: (Map[String, String], String, P, Seq[Header]) => F[Either[ResourceError, List[ResourceContents]]] +) + /** A compiled URI template that matches concrete URIs and extracts their `{variable}` values. */ final class UriTemplate private (regex: Regex, names: List[String]): /** Returns the extracted variables if `uri` matches, or `None` otherwise. */ diff --git a/server/src/test/scala/chimp/server/SecuredMcpServerStreamingTests.scala b/server/src/test/scala/chimp/server/SecuredMcpServerStreamingTests.scala index 6dc12c1e..785d9690 100644 --- a/server/src/test/scala/chimp/server/SecuredMcpServerStreamingTests.scala +++ b/server/src/test/scala/chimp/server/SecuredMcpServerStreamingTests.scala @@ -59,6 +59,20 @@ trait SecuredMcpServerStreamingTests[F[_]] extends AsyncFlatSpec with Matchers w .streaming .addStreamingTool(whoAmITool) + private def whoAmIPrompt: SecuredServerPrompt[F, User] = + prompt("whoAmI") + .description("Greets the caller") + .securedServerLogic[F, User]((_, user, _) => + monad.unit(GetPromptResult(messages = List(PromptMessage(Role.User, ToolContent.Text(text = user.email))))) + ) + + private def securedStreamingServerWithPrompt: SecuredStreamingMcpServer[F, String, String, User] = + McpServer[F]() + .serverSecurityLogicPure(auth.bearer[String](), statusCode(StatusCode.Unauthorized).and(stringBody))(securityLogic) + .streaming + .addStreamingTool(whoAmITool) + .addPrompt(whoAmIPrompt) + private def assertWhoAmIDeliversPrincipal(client: BidirectionalMcpClient[F]): F[Assertion] = val messages = ConcurrentLinkedQueue[Json]() val listener: ServerNotification => F[Unit] = { @@ -79,6 +93,12 @@ trait SecuredMcpServerStreamingTests[F[_]] extends AsyncFlatSpec with Matchers w it should "give the principal to a streaming tool when the security logic is effectful" in withSecuredStreamingServer(securedStreamingServerWithEffectfulSecurityLogic)(assertWhoAmIDeliversPrincipal) + it should "give the principal to the prompt logic" in + withSecuredStreamingServer(securedStreamingServerWithPrompt): client => + client + .getPrompt("whoAmI") + .map(_.messages shouldBe List(PromptMessage(Role.User, ToolContent.Text(text = "employee@example.com")))) + 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 diff --git a/server/src/test/scala/chimp/server/SecuredMcpServerTests.scala b/server/src/test/scala/chimp/server/SecuredMcpServerTests.scala index 3a4208eb..49738fa0 100644 --- a/server/src/test/scala/chimp/server/SecuredMcpServerTests.scala +++ b/server/src/test/scala/chimp/server/SecuredMcpServerTests.scala @@ -1,7 +1,7 @@ package chimp.server import chimp.client.{McpAuthorizationException, McpClient} -import chimp.protocol.{Implementation, ResourceContents, ToolContent} +import chimp.protocol.{GetPromptResult, Implementation, PromptMessage, ResourceContents, Role, ToolContent} import io.circe.{Codec, Json} import org.scalatest.{Assertion, RecoverMethods} import org.scalatest.flatspec.AsyncFlatSpec @@ -47,6 +47,33 @@ trait SecuredMcpServerTests[F[_]] extends AsyncFlatSpec with Matchers with Recov monad.unit(Right(List(ResourceContents.Text(uri = "test://greeting", text = "hello", mimeType = Some("text/plain"))))) ) + private def whoAmIPrompt: SecuredServerPrompt[F, User] = + prompt("whoAmI") + .description("Greets the caller.") + .argument("name", required = true) + .securedServerLogic[F, User]((args, user, _) => + monad.unit( + GetPromptResult(messages = + List(PromptMessage(Role.User, ToolContent.Text(text = s"Hello ${args.getOrElse("name", "world")} from ${user.email}"))) + ) + ) + ) + + private def whoAmIResource: SecuredServerResource[F, User] = + resource("test://whoami") + .name("whoami") + .mimeType("text/plain") + .securedServerLogic[F, User]((user, _) => + monad.unit(Right(List(ResourceContents.Text(uri = "test://whoami", text = user.email, mimeType = Some("text/plain"))))) + ) + + private def whoAmIResourceTemplate: SecuredServerResourceTemplate[F, User] = + resourceTemplate("test://user/{id}") + .name("user") + .securedServerLogic[F, User]((vars, uri, user, _) => + monad.unit(Right(List(ResourceContents.Text(uri = uri, text = s"${vars("id")} ${user.email}", mimeType = Some("text/plain"))))) + ) + private def securedServer: SecuredMcpServer[F, String, String, User] = McpServer[F]() .addTool(echoTool) @@ -67,6 +94,14 @@ trait SecuredMcpServerTests[F[_]] extends AsyncFlatSpec with Matchers with Recov .addResource(greetingResource) .addTool(whoAmITool) + private def securedServerWithPromptAndResources: SecuredMcpServer[F, String, String, User] = + McpServer[F]() + .addResource(greetingResource) + .serverSecurityLogicPure(auth.bearer[String](), statusCode(StatusCode.Unauthorized).and(stringBody))(securityLogic) + .addPrompts(whoAmIPrompt) + .addResources(whoAmIResource) + .addResourceTemplates(whoAmIResourceTemplate) + private def assertWhoAmIDeliversPrincipal(client: McpClient[F]): F[Assertion] = client .callTool("whoAmI", Json.obj("message" -> Json.fromString("hi"))) @@ -109,6 +144,47 @@ trait SecuredMcpServerTests[F[_]] extends AsyncFlatSpec with Matchers with Recov case ResourceContents.Text(_, text, _, _) => text shouldBe "hello" case other => fail(s"expected text contents, got $other") + it should "give the principal to the prompt logic" in + withSecuredServer(securedServerWithPromptAndResources): client => + client.serverCapabilities.prompts shouldBe defined + client + .getPrompt("whoAmI", Map("name" -> "Ada")) + .map: result => + result.messages shouldBe List( + PromptMessage(Role.User, ToolContent.Text(text = "Hello Ada from employee@example.com")) + ) + + it should "give the principal to the resource logic" in + withSecuredServer(securedServerWithPromptAndResources): client => + client + .listResources() + .flatMap: listed => + listed.resources.map(_.uri) shouldBe List("test://greeting", "test://whoami") + client + .readResource("test://whoami") + .map: result => + result.contents.head match + case ResourceContents.Text(_, text, _, _) => text shouldBe "employee@example.com" + case other => fail(s"expected text contents, got $other") + + it should "also serve a resource which does not need the principal" in + withSecuredServer(securedServerWithPromptAndResources): client => + 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 "give the principal to the resource template logic" in + withSecuredServer(securedServerWithPromptAndResources): client => + client + .readResource("test://user/42") + .map: result => + result.contents.head match + case ResourceContents.Text(_, text, _, _) => text shouldBe "42 employee@example.com" + 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 From f36a856b4d171f4ca2b9f507a30a46d5b8bab0c6 Mon Sep 17 00:00:00 2001 From: Maxime Perocheau Date: Mon, 21 Sep 2026 09:38:39 -0700 Subject: [PATCH 2/3] Address PR feedback --- .../main/scala/chimp/server/McpHandler.scala | 44 ++-- .../main/scala/chimp/server/McpServer.scala | 241 +++++------------- .../src/main/scala/chimp/server/Prompt.scala | 29 ++- .../main/scala/chimp/server/Resource.scala | 62 ++--- .../SecuredMcpServerStreamingTests.scala | 2 +- .../chimp/server/SecuredMcpServerTests.scala | 58 +++-- 6 files changed, 177 insertions(+), 259 deletions(-) diff --git a/server/src/main/scala/chimp/server/McpHandler.scala b/server/src/main/scala/chimp/server/McpHandler.scala index d4121445..38b50ee7 100644 --- a/server/src/main/scala/chimp/server/McpHandler.scala +++ b/server/src/main/scala/chimp/server/McpHandler.scala @@ -29,8 +29,10 @@ enum McpResponse: private[server] class McpHandler[F[_], C <: ServerContext[F]](server: McpServerDef[F, C]): private val logger = LoggerFactory.getLogger(classOf[McpHandler[?, ?]]) private val toolsByName = server.tools.map(tool => tool.name -> tool).toMap - private val hasResources = server.resourceDefinitions.nonEmpty || server.resourceTemplateDefinitions.nonEmpty - private val hasPrompts = server.promptDefinitions.nonEmpty + private val promptsByName = server.prompts.map(prompt => prompt.definition.name -> prompt).toMap + private val resourcesByUri = server.resources.map(resource => resource.definition.uri -> resource).toMap + private val hasResources = server.resources.nonEmpty || server.resourceTemplates.nonEmpty + private val hasPrompts = server.prompts.nonEmpty private val toolDefinitions = server.tools.map(toolToDefinition) private def toJsonSchema(toolSchema: ToolSchema): Json = toolSchema match @@ -82,10 +84,10 @@ private[server] class McpHandler[F[_], C <: ServerContext[F]](server: McpServerD case "tools/call" => handleToolsCall(params, id, headers, makeContext).map(jsonResponse) case "resources/list" if hasResources => - jsonResponse(JSONRPCMessage.Response(id = id, result = ListResourcesResult(server.resourceDefinitions).asJson)).unit + jsonResponse(JSONRPCMessage.Response(id = id, result = ListResourcesResult(server.resources.map(_.definition)).asJson)).unit case "resources/templates/list" if hasResources => jsonResponse( - JSONRPCMessage.Response(id = id, result = ListResourceTemplatesResult(server.resourceTemplateDefinitions).asJson) + JSONRPCMessage.Response(id = id, result = ListResourceTemplatesResult(server.resourceTemplates.map(_.definition)).asJson) ).unit case "resources/read" if hasResources => handleResourcesRead(params, id, headers, makeContext).map(jsonResponse) @@ -94,7 +96,7 @@ private[server] class McpHandler[F[_], C <: ServerContext[F]](server: McpServerD case "resources/unsubscribe" if server.subscriptions.isDefined => handleSubscribe(params, id, subscribe = false).map(jsonResponse) case "prompts/list" if hasPrompts => - jsonResponse(JSONRPCMessage.Response(id = id, result = ListPromptsResult(server.promptDefinitions).asJson)).unit + jsonResponse(JSONRPCMessage.Response(id = id, result = ListPromptsResult(server.prompts.map(_.definition)).asJson)).unit case "prompts/get" if hasPrompts => handlePromptsGet(params, id, headers, makeContext).map(jsonResponse) case "completion/complete" if server.completion.isDefined => @@ -184,15 +186,21 @@ private[server] class McpHandler[F[_], C <: ServerContext[F]](server: McpServerD makeContext: Option[ProgressToken] => C )(using MonadError[F]): F[JSONRPCMessage] = decodeParams[ReadResourceParams](params, id): params => - server.invokeResource(params.uri, makeContext(None), headers) match - case Some(result) => result.map(resourceReadResponse(id, params.uri)) - case None => - protocolError( - id, - JSONRPCErrorCodes.ResourceNotFound.code, - s"Resource not found: ${params.uri}", - Some(Json.obj("uri" -> Json.fromString(params.uri))) - ).unit + val context = makeContext(None) + resourcesByUri.get(params.uri) match + case Some(resource) => resource.read(context, headers).map(resourceReadResponse(id, params.uri)) + case None => + server.resourceTemplates.iterator + .map(template => (template, template.matcher.matchUri(params.uri))) + .collectFirst { case (template, Some(vars)) => template.read(vars, params.uri, context, headers) } match + case Some(result) => result.map(resourceReadResponse(id, params.uri)) + case None => + protocolError( + id, + JSONRPCErrorCodes.ResourceNotFound.code, + s"Resource not found: ${params.uri}", + Some(Json.obj("uri" -> Json.fromString(params.uri))) + ).unit private def decodeParams[P: Decoder](params: Option[Json], id: RequestId)(f: P => F[JSONRPCMessage])(using MonadError[F] @@ -226,9 +234,11 @@ private[server] class McpHandler[F[_], C <: ServerContext[F]](server: McpServerD makeContext: Option[ProgressToken] => C )(using MonadError[F]): F[JSONRPCMessage] = decodeParams[GetPromptParams](params, id): params => - server.invokePrompt(params.name, params.arguments.getOrElse(Map.empty), makeContext(None), headers) match - case Some(result) => - result.map(promptResult => JSONRPCMessage.Response(id = id, result = promptResult.asJson)) + promptsByName.get(params.name) match + case Some(prompt) => + prompt + .logic(params.arguments.getOrElse(Map.empty), makeContext(None), headers) + .map(promptResult => JSONRPCMessage.Response(id = id, result = promptResult.asJson)) case None => protocolError(id, JSONRPCErrorCodes.InvalidParams.code, s"Unknown prompt: ${params.name}").unit private def handleComplete(params: Option[Json], id: RequestId)(using MonadError[F]): F[JSONRPCMessage] = diff --git a/server/src/main/scala/chimp/server/McpServer.scala b/server/src/main/scala/chimp/server/McpServer.scala index fc0ffcde..591caa58 100644 --- a/server/src/main/scala/chimp/server/McpServer.scala +++ b/server/src/main/scala/chimp/server/McpServer.scala @@ -2,13 +2,10 @@ package chimp.server import chimp.protocol.* import chimp.server.transport.{SecuredServerHttpTransport, ServerHttpTransport} -import sttp.model.Header import sttp.monad.MonadError import sttp.tapir.server.ServerEndpoint import sttp.tapir.{EndpointInput, EndpointOutput} -import scala.annotation.{targetName, unused} - type CompletionHandler[F[_]] = (CompleteRef, CompleteArgument, Option[CompleteContext]) => F[Completion] type SetLoggingLevelHandler[F[_]] = LoggingLevel => F[Unit] @@ -18,6 +15,9 @@ case class ResourceSubscriptions[F[_]]( onUnsubscribe: UnsubscribeParams => F[Unit] ) +private def mergeBy[A, K](first: List[A], second: List[A])(key: A => K): List[A] = + (first ++ second).reverse.distinctBy(key).reverse + sealed trait McpServerDef[F[_], C <: ServerContext[F]]: def name: String def version: String @@ -25,48 +25,13 @@ sealed trait McpServerDef[F[_], C <: ServerContext[F]]: def showJsonSchemaMetadata: Boolean def originCheck: OriginCheck def tools: List[ServerTool[?, ?, F, C]] - def prompts: List[ServerPrompt[F]] - def resources: List[ServerResource[F]] - def resourceTemplates: List[ServerResourceTemplate[F]] + def prompts: List[ServerPrompt[F, C]] + def resources: List[ServerResource[F, C]] + def resourceTemplates: List[ServerResourceTemplate[F, C]] def completion: Option[CompletionHandler[F]] def loggingLevel: Option[SetLoggingLevelHandler[F]] def subscriptions: Option[ResourceSubscriptions[F]] - private[server] lazy val promptsByName: Map[String, ServerPrompt[F]] = - prompts.map(prompt => prompt.definition.name -> prompt).toMap - - private[server] lazy val resourcesByUri: Map[String, ServerResource[F]] = - resources.map(resource => resource.definition.uri -> resource).toMap - - private[server] def promptDefinitions: List[Prompt] = prompts.map(_.definition) - - private[server] def invokePrompt( - name: String, - args: Map[String, String], - @unused context: C, - headers: Seq[Header] - ): Option[F[GetPromptResult]] = - promptsByName.get(name).map(_.logic(args, headers)) - - private[server] def resourceDefinitions: List[Resource] = resources.map(_.definition) - - private[server] def resourceTemplateDefinitions: List[ResourceTemplate] = resourceTemplates.map(_.definition) - - private[server] def invokeResource( - uri: String, - @unused context: C, - headers: Seq[Header] - ): Option[F[Either[ResourceError, List[ResourceContents]]]] = - resourcesByUri - .get(uri) - .map(_.read(headers)) - .orElse(matchResourceTemplate(uri).map((template, vars) => template.read(vars, uri, headers))) - - private[server] def matchResourceTemplate(uri: String): Option[(ServerResourceTemplate[F], Map[String, String])] = - resourceTemplates.iterator - .map(template => (template, template.matcher.matchUri(uri))) - .collectFirst { case (template, Some(vars)) => (template, vars) } - case class McpServer[F[_]]( name: String = "Chimp MCP server", version: String = "1.0.0", @@ -74,9 +39,9 @@ case class McpServer[F[_]]( showJsonSchemaMetadata: Boolean = true, originCheck: OriginCheck = OriginCheck.localhostOnly, tools: List[ServerTool[?, ?, F, ServerContext[F]]] = Nil, - prompts: List[ServerPrompt[F]] = Nil, - resources: List[ServerResource[F]] = Nil, - resourceTemplates: List[ServerResourceTemplate[F]] = Nil, + prompts: List[ServerPrompt[F, ServerContext[F]]] = Nil, + resources: List[ServerResource[F, ServerContext[F]]] = Nil, + resourceTemplates: List[ServerResourceTemplate[F, ServerContext[F]]] = Nil, completion: Option[CompletionHandler[F]] = None, loggingLevel: Option[SetLoggingLevelHandler[F]] = None, subscriptions: Option[ResourceSubscriptions[F]] = None @@ -102,22 +67,22 @@ case class McpServer[F[_]]( def addTools(tools: ServerTool[?, ?, F, ServerContext[F]]*): McpServer[F] = copy(tools = this.tools ++ tools) - def addPrompt(prompt: ServerPrompt[F]): McpServer[F] = + def addPrompt(prompt: ServerPrompt[F, ServerContext[F]]): McpServer[F] = copy(prompts = prompts :+ prompt) - def addPrompts(prompts: ServerPrompt[F]*): McpServer[F] = + def addPrompts(prompts: ServerPrompt[F, ServerContext[F]]*): McpServer[F] = copy(prompts = this.prompts ++ prompts) - def addResource(resource: ServerResource[F]): McpServer[F] = + def addResource(resource: ServerResource[F, ServerContext[F]]): McpServer[F] = copy(resources = resources :+ resource) - def addResources(resources: ServerResource[F]*): McpServer[F] = + def addResources(resources: ServerResource[F, ServerContext[F]]*): McpServer[F] = copy(resources = this.resources ++ resources) - def addResourceTemplate(resourceTemplate: ServerResourceTemplate[F]): McpServer[F] = + def addResourceTemplate(resourceTemplate: ServerResourceTemplate[F, ServerContext[F]]): McpServer[F] = copy(resourceTemplates = resourceTemplates :+ resourceTemplate) - def addResourceTemplates(resourceTemplates: ServerResourceTemplate[F]*): McpServer[F] = + def addResourceTemplates(resourceTemplates: ServerResourceTemplate[F, ServerContext[F]]*): McpServer[F] = copy(resourceTemplates = this.resourceTemplates ++ resourceTemplates) def withCompletion(handler: CompletionHandler[F]): McpServer[F] = @@ -168,9 +133,9 @@ case class StreamingMcpServer[F[_]]( showJsonSchemaMetadata: Boolean = true, originCheck: OriginCheck = OriginCheck.localhostOnly, tools: List[ServerTool[?, ?, F, StreamingServerContext[F]]] = Nil, - prompts: List[ServerPrompt[F]] = Nil, - resources: List[ServerResource[F]] = Nil, - resourceTemplates: List[ServerResourceTemplate[F]] = Nil, + prompts: List[ServerPrompt[F, StreamingServerContext[F]]] = Nil, + resources: List[ServerResource[F, StreamingServerContext[F]]] = Nil, + resourceTemplates: List[ServerResourceTemplate[F, StreamingServerContext[F]]] = Nil, completion: Option[CompletionHandler[F]] = None, loggingLevel: Option[SetLoggingLevelHandler[F]] = None, subscriptions: Option[ResourceSubscriptions[F]] = None @@ -202,22 +167,22 @@ case class StreamingMcpServer[F[_]]( def addStreamingTools(tools: ServerTool[?, ?, F, StreamingServerContext[F]]*): StreamingMcpServer[F] = copy(tools = this.tools ++ tools) - def addPrompt(prompt: ServerPrompt[F]): StreamingMcpServer[F] = + def addPrompt(prompt: ServerPrompt[F, ServerContext[F]]): StreamingMcpServer[F] = copy(prompts = prompts :+ prompt) - def addPrompts(prompts: ServerPrompt[F]*): StreamingMcpServer[F] = + def addPrompts(prompts: ServerPrompt[F, ServerContext[F]]*): StreamingMcpServer[F] = copy(prompts = this.prompts ++ prompts) - def addResource(resource: ServerResource[F]): StreamingMcpServer[F] = + def addResource(resource: ServerResource[F, ServerContext[F]]): StreamingMcpServer[F] = copy(resources = resources :+ resource) - def addResources(resources: ServerResource[F]*): StreamingMcpServer[F] = + def addResources(resources: ServerResource[F, ServerContext[F]]*): StreamingMcpServer[F] = copy(resources = this.resources ++ resources) - def addResourceTemplate(resourceTemplate: ServerResourceTemplate[F]): StreamingMcpServer[F] = + def addResourceTemplate(resourceTemplate: ServerResourceTemplate[F, ServerContext[F]]): StreamingMcpServer[F] = copy(resourceTemplates = resourceTemplates :+ resourceTemplate) - def addResourceTemplates(resourceTemplates: ServerResourceTemplate[F]*): StreamingMcpServer[F] = + def addResourceTemplates(resourceTemplates: ServerResourceTemplate[F, ServerContext[F]]*): StreamingMcpServer[F] = copy(resourceTemplates = this.resourceTemplates ++ resourceTemplates) def withCompletion(handler: CompletionHandler[F]): StreamingMcpServer[F] = @@ -246,64 +211,29 @@ case class SecuredMcpServer[F[_], S, E, P]( errorOutput: EndpointOutput[E], securityLogic: MonadError[F] => S => F[Either[E, P]], securedTools: List[ServerTool[?, ?, F, SecuredServerContext[F, P]]] = Nil, - securedPrompts: List[SecuredServerPrompt[F, P]] = Nil, - securedResources: List[SecuredServerResource[F, P]] = Nil, - securedResourceTemplates: List[SecuredServerResourceTemplate[F, P]] = Nil + securedPrompts: List[ServerPrompt[F, SecuredServerContext[F, P]]] = Nil, + securedResources: List[ServerResource[F, SecuredServerContext[F, P]]] = Nil, + securedResourceTemplates: List[ServerResourceTemplate[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 - private lazy val securedPromptsByName: Map[String, SecuredServerPrompt[F, P]] = - securedPrompts.map(prompt => prompt.definition.name -> prompt).toMap - - private lazy val securedResourcesByUri: Map[String, SecuredServerResource[F, P]] = - securedResources.map(resource => resource.definition.uri -> resource).toMap - - override private[server] def promptDefinitions: List[Prompt] = - prompts.map(_.definition) ++ securedPrompts.map(_.definition) - - override private[server] def invokePrompt( - name: String, - args: Map[String, String], - context: SecuredServerContext[F, P], - headers: Seq[Header] - ): Option[F[GetPromptResult]] = - promptsByName - .get(name) - .map(_.logic(args, headers)) - .orElse(securedPromptsByName.get(name).map(_.logic(args, context.principal, headers))) - - override private[server] def resourceDefinitions: List[Resource] = - resources.map(_.definition) ++ securedResources.map(_.definition) - - override private[server] def resourceTemplateDefinitions: List[ResourceTemplate] = - resourceTemplates.map(_.definition) ++ securedResourceTemplates.map(_.definition) - - override private[server] def invokeResource( - uri: String, - context: SecuredServerContext[F, P], - headers: Seq[Header] - ): Option[F[Either[ResourceError, List[ResourceContents]]]] = - resourcesByUri - .get(uri) - .map(_.read(headers)) - .orElse(securedResourcesByUri.get(uri).map(_.read(context.principal, headers))) - .orElse(matchResourceTemplate(uri).map((template, vars) => template.read(vars, uri, headers))) - .orElse: - securedResourceTemplates.iterator - .map(template => (template, template.matcher.matchUri(uri))) - .collectFirst { case (template, Some(vars)) => template.read(vars, uri, context.principal, headers) } + def prompts: List[ServerPrompt[F, SecuredServerContext[F, P]]] = + mergeBy(server.prompts, securedPrompts)(_.definition.name) + + def resources: List[ServerResource[F, SecuredServerContext[F, P]]] = + mergeBy(server.resources, securedResources)(_.definition.uri) + + def resourceTemplates: List[ServerResourceTemplate[F, SecuredServerContext[F, P]]] = + mergeBy(server.resourceTemplates, securedResourceTemplates)(_.definition.uriTemplate) def name(value: String): SecuredMcpServer[F, S, E, P] = copy(server = server.name(value)) @@ -326,43 +256,26 @@ case class SecuredMcpServer[F[_], S, E, P]( 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 addPrompt(prompt: SecuredServerPrompt[F, P]): SecuredMcpServer[F, S, E, P] = + def addPrompt(prompt: ServerPrompt[F, SecuredServerContext[F, P]]): SecuredMcpServer[F, S, E, P] = copy(securedPrompts = securedPrompts :+ prompt) - def addPrompts(prompts: ServerPrompt[F]*): SecuredMcpServer[F, S, E, P] = - copy(server = server.addPrompts(prompts*)) - - @targetName("addSecuredPrompts") - def addPrompts(prompts: SecuredServerPrompt[F, P]*): SecuredMcpServer[F, S, E, P] = + def addPrompts(prompts: ServerPrompt[F, SecuredServerContext[F, P]]*): SecuredMcpServer[F, S, E, P] = copy(securedPrompts = securedPrompts ++ prompts) - def addResource(resource: ServerResource[F]): SecuredMcpServer[F, S, E, P] = - copy(server = server.addResource(resource)) - - def addResource(resource: SecuredServerResource[F, P]): SecuredMcpServer[F, S, E, P] = + def addResource(resource: ServerResource[F, SecuredServerContext[F, P]]): SecuredMcpServer[F, S, E, P] = copy(securedResources = securedResources :+ resource) - def addResources(resources: ServerResource[F]*): SecuredMcpServer[F, S, E, P] = - copy(server = server.addResources(resources*)) - - @targetName("addSecuredResources") - def addResources(resources: SecuredServerResource[F, P]*): SecuredMcpServer[F, S, E, P] = + def addResources(resources: ServerResource[F, SecuredServerContext[F, P]]*): SecuredMcpServer[F, S, E, P] = copy(securedResources = securedResources ++ resources) - def addResourceTemplate(resourceTemplate: ServerResourceTemplate[F]): SecuredMcpServer[F, S, E, P] = - copy(server = server.addResourceTemplate(resourceTemplate)) - - def addResourceTemplate(resourceTemplate: SecuredServerResourceTemplate[F, P]): SecuredMcpServer[F, S, E, P] = + def addResourceTemplate( + resourceTemplate: ServerResourceTemplate[F, SecuredServerContext[F, P]] + ): SecuredMcpServer[F, S, E, P] = copy(securedResourceTemplates = securedResourceTemplates :+ resourceTemplate) - def addResourceTemplates(resourceTemplates: ServerResourceTemplate[F]*): SecuredMcpServer[F, S, E, P] = - copy(server = server.addResourceTemplates(resourceTemplates*)) - - @targetName("addSecuredResourceTemplates") - def addResourceTemplates(resourceTemplates: SecuredServerResourceTemplate[F, P]*): SecuredMcpServer[F, S, E, P] = + def addResourceTemplates( + resourceTemplates: ServerResourceTemplate[F, SecuredServerContext[F, P]]* + ): SecuredMcpServer[F, S, E, P] = copy(securedResourceTemplates = securedResourceTemplates ++ resourceTemplates) def withCompletion(handler: CompletionHandler[F]): SecuredMcpServer[F, S, E, P] = @@ -390,9 +303,9 @@ case class SecuredStreamingMcpServer[F[_], S, E, P]( 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 prompts: List[ServerPrompt[F, SecuredStreamingServerContext[F, P]]] = server.prompts + def resources: List[ServerResource[F, SecuredStreamingServerContext[F, P]]] = server.resources + def resourceTemplates: List[ServerResourceTemplate[F, SecuredStreamingServerContext[F, P]]] = server.resourceTemplates def completion: Option[CompletionHandler[F]] = server.completion def loggingLevel: Option[SetLoggingLevelHandler[F]] = server.loggingLevel def subscriptions: Option[ResourceSubscriptions[F]] = server.subscriptions @@ -403,27 +316,6 @@ case class SecuredStreamingMcpServer[F[_], S, E, P]( def tools: List[ServerTool[?, ?, F, SecuredStreamingServerContext[F, P]]] = server.tools ++ streamingTools - override private[server] def promptDefinitions: List[Prompt] = server.promptDefinitions - - override private[server] def invokePrompt( - name: String, - args: Map[String, String], - context: SecuredStreamingServerContext[F, P], - headers: Seq[Header] - ): Option[F[GetPromptResult]] = - server.invokePrompt(name, args, context, headers) - - override private[server] def resourceDefinitions: List[Resource] = server.resourceDefinitions - - override private[server] def resourceTemplateDefinitions: List[ResourceTemplate] = server.resourceTemplateDefinitions - - override private[server] def invokeResource( - uri: String, - context: SecuredStreamingServerContext[F, P], - headers: Seq[Header] - ): Option[F[Either[ResourceError, List[ResourceContents]]]] = - server.invokeResource(uri, context, headers) - def name(value: String): SecuredStreamingMcpServer[F, S, E, P] = copy(server = server.name(value)) @@ -451,43 +343,26 @@ case class SecuredStreamingMcpServer[F[_], S, E, P]( 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 addPrompt(prompt: SecuredServerPrompt[F, P]): SecuredStreamingMcpServer[F, S, E, P] = + def addPrompt(prompt: ServerPrompt[F, SecuredServerContext[F, P]]): SecuredStreamingMcpServer[F, S, E, P] = copy(server = server.addPrompt(prompt)) - def addPrompts(prompts: ServerPrompt[F]*): SecuredStreamingMcpServer[F, S, E, P] = + def addPrompts(prompts: ServerPrompt[F, SecuredServerContext[F, P]]*): SecuredStreamingMcpServer[F, S, E, P] = copy(server = server.addPrompts(prompts*)) - @targetName("addSecuredPrompts") - def addPrompts(prompts: SecuredServerPrompt[F, P]*): SecuredStreamingMcpServer[F, S, E, P] = - copy(server = server.addPrompts(prompts*)) - - def addResource(resource: ServerResource[F]): SecuredStreamingMcpServer[F, S, E, P] = + def addResource(resource: ServerResource[F, SecuredServerContext[F, P]]): SecuredStreamingMcpServer[F, S, E, P] = copy(server = server.addResource(resource)) - def addResource(resource: SecuredServerResource[F, P]): 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*)) - - @targetName("addSecuredResources") - def addResources(resources: SecuredServerResource[F, P]*): SecuredStreamingMcpServer[F, S, E, P] = + def addResources(resources: ServerResource[F, SecuredServerContext[F, P]]*): 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 addResourceTemplate(resourceTemplate: SecuredServerResourceTemplate[F, P]): SecuredStreamingMcpServer[F, S, E, P] = + def addResourceTemplate( + resourceTemplate: ServerResourceTemplate[F, SecuredServerContext[F, P]] + ): 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*)) - - @targetName("addSecuredResourceTemplates") - def addResourceTemplates(resourceTemplates: SecuredServerResourceTemplate[F, P]*): SecuredStreamingMcpServer[F, S, E, P] = + def addResourceTemplates( + resourceTemplates: ServerResourceTemplate[F, SecuredServerContext[F, P]]* + ): SecuredStreamingMcpServer[F, S, E, P] = copy(server = server.addResourceTemplates(resourceTemplates*)) def withCompletion(handler: CompletionHandler[F]): SecuredStreamingMcpServer[F, S, E, P] = diff --git a/server/src/main/scala/chimp/server/Prompt.scala b/server/src/main/scala/chimp/server/Prompt.scala index bacc1dfb..8bbbc8c5 100644 --- a/server/src/main/scala/chimp/server/Prompt.scala +++ b/server/src/main/scala/chimp/server/Prompt.scala @@ -29,25 +29,31 @@ case class PartialPrompt( copy(arguments = arguments ++ args) /** Attaches effectful logic, with access to the request headers, producing the prompt's messages. */ - def serverLogic[F[_]](logic: (Map[String, String], Seq[Header]) => F[GetPromptResult]): ServerPrompt[F] = - ServerPrompt(definition, logic) + def serverLogic[F[_]]( + logic: (Map[String, String], Seq[Header]) => F[GetPromptResult] + ): ServerPrompt[F, ServerContext[F]] = + ServerPrompt(definition, (args, _, headers) => logic(args, headers)) /** Attaches effectful logic, with access to the principal; usable only on a [[SecuredMcpServer]]. */ def securedServerLogic[F[_], P]( logic: (Map[String, String], P, Seq[Header]) => F[GetPromptResult] - ): SecuredServerPrompt[F, P] = - SecuredServerPrompt(definition, logic) + ): ServerPrompt[F, SecuredServerContext[F, P]] = + ServerPrompt(definition, (args, context, headers) => logic(args, context.principal, headers)) /** Attaches synchronous logic that also receives the request headers. */ - def handleWithHeaders(logic: (Map[String, String], Seq[Header]) => GetPromptResult): ServerPrompt[Identity] = - ServerPrompt(definition, logic) + def handleWithHeaders( + logic: (Map[String, String], Seq[Header]) => GetPromptResult + ): ServerPrompt[Identity, ServerContext[Identity]] = + serverLogic[Identity](logic) /** Attaches synchronous logic over just the supplied argument values. */ - def handle(logic: Map[String, String] => GetPromptResult): ServerPrompt[Identity] = + def handle(logic: Map[String, String] => GetPromptResult): ServerPrompt[Identity, ServerContext[Identity]] = handleWithHeaders((args, _) => logic(args)) /** Attaches synchronous logic over the supplied argument values and the principal; usable only on a [[SecuredMcpServer]]. */ - def handleSecured[P](logic: (Map[String, String], P) => GetPromptResult): SecuredServerPrompt[Identity, P] = + def handleSecured[P]( + logic: (Map[String, String], P) => GetPromptResult + ): ServerPrompt[Identity, SecuredServerContext[Identity, P]] = securedServerLogic[Identity, P]((args, principal, _) => logic(args, principal)) private def definition: Prompt = @@ -56,10 +62,7 @@ case class PartialPrompt( end PartialPrompt /** A fully-defined prompt: its metadata plus the logic producing its messages. */ -case class ServerPrompt[F[_]](definition: Prompt, logic: (Map[String, String], Seq[Header]) => F[GetPromptResult]) - -/** A prompt whose logic also receives the principal made by a [[SecuredMcpServer]]. */ -case class SecuredServerPrompt[F[_], P]( +case class ServerPrompt[F[_], -C <: ServerContext[F]]( definition: Prompt, - logic: (Map[String, String], P, Seq[Header]) => F[GetPromptResult] + logic: (Map[String, String], C, Seq[Header]) => F[GetPromptResult] ) diff --git a/server/src/main/scala/chimp/server/Resource.scala b/server/src/main/scala/chimp/server/Resource.scala index 56901906..ec37ce84 100644 --- a/server/src/main/scala/chimp/server/Resource.scala +++ b/server/src/main/scala/chimp/server/Resource.scala @@ -41,25 +41,31 @@ case class PartialResource( copy(size = Some(value)) /** Attaches effectful logic, with access to the request headers, producing the resource's contents (or an error). */ - def serverLogic[F[_]](logic: Seq[Header] => F[Either[ResourceError, List[ResourceContents]]]): ServerResource[F] = - ServerResource(definition, logic) + def serverLogic[F[_]]( + logic: Seq[Header] => F[Either[ResourceError, List[ResourceContents]]] + ): ServerResource[F, ServerContext[F]] = + ServerResource(definition, (_, headers) => logic(headers)) /** Attaches effectful logic, with access to the principal; usable only on a [[SecuredMcpServer]]. */ def securedServerLogic[F[_], P]( logic: (P, Seq[Header]) => F[Either[ResourceError, List[ResourceContents]]] - ): SecuredServerResource[F, P] = - SecuredServerResource(definition, logic) + ): ServerResource[F, SecuredServerContext[F, P]] = + ServerResource(definition, (context, headers) => logic(context.principal, headers)) /** Attaches synchronous logic that also receives the request headers. */ - def handleWithHeaders(logic: Seq[Header] => Either[ResourceError, List[ResourceContents]]): ServerResource[Identity] = - ServerResource(definition, logic) + def handleWithHeaders( + logic: Seq[Header] => Either[ResourceError, List[ResourceContents]] + ): ServerResource[Identity, ServerContext[Identity]] = + serverLogic[Identity](logic) /** Attaches synchronous logic producing the resource's contents (or an error). */ - def handle(logic: () => Either[ResourceError, List[ResourceContents]]): ServerResource[Identity] = + def handle(logic: () => Either[ResourceError, List[ResourceContents]]): ServerResource[Identity, ServerContext[Identity]] = handleWithHeaders(_ => logic()) /** Attaches synchronous logic over the principal; usable only on a [[SecuredMcpServer]]. */ - def handleSecured[P](logic: P => Either[ResourceError, List[ResourceContents]]): SecuredServerResource[Identity, P] = + def handleSecured[P]( + logic: P => Either[ResourceError, List[ResourceContents]] + ): ServerResource[Identity, SecuredServerContext[Identity, P]] = securedServerLogic[Identity, P]((principal, _) => logic(principal)) private def definition: Resource = Resource(uri, name.getOrElse(uri), title, description, mimeType, size) @@ -67,12 +73,9 @@ case class PartialResource( end PartialResource /** A fully-defined resource: its metadata plus the logic reading its contents. */ -case class ServerResource[F[_]](definition: Resource, read: Seq[Header] => F[Either[ResourceError, List[ResourceContents]]]) - -/** A resource whose read logic also receives the principal made by a [[SecuredMcpServer]]. */ -case class SecuredServerResource[F[_], P]( +case class ServerResource[F[_], -C <: ServerContext[F]]( definition: Resource, - read: (P, Seq[Header]) => F[Either[ResourceError, List[ResourceContents]]] + read: (C, Seq[Header]) => F[Either[ResourceError, List[ResourceContents]]] ) /** A resource template being defined, before its read logic is attached. */ @@ -98,29 +101,35 @@ case class PartialResourceTemplate( /** Attaches effectful logic reading a matched URI; receives the extracted variables, the full URI, and the request headers. */ def serverLogic[F[_]]( logic: (Map[String, String], String, Seq[Header]) => F[Either[ResourceError, List[ResourceContents]]] - ): ServerResourceTemplate[F] = - ServerResourceTemplate(definition, UriTemplate.compile(uriTemplate), logic) + ): ServerResourceTemplate[F, ServerContext[F]] = + ServerResourceTemplate(definition, UriTemplate.compile(uriTemplate), (vars, uri, _, headers) => logic(vars, uri, headers)) /** Attaches effectful logic, with access to the principal; usable only on a [[SecuredMcpServer]]. */ def securedServerLogic[F[_], P]( logic: (Map[String, String], String, P, Seq[Header]) => F[Either[ResourceError, List[ResourceContents]]] - ): SecuredServerResourceTemplate[F, P] = - SecuredServerResourceTemplate(definition, UriTemplate.compile(uriTemplate), logic) + ): ServerResourceTemplate[F, SecuredServerContext[F, P]] = + ServerResourceTemplate( + definition, + UriTemplate.compile(uriTemplate), + (vars, uri, context, headers) => logic(vars, uri, context.principal, headers) + ) /** Attaches synchronous logic that also receives the request headers. */ def handleWithHeaders( logic: (Map[String, String], String, Seq[Header]) => Either[ResourceError, List[ResourceContents]] - ): ServerResourceTemplate[Identity] = - ServerResourceTemplate(definition, UriTemplate.compile(uriTemplate), logic) + ): ServerResourceTemplate[Identity, ServerContext[Identity]] = + serverLogic[Identity](logic) /** Attaches synchronous logic receiving the extracted variables and the full URI. */ - def handle(logic: (Map[String, String], String) => Either[ResourceError, List[ResourceContents]]): ServerResourceTemplate[Identity] = + def handle( + logic: (Map[String, String], String) => Either[ResourceError, List[ResourceContents]] + ): ServerResourceTemplate[Identity, ServerContext[Identity]] = handleWithHeaders((vars, uri, _) => logic(vars, uri)) /** Attaches synchronous logic over the extracted variables, the full URI, and the principal; usable only on a [[SecuredMcpServer]]. */ def handleSecured[P]( logic: (Map[String, String], String, P) => Either[ResourceError, List[ResourceContents]] - ): SecuredServerResourceTemplate[Identity, P] = + ): ServerResourceTemplate[Identity, SecuredServerContext[Identity, P]] = securedServerLogic[Identity, P]((vars, uri, principal, _) => logic(vars, uri, principal)) private def definition: ResourceTemplate = @@ -129,17 +138,10 @@ case class PartialResourceTemplate( end PartialResourceTemplate /** A fully-defined resource template: its metadata, a compiled URI matcher, and the logic reading matched URIs. */ -case class ServerResourceTemplate[F[_]]( - definition: ResourceTemplate, - matcher: UriTemplate, - read: (Map[String, String], String, Seq[Header]) => F[Either[ResourceError, List[ResourceContents]]] -) - -/** A resource template whose read logic also receives the principal made by a [[SecuredMcpServer]]. */ -case class SecuredServerResourceTemplate[F[_], P]( +case class ServerResourceTemplate[F[_], -C <: ServerContext[F]]( definition: ResourceTemplate, matcher: UriTemplate, - read: (Map[String, String], String, P, Seq[Header]) => F[Either[ResourceError, List[ResourceContents]]] + read: (Map[String, String], String, C, Seq[Header]) => F[Either[ResourceError, List[ResourceContents]]] ) /** A compiled URI template that matches concrete URIs and extracts their `{variable}` values. */ diff --git a/server/src/test/scala/chimp/server/SecuredMcpServerStreamingTests.scala b/server/src/test/scala/chimp/server/SecuredMcpServerStreamingTests.scala index 785d9690..750f7202 100644 --- a/server/src/test/scala/chimp/server/SecuredMcpServerStreamingTests.scala +++ b/server/src/test/scala/chimp/server/SecuredMcpServerStreamingTests.scala @@ -59,7 +59,7 @@ trait SecuredMcpServerStreamingTests[F[_]] extends AsyncFlatSpec with Matchers w .streaming .addStreamingTool(whoAmITool) - private def whoAmIPrompt: SecuredServerPrompt[F, User] = + private def whoAmIPrompt: ServerPrompt[F, SecuredServerContext[F, User]] = prompt("whoAmI") .description("Greets the caller") .securedServerLogic[F, User]((_, user, _) => diff --git a/server/src/test/scala/chimp/server/SecuredMcpServerTests.scala b/server/src/test/scala/chimp/server/SecuredMcpServerTests.scala index 49738fa0..50bf3c2d 100644 --- a/server/src/test/scala/chimp/server/SecuredMcpServerTests.scala +++ b/server/src/test/scala/chimp/server/SecuredMcpServerTests.scala @@ -39,7 +39,7 @@ trait SecuredMcpServerTests[F[_]] extends AsyncFlatSpec with Matchers with Recov .input[EchoInput] .securedServerLogic[F, User]((in, user, _) => monad.unit(ToolResult.text(s"${in.message} ${user.email}"))) - private def greetingResource: ServerResource[F] = + private def greetingResource: ServerResource[F, ServerContext[F]] = resource("test://greeting") .name("greeting") .mimeType("text/plain") @@ -47,7 +47,19 @@ trait SecuredMcpServerTests[F[_]] extends AsyncFlatSpec with Matchers with Recov monad.unit(Right(List(ResourceContents.Text(uri = "test://greeting", text = "hello", mimeType = Some("text/plain"))))) ) - private def whoAmIPrompt: SecuredServerPrompt[F, User] = + private def publicWhoAmIPrompt: ServerPrompt[F, ServerContext[F]] = + prompt("whoAmI") + .serverLogic[F]((_, _) => monad.unit(GetPromptResult(messages = List(PromptMessage(Role.User, ToolContent.Text(text = "public")))))) + + private def publicWhoAmIResource: ServerResource[F, ServerContext[F]] = + resource("test://whoami") + .serverLogic[F](_ => monad.unit(Right(List(ResourceContents.Text(uri = "test://whoami", text = "public"))))) + + private def publicWhoAmIResourceTemplate: ServerResourceTemplate[F, ServerContext[F]] = + resourceTemplate("test://user/{id}") + .serverLogic[F]((_, uri, _) => monad.unit(Right(List(ResourceContents.Text(uri = uri, text = "public"))))) + + private def whoAmIPrompt: ServerPrompt[F, SecuredServerContext[F, User]] = prompt("whoAmI") .description("Greets the caller.") .argument("name", required = true) @@ -59,7 +71,7 @@ trait SecuredMcpServerTests[F[_]] extends AsyncFlatSpec with Matchers with Recov ) ) - private def whoAmIResource: SecuredServerResource[F, User] = + private def whoAmIResource: ServerResource[F, SecuredServerContext[F, User]] = resource("test://whoami") .name("whoami") .mimeType("text/plain") @@ -67,7 +79,7 @@ trait SecuredMcpServerTests[F[_]] extends AsyncFlatSpec with Matchers with Recov monad.unit(Right(List(ResourceContents.Text(uri = "test://whoami", text = user.email, mimeType = Some("text/plain"))))) ) - private def whoAmIResourceTemplate: SecuredServerResourceTemplate[F, User] = + private def whoAmIResourceTemplate: ServerResourceTemplate[F, SecuredServerContext[F, User]] = resourceTemplate("test://user/{id}") .name("user") .securedServerLogic[F, User]((vars, uri, user, _) => @@ -96,7 +108,9 @@ trait SecuredMcpServerTests[F[_]] extends AsyncFlatSpec with Matchers with Recov private def securedServerWithPromptAndResources: SecuredMcpServer[F, String, String, User] = McpServer[F]() - .addResource(greetingResource) + .addPrompts(publicWhoAmIPrompt) + .addResources(greetingResource, publicWhoAmIResource) + .addResourceTemplates(publicWhoAmIResourceTemplate) .serverSecurityLogicPure(auth.bearer[String](), statusCode(StatusCode.Unauthorized).and(stringBody))(securityLogic) .addPrompts(whoAmIPrompt) .addResources(whoAmIResource) @@ -130,6 +144,12 @@ trait SecuredMcpServerTests[F[_]] extends AsyncFlatSpec with Matchers with Recov withSecuredServer(serverConfiguredAfterSecurity): client => monad.unit(client.serverInfo shouldBe Implementation("secured-server", "2.0.0")) + it should "expose complete prompt and resource collections" in Future { + securedServerWithPromptAndResources.prompts.map(_.definition.name) shouldBe List("whoAmI") + securedServerWithPromptAndResources.resources.map(_.definition.uri) shouldBe List("test://greeting", "test://whoami") + securedServerWithPromptAndResources.resourceTemplates.map(_.definition.uriTemplate) shouldBe List("test://user/{id}") + } + it should "serve a resource which was added after the security logic" in withSecuredServer(serverConfiguredAfterSecurity): client => client.serverCapabilities.resources shouldBe defined @@ -148,11 +168,15 @@ trait SecuredMcpServerTests[F[_]] extends AsyncFlatSpec with Matchers with Recov withSecuredServer(securedServerWithPromptAndResources): client => client.serverCapabilities.prompts shouldBe defined client - .getPrompt("whoAmI", Map("name" -> "Ada")) - .map: result => - result.messages shouldBe List( - PromptMessage(Role.User, ToolContent.Text(text = "Hello Ada from employee@example.com")) - ) + .listPrompts() + .flatMap: listed => + listed.prompts.map(_.name) shouldBe List("whoAmI") + client + .getPrompt("whoAmI", Map("name" -> "Ada")) + .map: result => + result.messages shouldBe List( + PromptMessage(Role.User, ToolContent.Text(text = "Hello Ada from employee@example.com")) + ) it should "give the principal to the resource logic" in withSecuredServer(securedServerWithPromptAndResources): client => @@ -179,11 +203,15 @@ trait SecuredMcpServerTests[F[_]] extends AsyncFlatSpec with Matchers with Recov it should "give the principal to the resource template logic" in withSecuredServer(securedServerWithPromptAndResources): client => client - .readResource("test://user/42") - .map: result => - result.contents.head match - case ResourceContents.Text(_, text, _, _) => text shouldBe "42 employee@example.com" - case other => fail(s"expected text contents, got $other") + .listResourceTemplates() + .flatMap: listed => + listed.resourceTemplates.map(_.uriTemplate) shouldBe List("test://user/{id}") + client + .readResource("test://user/42") + .map: result => + result.contents.head match + case ResourceContents.Text(_, text, _, _) => text shouldBe "42 employee@example.com" + 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] { From debed8510a0648ea0437145e67124dd96eed112c Mon Sep 17 00:00:00 2001 From: Maxime Perocheau Date: Mon, 21 Sep 2026 10:49:34 -0700 Subject: [PATCH 3/3] Address PR feedback --- server/src/main/scala/chimp/server/McpServer.scala | 6 ++++-- .../scala/chimp/server/SecuredMcpServerTests.scala | 13 ++++++++++--- 2 files changed, 14 insertions(+), 5 deletions(-) diff --git a/server/src/main/scala/chimp/server/McpServer.scala b/server/src/main/scala/chimp/server/McpServer.scala index 591caa58..c994f3d6 100644 --- a/server/src/main/scala/chimp/server/McpServer.scala +++ b/server/src/main/scala/chimp/server/McpServer.scala @@ -224,7 +224,8 @@ case class SecuredMcpServer[F[_], S, E, P]( 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 tools: List[ServerTool[?, ?, F, SecuredServerContext[F, P]]] = + mergeBy(server.tools, securedTools)(_.name) def prompts: List[ServerPrompt[F, SecuredServerContext[F, P]]] = mergeBy(server.prompts, securedPrompts)(_.definition.name) @@ -314,7 +315,8 @@ case class SecuredStreamingMcpServer[F[_], S, E, P]( 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 tools: List[ServerTool[?, ?, F, SecuredStreamingServerContext[F, P]]] = + mergeBy(server.tools, streamingTools)(_.name) def name(value: String): SecuredStreamingMcpServer[F, S, E, P] = copy(server = server.name(value)) diff --git a/server/src/test/scala/chimp/server/SecuredMcpServerTests.scala b/server/src/test/scala/chimp/server/SecuredMcpServerTests.scala index 50bf3c2d..97d153d1 100644 --- a/server/src/test/scala/chimp/server/SecuredMcpServerTests.scala +++ b/server/src/test/scala/chimp/server/SecuredMcpServerTests.scala @@ -33,6 +33,12 @@ trait SecuredMcpServerTests[F[_]] extends AsyncFlatSpec with Matchers with Recov .input[EchoInput] .serverLogic[F]((in, _) => monad.unit(ToolResult.text(in.message))) + private def publicWhoAmITool: ServerTool[EchoInput, NoStructuredOutput, F, ServerContext[F]] = + tool("whoAmI") + .description("Echoes a message.") + .input[EchoInput] + .serverLogic[F]((in, _) => monad.unit(ToolResult.text(s"${in.message} public"))) + private def whoAmITool: ServerTool[EchoInput, NoStructuredOutput, F, SecuredServerContext[F, User]] = tool("whoAmI") .description("Echoes a message and the caller's email.") @@ -88,13 +94,13 @@ trait SecuredMcpServerTests[F[_]] extends AsyncFlatSpec with Matchers with Recov private def securedServer: SecuredMcpServer[F, String, String, User] = McpServer[F]() - .addTool(echoTool) + .addTools(echoTool, publicWhoAmITool) .serverSecurityLogicPure(auth.bearer[String](), statusCode(StatusCode.Unauthorized).and(stringBody))(securityLogic) .addTool(whoAmITool) private def securedServerWithEffectfulSecurityLogic: SecuredMcpServer[F, String, String, User] = McpServer[F]() - .addTool(echoTool) + .addTools(echoTool, publicWhoAmITool) .serverSecurityLogic(auth.bearer[String](), statusCode(StatusCode.Unauthorized).and(stringBody))(securityLogicEffectful) .addTool(whoAmITool) @@ -134,7 +140,7 @@ trait SecuredMcpServerTests[F[_]] extends AsyncFlatSpec with Matchers with Recov client .listTools() .flatMap: tools => - tools.tools.map(_.name) should contain allOf ("echo", "whoAmI") + tools.tools.map(_.name) shouldBe List("echo", "whoAmI") client .callTool("echo", Json.obj("message" -> Json.fromString("hi"))) .map: result => @@ -145,6 +151,7 @@ trait SecuredMcpServerTests[F[_]] extends AsyncFlatSpec with Matchers with Recov monad.unit(client.serverInfo shouldBe Implementation("secured-server", "2.0.0")) it should "expose complete prompt and resource collections" in Future { + securedServer.tools.map(_.name) shouldBe List("echo", "whoAmI") securedServerWithPromptAndResources.prompts.map(_.definition.name) shouldBe List("whoAmI") securedServerWithPromptAndResources.resources.map(_.definition.uri) shouldBe List("test://greeting", "test://whoami") securedServerWithPromptAndResources.resourceTemplates.map(_.definition.uriTemplate) shouldBe List("test://user/{id}")