WIP: feat: MCP Tasks extension (full: protocol, server execution, client)

.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
+    secrets: inherit

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.

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"))

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)

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)

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(

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