diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 614be8f..4de0e2f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -64,6 +64,7 @@ jobs: if: github.event_name != 'pull_request' && (startsWith(github.ref, 'refs/tags/v')) secrets: inherit with: + java-version: "21" java-opts: "-Xmx4G" sttp-native: 1 @@ -76,4 +77,4 @@ jobs: if: github.event.pull_request.user.login == 'softwaremill-ci' needs: [ build, label ] uses: softwaremill/github-actions-workflows/.github/workflows/auto-merge.yml@main - secrets: inherit \ No newline at end of file + secrets: inherit diff --git a/client/src/main/scala/chimp/client/McpClient.scala b/client/src/main/scala/chimp/client/McpClient.scala index 8c33c40..24307ff 100644 --- a/client/src/main/scala/chimp/client/McpClient.scala +++ b/client/src/main/scala/chimp/client/McpClient.scala @@ -94,6 +94,24 @@ trait McpClient[F[_]]: */ def sendProgress(token: ProgressToken, progress: Double, total: Option[Double] = None, message: Option[String] = None): F[Unit] + /** Retrieves the current state of a task by its id (MCP Tasks extension, experimental). When a receiver answers `tools/call` with a + * [[chimp.protocol.CreateTaskResult]], the returned `taskId` is polled with this method until the task reaches a terminal state; the + * underlying result is then available in [[chimp.protocol.GetTaskResult.result]]. + */ + def getTask(taskId: TaskId): F[GetTaskResult] + + /** Requests cancellation of a task by its id (MCP Tasks extension, experimental). */ + def cancelTask(taskId: TaskId): F[Unit] + + /** Fulfils the input a task is waiting for while it is `InputRequired` (MCP Tasks extension, experimental). */ + def updateTask(taskId: TaskId, inputResponses: Map[String, Json]): F[Unit] + + /** Invokes a tool, declaring support for the Tasks extension (experimental). The server may answer directly + * ([[chimp.protocol.ToolCallResponse.Immediate]]) or, for a long-running call, defer with a task handle + * ([[chimp.protocol.ToolCallResponse.Deferred]]) that is then driven with [[getTask]] / [[cancelTask]] / [[updateTask]]. + */ + def callToolWithTasks(name: String, arguments: Json): F[ToolCallResponse] + /** An [[McpClient]] used over a [[chimp.client.transport.ClientBidirectionalTransport]], which additionally supports server-initiated * interactions: subscribing to resource updates, notifying the server about changes to the client's roots, and handling notifications * pushed by the server. diff --git a/client/src/main/scala/chimp/client/McpClientImpl.scala b/client/src/main/scala/chimp/client/McpClientImpl.scala index 0b96663..e2060a5 100644 --- a/client/src/main/scala/chimp/client/McpClientImpl.scala +++ b/client/src/main/scala/chimp/client/McpClientImpl.scala @@ -220,6 +220,20 @@ object McpClientImpl: val params = ProgressParams(progressToken = token, progress = progress, total = total, message = message).asJson sendNotification("notifications/progress", Some(params)) + override def getTask(taskId: TaskId): F[GetTaskResult] = + sendRequest[GetTaskResult]("tasks/get", Some(GetTaskParams(taskId).asJson)) + + override def cancelTask(taskId: TaskId): F[Unit] = + sendRequest[Json]("tasks/cancel", Some(CancelTaskParams(taskId).asJson)).map(_ => ()) + + override def updateTask(taskId: TaskId, inputResponses: Map[String, Json]): F[Unit] = + sendRequest[Json]("tasks/update", Some(UpdateTaskParams(taskId, inputResponses).asJson)).map(_ => ()) + + override def callToolWithTasks(name: String, arguments: Json): F[ToolCallResponse] = + requireServerCapability("tools/call", _.tools.isDefined): + val params = CallToolParams(name = name, arguments = arguments, _meta = Some(TasksExtension.clientCapabilityMeta)).asJson + sendRequest[ToolCallResponse]("tools/call", Some(params)) + protected def requireServerCapability[A](method: String, present: ServerCapabilities => Boolean)(action: => F[A]): F[A] = if present(serverCapabilities) then action else monad.error(McpProtocolException(s"Server did not negotiate the capability required for $method")) diff --git a/client/src/test/scala/chimp/client/TasksClientSpec.scala b/client/src/test/scala/chimp/client/TasksClientSpec.scala new file mode 100644 index 0000000..b2c71ab --- /dev/null +++ b/client/src/test/scala/chimp/client/TasksClientSpec.scala @@ -0,0 +1,86 @@ +package chimp.client + +import chimp.client.transport.ClientHttpTransport +import chimp.protocol.* +import io.circe.syntax.* +import org.scalatest.flatspec.AnyFlatSpec +import org.scalatest.matchers.should.Matchers +import sttp.client4.testing.SyncBackendStub +import sttp.client4.{GenericRequest, StringBody} +import sttp.model.StatusCode +import sttp.shared.Identity + +class TasksClientSpec extends AnyFlatSpec with Matchers: + + private val mcpUri = sttp.model.Uri.parse("http://localhost/mcp").toOption.get + private val clientInfo = Implementation(name = "chimp-test", version = "0.0.1") + + private def envelopeFor(method: String, request: GenericRequest[?, ?]): Boolean = + request.body match + case StringBody(s, _, _) => s.contains(s"\"$method\"") + case _ => false + + private val initEnvelope: String = + val initResult = InitializeResult( + protocolVersion = ProtocolVersion.Latest.name, + capabilities = ServerCapabilities(), + serverInfo = Implementation(name = "test-server", version = "1.0") + ) + (JSONRPCMessage.Response(id = RequestId(1), result = initResult.asJson): JSONRPCMessage).asJson.noSpaces + + private def client(backend: sttp.client4.testing.SyncBackendStub): McpClient[Identity] = + McpClient[Identity](ClientHttpTransport[Identity](backend, mcpUri), clientInfo, ProtocolVersion.Latest) + + it should "poll a task with tasks/get and expose the underlying result" in: + val task = GetTaskResult( + taskId = TaskId("t1"), + outcome = TaskOutcome.Completed(CallToolResult(content = List(ToolContent.Text(text = "done"))).asJson), + resultType = Some("complete") + ) + val taskEnvelope = (JSONRPCMessage.Response(id = RequestId(1), result = task.asJson): JSONRPCMessage).asJson.noSpaces + val backend = SyncBackendStub + .whenRequestMatches(envelopeFor("initialize", _)) + .thenRespondAdjust(initEnvelope) + .whenRequestMatches(envelopeFor("tasks/get", _)) + .thenRespondAdjust(taskEnvelope) + .whenAnyRequest + .thenRespondAdjust("", StatusCode.Accepted) + + val res = client(backend).getTask(TaskId("t1")) + res.status shouldBe TaskStatus.Completed + res.outcome match + case TaskOutcome.Completed(result) => + result.as[CallToolResult].toOption.map(_.content.head) shouldBe Some(ToolContent.Text("text", "done")) + case other => fail(s"expected Completed, got $other") + + it should "cancel a task with tasks/cancel" in: + val ack = TaskAck(taskId = Some(TaskId("t1")), status = Some(TaskStatus.Cancelled)) + val ackEnvelope = (JSONRPCMessage.Response(id = RequestId(1), result = ack.asJson): JSONRPCMessage).asJson.noSpaces + val backend = SyncBackendStub + .whenRequestMatches(envelopeFor("initialize", _)) + .thenRespondAdjust(initEnvelope) + .whenRequestMatches(envelopeFor("tasks/cancel", _)) + .thenRespondAdjust(ackEnvelope) + .whenAnyRequest + .thenRespondAdjust("", StatusCode.Accepted) + + noException should be thrownBy client(backend).cancelTask(TaskId("t1")) + + it should "declare task support and parse a task handle from callToolWithTasks" in: + val created = CreateTaskResult(taskId = TaskId("t9"), status = TaskStatus.Working) + val createdEnvelope = (JSONRPCMessage.Response(id = RequestId(1), result = created.asJson): JSONRPCMessage).asJson.noSpaces + val initResult = InitializeResult( + protocolVersion = ProtocolVersion.Latest.name, + capabilities = ServerCapabilities(tools = Some(ServerToolsCapability())), + serverInfo = Implementation(name = "s", version = "1") + ) + val toolsInitEnvelope = (JSONRPCMessage.Response(id = RequestId(1), result = initResult.asJson): JSONRPCMessage).asJson.noSpaces + val backend = SyncBackendStub + .whenRequestMatches(envelopeFor("initialize", _)) + .thenRespondAdjust(toolsInitEnvelope) + .whenRequestMatches(req => envelopeFor("tools/call", req) && envelopeFor(TasksExtension.Id, req)) + .thenRespondAdjust(createdEnvelope) + .whenAnyRequest + .thenRespondAdjust("", StatusCode.Accepted) + + client(backend).callToolWithTasks("slow", io.circe.Json.obj()) shouldBe ToolCallResponse.Deferred(created) diff --git a/core/src/main/scala/chimp/protocol/JsonRpc.scala b/core/src/main/scala/chimp/protocol/JsonRpc.scala index 95f04d3..aeae83e 100644 --- a/core/src/main/scala/chimp/protocol/JsonRpc.scala +++ b/core/src/main/scala/chimp/protocol/JsonRpc.scala @@ -67,3 +67,4 @@ enum JSONRPCErrorCodes(val code: Int): case InternalError extends JSONRPCErrorCodes(-32603) case InvocationError extends JSONRPCErrorCodes(-32000) case ResourceNotFound extends JSONRPCErrorCodes(-32002) + case MissingRequiredClientCapability extends JSONRPCErrorCodes(-32003) diff --git a/core/src/main/scala/chimp/protocol/Lifecycle.scala b/core/src/main/scala/chimp/protocol/Lifecycle.scala index ed859a1..a398d3c 100644 --- a/core/src/main/scala/chimp/protocol/Lifecycle.scala +++ b/core/src/main/scala/chimp/protocol/Lifecycle.scala @@ -23,7 +23,8 @@ final case class ServerCapabilities( completions: Option[Json] = None, prompts: Option[ServerPromptsCapability] = None, resources: Option[ServerResourcesCapability] = None, - tools: Option[ServerToolsCapability] = None + tools: Option[ServerToolsCapability] = None, + extensions: Option[Map[String, Json]] = None ) derives Codec final case class InitializeParams( diff --git a/core/src/main/scala/chimp/protocol/Tasks.scala b/core/src/main/scala/chimp/protocol/Tasks.scala new file mode 100644 index 0000000..928bd82 --- /dev/null +++ b/core/src/main/scala/chimp/protocol/Tasks.scala @@ -0,0 +1,187 @@ +package chimp.protocol + +import io.circe.{Codec, Decoder, Encoder, Json} +import io.circe.syntax.* + +import java.time.Instant +import scala.concurrent.duration.{DurationLong, FiniteDuration} + +// the Tasks wire format fixes durations as an integer number of milliseconds (ttlMs, pollIntervalMs); this file-private codec keeps the +// FiniteDuration domain type without leaking into the wider protocol scope +private given Codec[FiniteDuration] = + Codec.from(Decoder.decodeLong.map(_.millis), Encoder.encodeLong.contramap(_.toMillis)) + +/** Identifier of a task, generated by the receiver with enough entropy to prevent enumeration. */ +opaque type TaskId = String +object TaskId: + def apply(value: String): TaskId = value + extension (taskId: TaskId) def value: String = taskId + given Codec[TaskId] = Codec.from(Decoder.decodeString, Encoder.encodeString) + +/** The MCP Tasks extension (SEP-2663, identifier `io.modelcontextprotocol/tasks`): durable handles that let a receiver answer a request + * with a task, which the requestor then polls and later collects the result of. Experimental; the wire format follows the reference + * extension and may change. + */ +object TasksExtension: + val Id: String = "io.modelcontextprotocol/tasks" + + /** `_meta` key under which a client declares its per-request capabilities. */ + val ClientCapabilitiesMetaKey: String = "io.modelcontextprotocol/clientCapabilities" + + /** The `_meta` entry a client adds to a request to declare support for the Tasks extension, so the server may answer with a task. */ + def clientCapabilityMeta: Map[String, Json] = + Map(ClientCapabilitiesMetaKey -> Json.obj("extensions" -> Json.obj(Id -> Json.obj()))) + + /** Whether the given request `_meta` declares support for the Tasks extension. */ + def declaredIn(meta: Option[Map[String, Json]]): Boolean = + meta + .flatMap(_.get(ClientCapabilitiesMetaKey)) + .flatMap(_.hcursor.downField("extensions").downField(Id).focus) + .isDefined + +/** State of a task. Terminal states are `Completed`, `Failed` and `Cancelled`. */ +enum TaskStatus: + case Working, InputRequired, Completed, Failed, Cancelled + +object TaskStatus: + private val toWire: Map[TaskStatus, String] = Map( + Working -> "working", + InputRequired -> "input_required", + Completed -> "completed", + Failed -> "failed", + Cancelled -> "cancelled" + ) + private val fromWire: Map[String, TaskStatus] = toWire.map((k, v) => v -> k) + + def isTerminal(status: TaskStatus): Boolean = status match + case Completed | Failed | Cancelled => true + case Working | InputRequired => false + + given Encoder[TaskStatus] = Encoder.instance(status => Json.fromString(toWire(status))) + given Decoder[TaskStatus] = Decoder.decodeString.emap(s => fromWire.get(s).toRight(s"Unknown task status: $s")) + +/** Result returned when a receiver answers a request with a task instead of the request's normal result. `ttlMs` and `pollIntervalMs` carry + * their unit in the name because that is the wire field name; the values are typed as [[scala.concurrent.duration.FiniteDuration]] and + * serialized as integer milliseconds. + */ +final case class CreateTaskResult( + taskId: TaskId, + status: TaskStatus, + createdAt: Option[Instant] = None, + lastUpdatedAt: Option[Instant] = None, + ttlMs: Option[FiniteDuration] = None, + pollIntervalMs: Option[FiniteDuration] = None, + statusMessage: Option[String] = None, + resultType: String = "task", + _meta: Option[Map[String, Json]] = None +) derives Codec + +/** The response to a `tools/call` made with task support declared: the receiver either answers immediately with the tool's + * [[CallToolResult]], or defers by returning a [[CreateTaskResult]] task handle to poll. + */ +enum ToolCallResponse: + case Immediate(result: CallToolResult) + case Deferred(task: CreateTaskResult) + +object ToolCallResponse: + given Decoder[ToolCallResponse] = Decoder.instance: c => + c.get[Option[String]]("resultType") + .flatMap: + case Some("task") => c.as[CreateTaskResult].map(ToolCallResponse.Deferred(_)) + case _ => c.as[CallToolResult].map(ToolCallResponse.Immediate(_)) + +final case class GetTaskParams(taskId: TaskId, _meta: Option[Map[String, Json]] = None) derives Codec +final case class GetTaskRequest(method: String = "tasks/get", params: GetTaskParams) derives Codec + +/** The status of a task together with the data specific to that status. Modelling it as a sealed type means `result`, `error` and + * `inputRequests` can only appear with the status they belong to. `result` is raw JSON, since its shape depends on the request the task + * stands for. + */ +enum TaskOutcome: + case Working + case InputRequired(inputRequests: Map[String, Json]) + case Completed(result: Json) + case Failed(error: JSONRPCErrorObject) + case Cancelled + + def status: TaskStatus = this match + case TaskOutcome.Working => TaskStatus.Working + case TaskOutcome.InputRequired(_) => TaskStatus.InputRequired + case TaskOutcome.Completed(_) => TaskStatus.Completed + case TaskOutcome.Failed(_) => TaskStatus.Failed + case TaskOutcome.Cancelled => TaskStatus.Cancelled + +/** Detailed task state returned by `tasks/get`. The [[TaskOutcome]] carries the `status` and its status-specific data. On the wire the + * outcome is flattened: `status` plus, where applicable, `result` / `error` / `inputRequests`. + */ +final case class GetTaskResult( + taskId: TaskId, + outcome: TaskOutcome, + createdAt: Option[Instant] = None, + lastUpdatedAt: Option[Instant] = None, + ttlMs: Option[FiniteDuration] = None, + pollIntervalMs: Option[FiniteDuration] = None, + statusMessage: Option[String] = None, + resultType: Option[String] = None, + _meta: Option[Map[String, Json]] = None +): + def status: TaskStatus = outcome.status + +object GetTaskResult: + given Encoder[GetTaskResult] = Encoder.instance: task => + val base = Json.obj( + "taskId" -> task.taskId.asJson, + "status" -> task.status.asJson, + "createdAt" -> task.createdAt.asJson, + "lastUpdatedAt" -> task.lastUpdatedAt.asJson, + "ttlMs" -> task.ttlMs.asJson, + "pollIntervalMs" -> task.pollIntervalMs.asJson, + "statusMessage" -> task.statusMessage.asJson, + "resultType" -> task.resultType.asJson, + "_meta" -> task._meta.asJson + ) + val payload = task.outcome match + case TaskOutcome.Completed(result) => Json.obj("result" -> result) + case TaskOutcome.Failed(error) => Json.obj("error" -> error.asJson) + case TaskOutcome.InputRequired(inputRequests) => Json.obj("inputRequests" -> inputRequests.asJson) + case TaskOutcome.Working | TaskOutcome.Cancelled => Json.obj() + base.deepMerge(payload) + + given Decoder[GetTaskResult] = Decoder.instance: c => + for + taskId <- c.get[TaskId]("taskId") + status <- c.get[TaskStatus]("status") + createdAt <- c.get[Option[Instant]]("createdAt") + lastUpdatedAt <- c.get[Option[Instant]]("lastUpdatedAt") + ttlMs <- c.get[Option[FiniteDuration]]("ttlMs") + pollIntervalMs <- c.get[Option[FiniteDuration]]("pollIntervalMs") + statusMessage <- c.get[Option[String]]("statusMessage") + resultType <- c.get[Option[String]]("resultType") + meta <- c.get[Option[Map[String, Json]]]("_meta") + outcome <- status match + case TaskStatus.Working => Right(TaskOutcome.Working) + case TaskStatus.Cancelled => Right(TaskOutcome.Cancelled) + case TaskStatus.Completed => c.get[Json]("result").map(result => TaskOutcome.Completed(result)) + case TaskStatus.Failed => c.get[JSONRPCErrorObject]("error").map(error => TaskOutcome.Failed(error)) + case TaskStatus.InputRequired => c.get[Map[String, Json]]("inputRequests").map(reqs => TaskOutcome.InputRequired(reqs)) + yield GetTaskResult(taskId, outcome, createdAt, lastUpdatedAt, ttlMs, pollIntervalMs, statusMessage, resultType, meta) + +final case class CancelTaskParams(taskId: TaskId, _meta: Option[Map[String, Json]] = None) derives Codec +final case class CancelTaskRequest(method: String = "tasks/cancel", params: CancelTaskParams) derives Codec + +final case class UpdateTaskParams(taskId: TaskId, inputResponses: Map[String, Json], _meta: Option[Map[String, Json]] = None) derives Codec +final case class UpdateTaskRequest(method: String = "tasks/update", params: UpdateTaskParams) derives Codec + +/** Acknowledgement returned by `tasks/cancel` and `tasks/update`. */ +final case class TaskAck( + taskId: Option[TaskId] = None, + status: Option[TaskStatus] = None, + resultType: String = "complete", + _meta: Option[Map[String, Json]] = None +) derives Codec + +/** Notification pushed by a receiver that supports task subscriptions; carries the same fields as a `tasks/get` result. */ +final case class TaskStatusNotification( + method: String = "notifications/tasks", + params: GetTaskResult +) derives Codec diff --git a/core/src/test/scala/chimp/protocol/TasksSpec.scala b/core/src/test/scala/chimp/protocol/TasksSpec.scala new file mode 100644 index 0000000..ca9bb56 --- /dev/null +++ b/core/src/test/scala/chimp/protocol/TasksSpec.scala @@ -0,0 +1,71 @@ +package chimp.protocol + +import io.circe.Json +import io.circe.parser.decode +import io.circe.syntax.* +import org.scalatest.flatspec.AnyFlatSpec +import org.scalatest.matchers.should.Matchers + +import java.time.Instant +import scala.concurrent.duration.* + +class TasksSpec extends AnyFlatSpec with Matchers: + + it should "decode a CreateTaskResult from the spec example" in: + val json = + """ + { + "resultType": "task", + "taskId": "786512e2-9e0d-44bd-8f29-789f320fe840", + "status": "working", + "createdAt": "2025-11-25T10:30:00Z", + "lastUpdatedAt": "2025-11-25T10:50:00Z", + "ttlMs": 3600000, + "pollIntervalMs": 5000 + } + """ + val res = decode[CreateTaskResult](json) + res.map(_.taskId) shouldBe Right(TaskId("786512e2-9e0d-44bd-8f29-789f320fe840")) + res.map(_.status) shouldBe Right(TaskStatus.Working) + res.map(_.createdAt) shouldBe Right(Some(Instant.parse("2025-11-25T10:30:00Z"))) + res.map(_.ttlMs) shouldBe Right(Some(1.hour)) + res.map(_.pollIntervalMs) shouldBe Right(Some(5.seconds)) + + it should "encode durations as integer milliseconds on the wire" in: + val created = + CreateTaskResult( + taskId = TaskId("t"), + status = TaskStatus.Working, + ttlMs = Some(1.hour), + pollIntervalMs = Some(5.seconds) + ) + val json = created.asJson + json.hcursor.downField("ttlMs").as[Long] shouldBe Right(3600000L) + json.hcursor.downField("pollIntervalMs").as[Long] shouldBe Right(5000L) + + it should "encode and decode task status with the spec wire strings" in: + (TaskStatus.InputRequired: TaskStatus).asJson shouldBe Json.fromString("input_required") + (TaskStatus.Cancelled: TaskStatus).asJson shouldBe Json.fromString("cancelled") + decode[TaskStatus](""" "completed" """.trim) shouldBe Right(TaskStatus.Completed) + + it should "reject an unknown task status" in: + decode[TaskStatus](""" "bogus" """.trim).isLeft shouldBe true + + it should "round-trip a completed GetTaskResult carrying the tool result" in: + val toolResult = CallToolResult(content = List(ToolContent.Text(text = "Hello, Luca!"))).asJson + val task = GetTaskResult( + taskId = TaskId("t1"), + outcome = TaskOutcome.Completed(toolResult), + resultType = Some("complete") + ) + decode[GetTaskResult](task.asJson.noSpaces) shouldBe Right(task) + + it should "flatten the outcome onto the wire and reject a completed task without a result" in: + val completed = GetTaskResult(taskId = TaskId("t1"), outcome = TaskOutcome.Completed(Json.obj("k" -> Json.fromInt(1)))) + completed.asJson.hcursor.downField("status").as[String] shouldBe Right("completed") + completed.asJson.hcursor.downField("result").as[Json] shouldBe Right(Json.obj("k" -> Json.fromInt(1))) + decode[GetTaskResult]("""{ "taskId": "t1", "status": "completed" }""").isLeft shouldBe true + + it should "mark only completed, failed and cancelled as terminal" in: + TaskStatus.values.filter(TaskStatus.isTerminal).toSet shouldBe + Set(TaskStatus.Completed, TaskStatus.Failed, TaskStatus.Cancelled) diff --git a/docs/client/capabilities.md b/docs/client/capabilities.md index bede01f..f4e725f 100644 --- a/docs/client/capabilities.md +++ b/docs/client/capabilities.md @@ -43,3 +43,19 @@ def listen(client: BidirectionalMcpClient[Task]): Task[Unit] = case _ => ZIO.unit } ``` + +## Tasks (experimental) + +The [Tasks extension](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/main/seps/2663-tasks-extension.md) (`io.modelcontextprotocol/tasks`) lets a server answer a long-running request with a durable task handle that the client polls and later collects the result of. The client supports the requestor side with `getTask`, `cancelTask` and `updateTask`. This part of the protocol is experimental and its wire format may change. + +```scala mdoc:compile-only +import chimp.client.* +import chimp.protocol.* +import zio.{Task, ZIO} + +def awaitResult(client: McpClient[Task], taskId: TaskId): Task[GetTaskResult] = + client.getTask(taskId).flatMap { task => + if TaskStatus.isTerminal(task.status) then ZIO.succeed(task) + else awaitResult(client, taskId) + } +``` diff --git a/docs/server/capabilities.md b/docs/server/capabilities.md index f4f34cf..d3df9f4 100644 --- a/docs/server/capabilities.md +++ b/docs/server/capabilities.md @@ -29,3 +29,56 @@ val server = StreamingMcpServer[Identity]().addStreamingTool(work) ``` Server-wide capabilities are enabled by registering a handler — only what you wire up is advertised: `.withCompletion`, `.withLoggingLevel`, `.withSubscriptions`. + +## Tasks (experimental) + +With the [Tasks extension](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/main/seps/2663-tasks-extension.md) (`io.modelcontextprotocol/tasks`) the server answers a long-running `tools/call` with a durable task handle instead of blocking. The client (which must declare the extension in its per-request `_meta`) then polls with `tasks/get` and collects the result once the task is `completed`. + +Enable it with `.withTasks`, providing a `TaskStore` and a `TaskExecutor` for the effect type. The synchronous server uses a thread-pool executor: + +```scala mdoc:compile-only +import chimp.server.* +import io.circe.Codec +import sttp.monad.{IdentityMonad, MonadError} +import sttp.shared.Identity +import sttp.tapir.Schema + +given MonadError[Identity] = IdentityMonad + +case class ReportInput(days: Int) derives Codec, Schema + +val report = tool("report").input[ReportInput].handle(in => ToolResult.text(s"report for ${in.days} days")) + +val server = McpServer(tools = List(report)) + .withTasks(TaskSupport(TaskStore.inMemory[Identity], TaskExecutor.threadPool())) +``` + +The server runs the tool in the background, transitions the task to `completed` (with the tool's `CallToolResult`) or `failed`, and answers `tasks/cancel` by interrupting the worker. `useTask` and `requireTask` on `TaskSupport` control, per tool, whether a task is offered or required; a required task with no client support fails with `-32003`. + +### Requesting input while running + +A tool registered with `addTaskTool` (defined with `.taskLogic`) receives a `TaskContext` and can ask the client for input mid-run. `requestInput` moves the task to `input_required`, surfacing the request under a key via `tasks/get`; the tool parks until the client answers with `tasks/update`, then resumes: + +```scala mdoc:compile-only +import chimp.server.* +import io.circe.{Codec, Json} +import sttp.monad.{IdentityMonad, MonadError} +import sttp.shared.Identity +import sttp.tapir.Schema + +given MonadError[Identity] = IdentityMonad + +case class ReviewInput(text: String) derives Codec, Schema + +val review = tool("review") + .input[ReviewInput] + .taskLogic[Identity]: (in, ctx, _) => + val decision = ctx.requestInput("approve", Json.fromString(s"Approve: ${in.text}?")) + ToolResult.text(s"decision: ${decision.noSpaces}") + +val server = McpServer() + .withTasks(TaskSupport(TaskStore.inMemory[Identity], TaskExecutor.threadPool())) + .addTaskTool(review) +``` + +A task tool always runs as a task and requires the client to declare task support, so calling it without that support fails with `-32003`. diff --git a/server/src/main/scala/chimp/server/McpHandler.scala b/server/src/main/scala/chimp/server/McpHandler.scala index 585dcee..6f59274 100644 --- a/server/src/main/scala/chimp/server/McpHandler.scala +++ b/server/src/main/scala/chimp/server/McpHandler.scala @@ -29,10 +29,12 @@ 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 taskToolsByName = server.taskTools.map(tool => tool.name -> tool).toMap + private val inputCoordinator = TaskInputCoordinator() 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 toolDefinitions = server.tools.map(toolToDefinition) + private val toolDefinitions = server.tools.map(toolToDefinition) ++ server.taskTools.map(toolToDefinition) private def toJsonSchema(toolSchema: ToolSchema): Json = toolSchema match case ToolSchema.Derived(schema) => @@ -40,7 +42,7 @@ private[server] class McpHandler[F[_], C <: ServerContext[F]](server: McpServerD (if server.showJsonSchemaMetadata then base else base.copy($schema = None)).asJson case ToolSchema.Raw(json) => json - private def toolToDefinition(tool: ServerTool[?, ?, F, C]): ToolDefinition = + private def toolToDefinition(tool: ServerTool[?, ?, F, ?]): ToolDefinition = ToolDefinition( name = tool.name, description = tool.description, @@ -82,6 +84,12 @@ private[server] class McpHandler[F[_], C <: ServerContext[F]](server: McpServerD jsonResponse(JSONRPCMessage.Response(id = id, result = ListToolsResponse(toolDefinitions).asJson)).unit case "tools/call" => handleToolsCall(params, id, headers, makeContext).map(jsonResponse) + case "tasks/get" if server.tasks.isDefined => + handleTasksGet(params, id).map(jsonResponse) + case "tasks/cancel" if server.tasks.isDefined => + handleTasksCancel(params, id).map(jsonResponse) + case "tasks/update" if server.tasks.isDefined => + handleTasksUpdate(params, id).map(jsonResponse) case "resources/list" if hasResources => jsonResponse(JSONRPCMessage.Response(id = id, result = ListResourcesResult(server.resources.map(_.definition)).asJson)).unit case "resources/templates/list" if hasResources => @@ -126,7 +134,8 @@ private[server] class McpHandler[F[_], C <: ServerContext[F]](server: McpServerD prompts = Option.when(server.prompts.nonEmpty)(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))) + tools = Option.when(server.tools.nonEmpty)(ServerToolsCapability(listChanged = Some(false))), + extensions = server.tasks.map(_ => Map(TasksExtension.Id -> Json.obj())) ) val result = InitializeResult( protocolVersion = negotiated.name, @@ -137,46 +146,189 @@ private[server] class McpHandler[F[_], C <: ServerContext[F]](server: McpServerD JSONRPCMessage.Response(id = id, result = result.asJson) private def handleToolsCall(params: Option[Json], id: RequestId, headers: Seq[Header], makeContext: Option[ProgressToken] => C)(using - MonadError[F] + m: MonadError[F] ): F[JSONRPCMessage] = val name = params.flatMap(_.hcursor.downField("name").as[String].toOption) val arguments = params.flatMap(_.hcursor.downField("arguments").focus).getOrElse(Json.obj()) val progressToken = params.flatMap(_.hcursor.downField("_meta").downField("progressToken").as[ProgressToken].toOption) + val requestMeta = params.flatMap(_.hcursor.downField("_meta").as[Map[String, Json]].toOption) + val clientSupportsTasks = TasksExtension.declaredIn(requestMeta) + + def invalidArguments(error: DecodingFailure): F[JSONRPCMessage] = + protocolError( + id, + JSONRPCErrorCodes.InvalidParams.code, + s"Invalid arguments: ${error.getMessage}. Input: ${arguments.noSpaces.take(200)}" + ).unit + + def missingTaskCapability(toolName: String): F[JSONRPCMessage] = + protocolError( + id, + JSONRPCErrorCodes.MissingRequiredClientCapability.code, + s"Tool '$toolName' requires the ${TasksExtension.Id} client capability" + ).unit + name match case Some(name) => - toolsByName.get(name) match - case Some(tool) => - tool.inputDecoder.decodeJson(arguments) match - case Right(input) => - val context = makeContext(progressToken) - tool - .logic(input, context, headers) - .map: result => - toolCallResponse(id, result) - case Left(decodingError) => - val snippet = arguments.noSpaces.take(200) + taskToolsByName.get(name) match + // a task tool always runs as a task and needs both server task support and the client capability + case Some(taskTool) => + server.tasks match + case None => protocolError( id, - JSONRPCErrorCodes.InvalidParams.code, - s"Invalid arguments: ${decodingError.getMessage}. Input: $snippet" + JSONRPCErrorCodes.InternalError.code, + s"Tool '$name' runs as a task, but task support is not configured" ).unit - case None => protocolError(id, JSONRPCErrorCodes.MethodNotFound.code, s"Unknown tool: $name").unit + case Some(_) if !clientSupportsTasks => missingTaskCapability(name) + case Some(support) => + taskTool.inputDecoder.decodeJson(arguments) match + case Right(input) => runTask(support, id)(taskId => taskTool.logic(input, makeTaskContext(support, taskId), headers)) + case Left(error) => invalidArguments(error) + case None => + toolsByName.get(name) match + case Some(tool) => + tool.inputDecoder.decodeJson(arguments) match + case Right(input) => + val context = makeContext(progressToken) + server.tasks match + case Some(support) if support.requireTask(name) && !clientSupportsTasks => missingTaskCapability(name) + case Some(support) if clientSupportsTasks && support.useTask(name) => + runTask(support, id)(_ => tool.logic(input, context, headers)) + case _ => + tool.logic(input, context, headers).map(result => toolCallResponse(id, result)) + case Left(error) => invalidArguments(error) + case None => protocolError(id, JSONRPCErrorCodes.MethodNotFound.code, s"Unknown tool: $name").unit case None => protocolError(id, JSONRPCErrorCodes.InvalidParams.code, "Missing tool name").unit - private def toolCallResponse(id: RequestId, result: ToolResult[?]): JSONRPCMessage = + private def toCallToolResult(result: ToolResult[?]): CallToolResult = // for backwards compatibility, structured output is serialized into a text block, unless the tool returned content of its own val content = result.structuredContent match case Some(json) if result.content.isEmpty => List(ToolContent.Text(text = json.noSpaces)) case _ => result.content - JSONRPCMessage.Response( - id = id, - result = CallToolResult( - content = content, - structuredContent = result.structuredContent, - isError = result.isError - ).asJson + CallToolResult(content = content, structuredContent = result.structuredContent, isError = result.isError) + + private def toolCallResponse(id: RequestId, result: ToolResult[?]): JSONRPCMessage = + JSONRPCMessage.Response(id = id, result = toCallToolResult(result).asJson) + + private def runTask[O](support: TaskSupport[F], id: RequestId)(compute: TaskId => F[ToolResult[O]])(using + m: MonadError[F] + ): F[JSONRPCMessage] = + val taskId = TaskId(java.util.UUID.randomUUID().toString) + val now = java.time.Instant.now() + val initial = GetTaskResult( + taskId = taskId, + outcome = TaskOutcome.Working, + createdAt = Some(now), + lastUpdatedAt = Some(now), + ttlMs = support.ttl, + pollIntervalMs = support.pollInterval, + resultType = Some("complete") ) + // start and handleError both take their body by-name, so a synchronous (Identity) tool that throws is caught here too + support.store + .create(initial) + .flatMap { _ => + support.executor.start(taskId)( + m.handleError( + m.flatMap(compute(taskId))(result => finishTask(support, taskId, TaskOutcome.Completed(toCallToolResult(result).asJson))) + ) { case t => + finishTask( + support, + taskId, + TaskOutcome.Failed(JSONRPCErrorObject(JSONRPCErrorCodes.InternalError.code, Option(t.getMessage).getOrElse("Task failed"))) + ) + } + ) + } + .map { _ => + JSONRPCMessage.Response( + id = id, + result = CreateTaskResult( + taskId = taskId, + status = TaskStatus.Working, + createdAt = Some(now), + lastUpdatedAt = Some(now), + ttlMs = support.ttl, + pollIntervalMs = support.pollInterval + ).asJson + ) + } + + private def makeTaskContext(support: TaskSupport[F], taskId: TaskId)(using m: MonadError[F]): TaskContext[F] = + new TaskContext[F]: + def requestInput(key: String, request: Json): F[Json] = + // register the waiter before advertising input_required, so a fast tasks/update is never lost + m.flatMap(m.eval(inputCoordinator.register(taskId, key))) { waiter => + m.flatMap(setInputRequired(support, taskId, key, request)) { _ => + // waiter.get blocks the worker until tasks/update delivers the answer; fine on the virtual-thread executor + m.flatMap(m.eval(waiter.get()))(response => m.map(resolveInput(support, taskId, key))(_ => response)) + } + } + + private def setInputRequired(support: TaskSupport[F], taskId: TaskId, key: String, request: Json)(using MonadError[F]): F[Unit] = + support.store + .update(taskId): current => + val outstanding = current.outcome match + case TaskOutcome.InputRequired(requests) => requests + case _ => Map.empty[String, Json] + current.copy(outcome = TaskOutcome.InputRequired(outstanding + (key -> request)), lastUpdatedAt = Some(java.time.Instant.now())) + .map(_ => ()) + + private def resolveInput(support: TaskSupport[F], taskId: TaskId, key: String)(using MonadError[F]): F[Unit] = + support.store + .update(taskId): current => + current.outcome match + case TaskOutcome.InputRequired(requests) => + val remaining = requests - key + val outcome = if remaining.isEmpty then TaskOutcome.Working else TaskOutcome.InputRequired(remaining) + current.copy(outcome = outcome, lastUpdatedAt = Some(java.time.Instant.now())) + case _ => current + .map(_ => ()) + + // only transition a task that is still working, so a cancellation is not overwritten by a late completion + private def finishTask(support: TaskSupport[F], taskId: TaskId, outcome: TaskOutcome)(using MonadError[F]): F[Unit] = + support.store + .update(taskId): current => + if current.status == TaskStatus.Working then current.copy(outcome = outcome, lastUpdatedAt = Some(java.time.Instant.now())) + else current + .map(_ => ()) + + private def handleTasksGet(params: Option[Json], id: RequestId)(using MonadError[F]): F[JSONRPCMessage] = + decodeParams[GetTaskParams](params, id): p => + server.tasks.get.store + .get(p.taskId) + .map: + case Some(task) => JSONRPCMessage.Response(id = id, result = task.asJson) + case None => protocolError(id, JSONRPCErrorCodes.InvalidParams.code, s"Unknown task: ${p.taskId}") + + private def handleTasksCancel(params: Option[Json], id: RequestId)(using MonadError[F]): F[JSONRPCMessage] = + decodeParams[CancelTaskParams](params, id): p => + val support = server.tasks.get + support.store + .update(p.taskId): current => + if TaskStatus.isTerminal(current.status) then current + else current.copy(outcome = TaskOutcome.Cancelled, lastUpdatedAt = Some(java.time.Instant.now())) + .flatMap: + case Some(_) => + inputCoordinator.cancel(p.taskId) + support.executor.cancel(p.taskId).map(_ => taskAck(id, p.taskId, TaskStatus.Cancelled)) + case None => protocolError(id, JSONRPCErrorCodes.InvalidParams.code, s"Unknown task: ${p.taskId}").unit + + private def handleTasksUpdate(params: Option[Json], id: RequestId)(using MonadError[F]): F[JSONRPCMessage] = + decodeParams[UpdateTaskParams](params, id): p => + server.tasks.get.store + .get(p.taskId) + .map: + case Some(task) => + // hand each response to the waiting task tool; it transitions the task back to working itself + p.inputResponses.foreach((key, response) => inputCoordinator.deliverInput(p.taskId, key, response)) + taskAck(id, task.taskId, task.status) + case None => protocolError(id, JSONRPCErrorCodes.InvalidParams.code, s"Unknown task: ${p.taskId}") + + private def taskAck(id: RequestId, taskId: TaskId, status: TaskStatus): JSONRPCMessage = + JSONRPCMessage.Response(id = id, result = TaskAck(taskId = Some(taskId), status = Some(status)).asJson) private def handleResourcesRead(params: Option[Json], id: RequestId, headers: Seq[Header])(using MonadError[F]): F[JSONRPCMessage] = decodeParams[ReadResourceParams](params, id): params => diff --git a/server/src/main/scala/chimp/server/McpServer.scala b/server/src/main/scala/chimp/server/McpServer.scala index 75fca77..71232e4 100644 --- a/server/src/main/scala/chimp/server/McpServer.scala +++ b/server/src/main/scala/chimp/server/McpServer.scala @@ -26,6 +26,8 @@ sealed trait McpServerDef[F[_], C <: ServerContext[F]]: def completion: Option[CompletionHandler[F]] def loggingLevel: Option[SetLoggingLevelHandler[F]] def subscriptions: Option[ResourceSubscriptions[F]] + def tasks: Option[TaskSupport[F]] + def taskTools: List[ServerTool[?, ?, F, TaskContext[F]]] case class McpServer[F[_]]( name: String = "Chimp MCP server", @@ -39,7 +41,9 @@ case class McpServer[F[_]]( resourceTemplates: List[ServerResourceTemplate[F]] = Nil, completion: Option[CompletionHandler[F]] = None, loggingLevel: Option[SetLoggingLevelHandler[F]] = None, - subscriptions: Option[ResourceSubscriptions[F]] = None + subscriptions: Option[ResourceSubscriptions[F]] = None, + tasks: Option[TaskSupport[F]] = None, + taskTools: List[ServerTool[?, ?, F, TaskContext[F]]] = Nil ) extends McpServerDef[F, ServerContext[F]]: def name(value: String): McpServer[F] = copy(name = value) @@ -89,6 +93,13 @@ case class McpServer[F[_]]( def withSubscriptions(handler: ResourceSubscriptions[F]): McpServer[F] = copy(subscriptions = Some(handler)) + def withTasks(support: TaskSupport[F]): McpServer[F] = + copy(tasks = Some(support)) + + /** Registers a tool that runs as a task and can request input from the client while running (Tasks extension). */ + def addTaskTool(tool: ServerTool[?, ?, F, TaskContext[F]]): McpServer[F] = + copy(taskTools = taskTools :+ tool) + def endpoint(path: List[String]): ServerEndpoint[Any, F] = ServerHttpTransport(path).serve(this) def streaming: StreamingMcpServer[F] = @@ -104,7 +115,9 @@ case class McpServer[F[_]]( resourceTemplates, completion, loggingLevel, - subscriptions + subscriptions, + tasks, + taskTools ) case class StreamingMcpServer[F[_]]( @@ -119,7 +132,9 @@ case class StreamingMcpServer[F[_]]( resourceTemplates: List[ServerResourceTemplate[F]] = Nil, completion: Option[CompletionHandler[F]] = None, loggingLevel: Option[SetLoggingLevelHandler[F]] = None, - subscriptions: Option[ResourceSubscriptions[F]] = None + subscriptions: Option[ResourceSubscriptions[F]] = None, + tasks: Option[TaskSupport[F]] = None, + taskTools: List[ServerTool[?, ?, F, TaskContext[F]]] = Nil ) extends McpServerDef[F, StreamingServerContext[F]]: def name(value: String): StreamingMcpServer[F] = copy(name = value) @@ -174,3 +189,10 @@ case class StreamingMcpServer[F[_]]( def withSubscriptions(handler: ResourceSubscriptions[F]): StreamingMcpServer[F] = copy(subscriptions = Some(handler)) + + def withTasks(support: TaskSupport[F]): StreamingMcpServer[F] = + copy(tasks = Some(support)) + + /** Registers a tool that runs as a task and can request input from the client while running (Tasks extension). */ + def addTaskTool(tool: ServerTool[?, ?, F, TaskContext[F]]): StreamingMcpServer[F] = + copy(taskTools = taskTools :+ tool) diff --git a/server/src/main/scala/chimp/server/ServerContext.scala b/server/src/main/scala/chimp/server/ServerContext.scala index c764818..e8cc027 100644 --- a/server/src/main/scala/chimp/server/ServerContext.scala +++ b/server/src/main/scala/chimp/server/ServerContext.scala @@ -10,6 +10,13 @@ trait ServerContext[F[_]] object ServerContext: def noop[F[_]]: ServerContext[F] = new ServerContext[F] {} +/** A [[ServerContext]] for a tool running as a task, letting it request input from the client mid-execution (Tasks extension, + * experimental). Requesting input moves the task to `input_required`, surfacing `key -> request` via `tasks/get`; the tool resumes with + * the response once the client answers with `tasks/update`. + */ +trait TaskContext[F[_]] extends ServerContext[F]: + def requestInput(key: String, request: Json): F[Json] + trait StreamingServerContext[F[_]] extends ServerContext[F]: def reportProgress(progress: Double, total: Option[Double] = None, message: Option[String] = None): F[Unit] def log(level: LoggingLevel, data: Json, logger: Option[String] = None): F[Unit] diff --git a/server/src/main/scala/chimp/server/TaskSupport.scala b/server/src/main/scala/chimp/server/TaskSupport.scala new file mode 100644 index 0000000..d6219db --- /dev/null +++ b/server/src/main/scala/chimp/server/TaskSupport.scala @@ -0,0 +1,106 @@ +package chimp.server + +import chimp.protocol.{GetTaskResult, TaskId} +import io.circe.Json +import sttp.monad.MonadError +import sttp.shared.Identity + +import java.util.concurrent.{CompletableFuture, ConcurrentHashMap, ExecutorService, Executors, Future as JavaFuture} +import scala.concurrent.duration.{DurationInt, FiniteDuration} + +/** Durable-ish store of task state for the Tasks extension, addressable by task id. The default in-memory implementation keeps tasks for + * the lifetime of the process. + */ +trait TaskStore[F[_]]: + def create(task: GetTaskResult): F[Unit] + def get(taskId: TaskId): F[Option[GetTaskResult]] + + /** Applies `f` to the stored task if present, atomically, and returns the updated task. */ + def update(taskId: TaskId)(f: GetTaskResult => GetTaskResult): F[Option[GetTaskResult]] + +object TaskStore: + def inMemory[F[_]](using m: MonadError[F]): TaskStore[F] = new TaskStore[F]: + private val tasks = ConcurrentHashMap[TaskId, GetTaskResult]() + + def create(task: GetTaskResult): F[Unit] = m.eval: + val _ = tasks.put(task.taskId, task) + () + + def get(taskId: TaskId): F[Option[GetTaskResult]] = m.eval(Option(tasks.get(taskId))) + + def update(taskId: TaskId)(f: GetTaskResult => GetTaskResult): F[Option[GetTaskResult]] = m.eval: + Option(tasks.computeIfPresent(taskId, (_, current) => f(current))) + +/** Runs task bodies in the background and supports best-effort cancellation. The body is passed as a thunk so that, on eager effect types + * such as `Identity`, it is only run on the background worker rather than at the call site. + */ +trait TaskExecutor[F[_]]: + def start(taskId: TaskId)(body: => F[Unit]): F[Unit] + def cancel(taskId: TaskId): F[Unit] + +object TaskExecutor: + + /** A [[TaskExecutor]] for synchronous (`Identity`) servers, such as the Netty sync server, backed by an `ExecutorService`. Defaults to a + * virtual-thread-per-task executor (JDK 21+), which suits the blocking tool logic that tasks run and scales to many concurrent tasks. + * Cancellation interrupts the worker thread. + */ + def threadPool(pool: ExecutorService = Executors.newVirtualThreadPerTaskExecutor()): TaskExecutor[Identity] = new TaskExecutor[Identity]: + private val running = ConcurrentHashMap[TaskId, JavaFuture[?]]() + + def start(taskId: TaskId)(body: => Identity[Unit]): Identity[Unit] = + val future = pool.submit(new Runnable: + def run(): Unit = + try body + finally + val _ = running.remove(taskId)) + val _ = running.put(taskId, future) + () + + def cancel(taskId: TaskId): Identity[Unit] = + val _ = Option(running.remove(taskId)).foreach(_.cancel(true)) + () + +/** Bundles everything a server needs to answer requests with tasks (Tasks extension, experimental). + * + * @param useTask + * Given a tool name, whether to answer its `tools/call` with a task when the client declares task support. Defaults to always. + * @param requireTask + * Given a tool name, whether a task is required; if the client does not declare task support, the call fails with `-32003`. Defaults to + * never. + */ +final case class TaskSupport[F[_]]( + store: TaskStore[F], + executor: TaskExecutor[F], + ttl: Option[FiniteDuration] = Some(1.hour), + pollInterval: Option[FiniteDuration] = Some(1.second), + useTask: String => Boolean = (_: String) => true, + requireTask: String => Boolean = (_: String) => false +) + +/** Coordinates the `input_required` flow within a process: a running task tool parks on the future from [[register]] for a `(taskId, key)` + * while `tasks/update` hands the answer over with [[deliverInput]]. Cancelling a task unblocks any of its waiters. + */ +private[server] final class TaskInputCoordinator: + private val pending = ConcurrentHashMap[(TaskId, String), CompletableFuture[Json]]() + + /** Registers interest in an answer for `(taskId, key)` and returns the future to block on. Call this before advertising the request + * (moving the task to `input_required`), so a fast `tasks/update` is never delivered before there is a waiter. + */ + def register(taskId: TaskId, key: String): CompletableFuture[Json] = + pending.computeIfAbsent((taskId, key), _ => CompletableFuture[Json]()) + + /** Hands `response` to a waiter, if any. Returns whether a waiter was present. */ + def deliverInput(taskId: TaskId, key: String, response: Json): Boolean = + Option(pending.remove((taskId, key))).exists { future => + future.complete(response) + true + } + + /** Unblocks all waiters for a cancelled task by failing their futures. */ + def cancel(taskId: TaskId): Unit = + val iterator = pending.entrySet().iterator() + while iterator.hasNext do + val entry = iterator.next() + if entry.getKey._1 == taskId then + entry.getValue.completeExceptionally(InterruptedException("task cancelled")) + iterator.remove() diff --git a/server/src/main/scala/chimp/server/Tool.scala b/server/src/main/scala/chimp/server/Tool.scala index 1e4f489..f09c28d 100644 --- a/server/src/main/scala/chimp/server/Tool.scala +++ b/server/src/main/scala/chimp/server/Tool.scala @@ -109,6 +109,12 @@ case class Tool[I, O]( ): ServerTool[I, O, F, StreamingServerContext[F]] = ServerTool(name, description, inputSchema, inputDecoder, outputSchema, annotations, logic) + /** Attaches effectful logic with access to a [[TaskContext]], so the tool can request input from the client while running as a task + * (Tasks extension). Register it with `addTaskTool`; such a tool is always answered with a task. + */ + def taskLogic[F[_]](logic: (I, TaskContext[F], Seq[Header]) => F[ToolResult[O]]): ServerTool[I, O, F, TaskContext[F]] = + ServerTool(name, description, inputSchema, inputDecoder, outputSchema, annotations, logic) + /** Attaches synchronous logic that also receives the request headers. */ def handleWithHeaders(logic: (I, Seq[Header]) => ToolResult[O]): ServerTool[I, O, Identity, ServerContext[Identity]] = ServerTool(name, description, inputSchema, inputDecoder, outputSchema, annotations, (i, _, headers) => logic(i, headers)) diff --git a/server/src/test/scala/chimp/server/TaskServerSpec.scala b/server/src/test/scala/chimp/server/TaskServerSpec.scala new file mode 100644 index 0000000..7a34762 --- /dev/null +++ b/server/src/test/scala/chimp/server/TaskServerSpec.scala @@ -0,0 +1,172 @@ +package chimp.server + +import chimp.protocol.* +import chimp.protocol.JSONRPCMessage.given +import io.circe.Json +import io.circe.syntax.* +import org.scalatest.flatspec.AnyFlatSpec +import org.scalatest.matchers.should.Matchers +import sttp.monad.{IdentityMonad, MonadError} +import sttp.shared.Identity +import sttp.tapir.Schema + +class TaskServerSpec extends AnyFlatSpec with Matchers: + import JSONRPCMessage.* + + private given MonadError[Identity] = IdentityMonad + + case class TIn(message: String) derives Schema, io.circe.Codec + + private val instant = tool("instant").input[TIn].handle(in => ToolResult.text(s"echo:${in.message}")) + private val slow = tool("slow") + .input[TIn] + .handle: in => + Thread.sleep(150) + ToolResult.text(s"slow:${in.message}") + private val boom = tool("boom").input[TIn].handle(_ => throw new RuntimeException("boom")) + private val forever = tool("forever") + .input[TIn] + .handle: _ => + Thread.sleep(10000) + ToolResult.text("done") + + private val askThenEcho = tool("askThenEcho") + .input[TIn] + .taskLogic[Identity]: (in, ctx, _) => + val answer = ctx.requestInput("q1", Json.fromString(s"need input for ${in.message}")) + ToolResult.text(s"got:${answer.noSpaces}") + + private def handlerWith(requireTask: String => Boolean = _ => false): McpHandler[Identity, ServerContext[Identity]] = + val support = TaskSupport[Identity]( + store = TaskStore.inMemory[Identity], + executor = TaskExecutor.threadPool(), + requireTask = requireTask + ) + McpHandler(McpServer(tools = List(instant, slow, boom, forever)).withTasks(support).addTaskTool(askThenEcho)) + + private def resultJson(response: McpResponse): Json = + val json = response match + case McpResponse.JsonResponse(j) => j + case McpResponse.EmptyAcceptResponse => fail("expected JsonResponse") + json.as[JSONRPCMessage].getOrElse(fail("decode JSONRPCMessage")) match + case Response(_, _, result) => result + case other => fail(s"expected Response, got $other") + + private def errorObj(response: McpResponse): JSONRPCErrorObject = + val json = response match + case McpResponse.JsonResponse(j) => j + case McpResponse.EmptyAcceptResponse => fail("expected JsonResponse") + json.as[JSONRPCMessage].getOrElse(fail("decode JSONRPCMessage")) match + case Error(_, _, err) => err + case other => fail(s"expected Error, got $other") + + private def callToolReq(name: String, withTasks: Boolean): Json = + val meta = Option.when(withTasks)(TasksExtension.clientCapabilityMeta) + val params = CallToolParams(name = name, arguments = TIn("hi").asJson, _meta = meta).asJson + (Request(method = "tools/call", params = Some(params), id = RequestId("call")): JSONRPCMessage).asJson + + private def pollUntil(handler: McpHandler[Identity, ServerContext[Identity]], taskId: TaskId)( + predicate: GetTaskResult => Boolean + ): GetTaskResult = + var last = GetTaskResult(taskId = taskId, outcome = TaskOutcome.Working) + var done = false + var i = 0 + while !done && i < 200 do + val req = (Request(method = "tasks/get", params = Some(GetTaskParams(taskId).asJson), id = RequestId("get")): JSONRPCMessage).asJson + last = resultJson(handler.handleJsonRpc(req, Seq.empty)).as[GetTaskResult].getOrElse(fail("decode GetTaskResult")) + if predicate(last) then done = true + else + Thread.sleep(20) + i += 1 + last + + private def pollTask(handler: McpHandler[Identity, ServerContext[Identity]], taskId: TaskId): GetTaskResult = + pollUntil(handler, taskId)(task => TaskStatus.isTerminal(task.status)) + + "a task-enabled server" should "run tools/call synchronously when the client does not declare task support" in: + val handler = handlerWith() + val result = resultJson(handler.handleJsonRpc(callToolReq("instant", withTasks = false), Seq.empty)) + val call = result.as[CallToolResult].getOrElse(fail("decode CallToolResult")) + call.content.head shouldBe ToolContent.Text("text", "echo:hi") + + it should "answer with a task and complete it when the client declares support" in: + val handler = handlerWith() + val created = resultJson(handler.handleJsonRpc(callToolReq("slow", withTasks = true), Seq.empty)) + .as[CreateTaskResult] + .getOrElse(fail("decode CreateTaskResult")) + created.status shouldBe TaskStatus.Working + created.taskId.value should not be empty + + val finished = pollTask(handler, created.taskId) + finished.outcome match + case TaskOutcome.Completed(result) => + result.as[CallToolResult].toOption.map(_.content.head) shouldBe Some(ToolContent.Text("text", "slow:hi")) + case other => fail(s"expected Completed, got $other") + + it should "report a failed task when the tool throws" in: + val handler = handlerWith() + val created = resultJson(handler.handleJsonRpc(callToolReq("boom", withTasks = true), Seq.empty)) + .as[CreateTaskResult] + .getOrElse(fail("decode CreateTaskResult")) + + val finished = pollTask(handler, created.taskId) + finished.outcome match + case TaskOutcome.Failed(_) => succeed + case other => fail(s"expected Failed, got $other") + + it should "cancel a running task" in: + val handler = handlerWith() + val created = resultJson(handler.handleJsonRpc(callToolReq("forever", withTasks = true), Seq.empty)) + .as[CreateTaskResult] + .getOrElse(fail("decode CreateTaskResult")) + + val cancelReq = + (Request(method = "tasks/cancel", params = Some(CancelTaskParams(created.taskId).asJson), id = RequestId("c")): JSONRPCMessage).asJson + val ack = resultJson(handler.handleJsonRpc(cancelReq, Seq.empty)).as[TaskAck].getOrElse(fail("decode TaskAck")) + ack.status shouldBe Some(TaskStatus.Cancelled) + + pollTask(handler, created.taskId).status shouldBe TaskStatus.Cancelled + + it should "reject a required-task call that lacks the client capability with -32003" in: + val handler = handlerWith(requireTask = _ == "slow") + val err = errorObj(handler.handleJsonRpc(callToolReq("slow", withTasks = false), Seq.empty)) + err.code shouldBe JSONRPCErrorCodes.MissingRequiredClientCapability.code + + it should "advertise the tasks extension capability on initialize" in: + val handler = handlerWith() + val req = (Request(method = "initialize", id = RequestId("i")): JSONRPCMessage).asJson + val result = resultJson(handler.handleJsonRpc(req, Seq.empty)).as[InitializeResult].getOrElse(fail("decode InitializeResult")) + result.capabilities.extensions.map(_.keySet) shouldBe Some(Set(TasksExtension.Id)) + + it should "run an input_required task tool: request input, answer via tasks/update, then complete" in: + val handler = handlerWith() + val created = resultJson(handler.handleJsonRpc(callToolReq("askThenEcho", withTasks = true), Seq.empty)) + .as[CreateTaskResult] + .getOrElse(fail("decode CreateTaskResult")) + created.status shouldBe TaskStatus.Working + + // the tool has requested input and parked, surfacing the request via tasks/get + val waiting = pollUntil(handler, created.taskId)(_.status == TaskStatus.InputRequired) + waiting.outcome match + case TaskOutcome.InputRequired(requests) => requests.keySet shouldBe Set("q1") + case other => fail(s"expected InputRequired, got $other") + + // answer it + val updateReq = (Request( + method = "tasks/update", + params = Some(UpdateTaskParams(created.taskId, Map("q1" -> Json.fromString("42"))).asJson), + id = RequestId("u") + ): JSONRPCMessage).asJson + val ack = resultJson(handler.handleJsonRpc(updateReq, Seq.empty)).as[TaskAck].getOrElse(fail("decode TaskAck")) + ack.taskId shouldBe Some(created.taskId) + + // the tool resumes and completes with the answer folded in + pollTask(handler, created.taskId).outcome match + case TaskOutcome.Completed(result) => + result.as[CallToolResult].toOption.map(_.content.head) shouldBe Some(ToolContent.Text("text", "got:\"42\"")) + case other => fail(s"expected Completed, got $other") + + it should "reject the input_required tool when the client does not declare task support" in: + val handler = handlerWith() + val err = errorObj(handler.handleJsonRpc(callToolReq("askThenEcho", withTasks = false), Seq.empty)) + err.code shouldBe JSONRPCErrorCodes.MissingRequiredClientCapability.code