diff --git a/docs/server/prompts.md b/docs/server/prompts.md index 34128d3..0e818c1 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 7bbe3e8..86f823f 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 c63f21a..d4b7dd6 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 585dcee..38b50ee 100644 --- a/server/src/main/scala/chimp/server/McpHandler.scala +++ b/server/src/main/scala/chimp/server/McpHandler.scala @@ -32,6 +32,7 @@ private[server] class McpHandler[F[_], C <: ServerContext[F]](server: McpServerD 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 @@ -89,15 +90,15 @@ private[server] class McpHandler[F[_], C <: ServerContext[F]](server: McpServerD JSONRPCMessage.Response(id = id, result = ListResourceTemplatesResult(server.resourceTemplates.map(_.definition)).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 => + case "prompts/list" if hasPrompts => 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/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 +124,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,17 +179,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 => + val context = makeContext(None) resourcesByUri.get(params.uri) match - case Some(resource) => resource.read(headers).map(resourceReadResponse(id, params.uri)) + case Some(resource) => resource.read(context, headers).map(resourceReadResponse(id, params.uri)) case None => - val templateMatch = server.resourceTemplates.iterator + 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 => + .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, @@ -221,13 +227,18 @@ 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)) + .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 49f380c..c994f3d 100644 --- a/server/src/main/scala/chimp/server/McpServer.scala +++ b/server/src/main/scala/chimp/server/McpServer.scala @@ -15,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 @@ -22,9 +25,9 @@ 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]] @@ -36,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 @@ -64,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] = @@ -94,7 +97,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]] @@ -130,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 @@ -164,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] = @@ -192,8 +195,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,21 +210,31 @@ 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[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 + 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) + + 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)) @@ -244,23 +257,27 @@ 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: 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*)) + 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: 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*)) + 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: 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*)) + 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] = copy(server = server.withCompletion(handler)) @@ -287,9 +304,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 @@ -298,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)) @@ -327,22 +345,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] = + 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*)) - 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 addResources(resources: ServerResource[F]*): 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] = + 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] = + 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 6e42081..8bbbc8c 100644 --- a/server/src/main/scala/chimp/server/Prompt.scala +++ b/server/src/main/scala/chimp/server/Prompt.scala @@ -29,21 +29,40 @@ 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] + ): 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 + ): ServerPrompt[Identity, SecuredServerContext[Identity, P]] = + securedServerLogic[Identity, P]((args, principal, _) => logic(args, principal)) + private def definition: Prompt = Prompt(name, title, description, Option.when(arguments.nonEmpty)(arguments)) 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]) +case class ServerPrompt[F[_], -C <: ServerContext[F]]( + definition: Prompt, + 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 60ec2d2..ec37ce8 100644 --- a/server/src/main/scala/chimp/server/Resource.scala +++ b/server/src/main/scala/chimp/server/Resource.scala @@ -41,23 +41,42 @@ 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]]] + ): 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]] + ): ServerResource[Identity, SecuredServerContext[Identity, P]] = + securedServerLogic[Identity, P]((principal, _) => logic(principal)) + private def definition: Resource = Resource(uri, name.getOrElse(uri), title, description, mimeType, size) 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]]]) +case class ServerResource[F[_], -C <: ServerContext[F]]( + definition: Resource, + read: (C, Seq[Header]) => F[Either[ResourceError, List[ResourceContents]]] +) /** A resource template being defined, before its read logic is attached. */ case class PartialResourceTemplate( @@ -82,29 +101,47 @@ 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]]] + ): 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]] + ): ServerResourceTemplate[Identity, SecuredServerContext[Identity, P]] = + securedServerLogic[Identity, P]((vars, uri, principal, _) => logic(vars, uri, principal)) + private def definition: ResourceTemplate = ResourceTemplate(uriTemplate, name.getOrElse(uriTemplate), title, description, mimeType) end PartialResourceTemplate /** A fully-defined resource template: its metadata, a compiled URI matcher, and the logic reading matched URIs. */ -case class ServerResourceTemplate[F[_]]( +case class ServerResourceTemplate[F[_], -C <: ServerContext[F]]( definition: ResourceTemplate, matcher: UriTemplate, - read: (Map[String, String], String, 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 6dc12c1..750f720 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: ServerPrompt[F, SecuredServerContext[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 3a4208e..97d153d 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 @@ -33,13 +33,19 @@ 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.") .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,15 +53,54 @@ 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 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) + .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: ServerResource[F, SecuredServerContext[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: ServerResourceTemplate[F, SecuredServerContext[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) + .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) @@ -67,6 +112,16 @@ trait SecuredMcpServerTests[F[_]] extends AsyncFlatSpec with Matchers with Recov .addResource(greetingResource) .addTool(whoAmITool) + private def securedServerWithPromptAndResources: SecuredMcpServer[F, String, String, User] = + McpServer[F]() + .addPrompts(publicWhoAmIPrompt) + .addResources(greetingResource, publicWhoAmIResource) + .addResourceTemplates(publicWhoAmIResourceTemplate) + .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"))) @@ -85,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 => @@ -95,6 +150,13 @@ 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 { + 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}") + } + it should "serve a resource which was added after the security logic" in withSecuredServer(serverConfiguredAfterSecurity): client => client.serverCapabilities.resources shouldBe defined @@ -109,6 +171,55 @@ 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 + .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 => + 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 + .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] { Future(withSecuredServer(securedServer, token = "wrong")(_ => monad.unit(succeed))).flatten