diff --git a/contributing/samples/mcp/mcp_tasks_agent/__init__.py b/contributing/samples/mcp/mcp_tasks_agent/__init__.py new file mode 100644 index 00000000000..4015e47d6e4 --- /dev/null +++ b/contributing/samples/mcp/mcp_tasks_agent/__init__.py @@ -0,0 +1,15 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from . import agent diff --git a/contributing/samples/mcp/mcp_tasks_agent/agent.py b/contributing/samples/mcp/mcp_tasks_agent/agent.py new file mode 100644 index 00000000000..c19e18513b0 --- /dev/null +++ b/contributing/samples/mcp/mcp_tasks_agent/agent.py @@ -0,0 +1,70 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Sample agent for the MCP Tasks extension. + +An operation that takes minutes cannot be held open across an HTTP write +timeout, and the usual workaround -- splitting one tool into start, status and +result, and letting the model run the polling loop -- costs a model turn per +poll and only happens if the prompt says so. + +With `enable_tasks=True`, the server decides per call that the work is +long-running and answers with a task handle; the toolset polls it to +completion and returns the result. The agent below sees one tool and one +result, exactly as if the call had answered inline. + +Needs MCP SDK 2.x installed (`mcp>=2,<3`); the extension seam this rides on +does not exist in 1.x. + +Start the server first, in another terminal: + + python contributing/samples/mcp/mcp_tasks_agent/task_server.py + +Then: + + adk run contributing/samples/mcp/mcp_tasks_agent + +Ask it to run the slow operation. The call takes about twenty seconds and +comes back as an ordinary tool result. +""" + +from __future__ import annotations + +from google.adk.agents.llm_agent import LlmAgent +from google.adk.tools.mcp_tool.mcp_session_manager import StreamableHTTPConnectionParams +from google.adk.tools.mcp_tool.mcp_toolset import McpToolset + +root_agent = LlmAgent( + model="gemini-2.0-flash", + name="tasks_agent", + instruction="""\ +You run long operations for the user with the `slow_operation` tool. + +Call it once and report what it returns. It takes a while; that is expected, +and you do not need to poll or call it again.""", + tools=[ + McpToolset( + connection_params=StreamableHTTPConnectionParams( + url="http://localhost:3000/mcp", + # The connection budget, not a cap on the operation: a task is + # polled over as many requests as it takes. + timeout=10.0, + ), + # Accept a task handle instead of a blocking call, when the server + # offers one. Off by default; servers without the extension are + # unaffected. + enable_tasks=True, + ) + ], +) diff --git a/contributing/samples/mcp/mcp_tasks_agent/task_server.py b/contributing/samples/mcp/mcp_tasks_agent/task_server.py new file mode 100644 index 00000000000..31ffe0839f3 --- /dev/null +++ b/contributing/samples/mcp/mcp_tasks_agent/task_server.py @@ -0,0 +1,210 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""An MCP server that answers a slow tool call with a task. + +Servers implementing the Tasks extension are still scarce, so this one exists +to have something to point `enable_tasks=True` at. It keeps a slow tool that +takes longer than any sensible HTTP write timeout, hands back a task handle +instead of blocking, and serves `tasks/get` and `tasks/cancel` against an +in-memory store. + +Run it with: + + python contributing/samples/mcp/mcp_tasks_agent/task_server.py +""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass +from dataclasses import field +import datetime +from typing import Any +from typing import Literal +import uuid + +from mcp.server.extension import Extension +from mcp.server.extension import MethodBinding +from mcp.server.mcpserver import MCPServer +import mcp_types + +TASKS_EXTENSION_ID = "io.modelcontextprotocol/tasks" + +# Long enough that holding the request open would be the wrong answer, short +# enough to watch. +_WORK_DURATION_SECONDS = 20.0 +_POLL_INTERVAL_MS = 2000 +_TTL_MS = 300000 + + +def _now() -> str: + return ( + datetime.datetime.now(datetime.timezone.utc) + .isoformat(timespec="seconds") + .replace("+00:00", "Z") + ) + + +@dataclass +class _Task: + """One tracked operation, and whatever it has produced so far.""" + + task_id: str + status: str = "working" + status_message: str | None = None + created_at: str = field(default_factory=_now) + last_updated_at: str = field(default_factory=_now) + result: dict[str, Any] | None = None + error: dict[str, Any] | None = None + worker: asyncio.Task[None] | None = None + + def to_wire(self) -> dict[str, Any]: + wire: dict[str, Any] = { + "taskId": self.task_id, + "status": self.status, + "createdAt": self.created_at, + "lastUpdatedAt": self.last_updated_at, + "ttlMs": _TTL_MS, + "pollIntervalMs": _POLL_INTERVAL_MS, + } + if self.status_message is not None: + wire["statusMessage"] = self.status_message + if self.result is not None: + wire["result"] = self.result + if self.error is not None: + wire["error"] = self.error + return wire + + +class _TaskParams(mcp_types.RequestParams): + task_id: str + + +class TasksExtension(Extension): + """Serves the Tasks extension over an in-memory task store.""" + + identifier = TASKS_EXTENSION_ID + + def __init__(self) -> None: + self._tasks: dict[str, _Task] = {} + + def methods(self): + return ( + MethodBinding( + method="tasks/get", + params_type=_TaskParams, + handler=self._handle_get, + ), + MethodBinding( + method="tasks/cancel", + params_type=_TaskParams, + handler=self._handle_cancel, + ), + ) + + async def _handle_get(self, ctx: Any, params: _TaskParams) -> dict[str, Any]: + del ctx + task = self._tasks.get(params.task_id) + if task is None: + raise mcp_types.MCPError( + code=mcp_types.jsonrpc.INVALID_PARAMS, + message=f"unknown taskId {params.task_id!r}", + ) + return task.to_wire() + + async def _handle_cancel( + self, ctx: Any, params: _TaskParams + ) -> dict[str, Any]: + del ctx + task = self._tasks.get(params.task_id) + if task is not None and task.worker is not None: + task.worker.cancel() + return {} + + async def intercept_tool_call(self, params, ctx, call_next): + """Turns a call to the slow tool into a task, when the client can take one. + + A server must not hand a task to a client that did not advertise the + extension on this very request, so anything else falls through to the + ordinary handler. + """ + if params.name != "slow_operation" or not _client_supports_tasks(ctx): + return await call_next(ctx) + + task = _Task(task_id=str(uuid.uuid4())) + task.status_message = "The operation is now in progress." + self._tasks[task.task_id] = task + task.worker = asyncio.create_task(self._run(task, params)) + + return {"resultType": "task", **task.to_wire()} + + async def _run(self, task: _Task, params: Any) -> None: + """Does the slow work, then leaves the outcome on the task.""" + label = (params.arguments or {}).get("label", "the job") + try: + await asyncio.sleep(_WORK_DURATION_SECONDS) + except asyncio.CancelledError: + task.status = "cancelled" + task.last_updated_at = _now() + raise + except Exception as e: # pylint: disable=broad-except + task.status = "failed" + task.error = {"code": -32000, "message": str(e)} + task.last_updated_at = _now() + return + task.status = "completed" + task.status_message = "Done." + message = f"Finished {label} after {int(_WORK_DURATION_SECONDS)} seconds." + # A task's result is an ordinary tool result and is held to the tool's + # output schema, so it needs the structured half too -- the client + # validates it exactly as it would an inline answer. + task.result = { + "content": [{"type": "text", "text": message}], + "structuredContent": {"result": message}, + } + task.last_updated_at = _now() + + +def _client_supports_tasks(ctx: Any) -> bool: + """Whether this request advertised the tasks extension.""" + meta = getattr(ctx, "meta", None) or {} + capabilities = meta.get("io.modelcontextprotocol/clientCapabilities") or {} + extensions = capabilities.get("extensions") or {} + return TASKS_EXTENSION_ID in extensions + + +server = MCPServer( + name="Task Server", + instructions="Runs an operation that takes longer than one request.", + extensions=[TasksExtension()], +) + + +@server.tool() +async def slow_operation(label: str = "the job") -> str: + """Runs a long operation. Served as a task when the client supports it. + + Args: + label: What to call the job in the final message. + """ + # Reached only by a client without the tasks extension, which has no choice + # but to hold the request open. + await asyncio.sleep(_WORK_DURATION_SECONDS) + return f"Finished {label} after {int(_WORK_DURATION_SECONDS)} seconds." + + +if __name__ == "__main__": + print("Task server listening on http://localhost:3000/mcp") + server.run(transport="streamable-http", host="localhost", port=3000) diff --git a/docs/guides/README.md b/docs/guides/README.md index c283efb1408..eb5bda335a6 100644 --- a/docs/guides/README.md +++ b/docs/guides/README.md @@ -63,6 +63,7 @@ This directory contains specific developer guides for the ADK Python implementat * [State](sessions/state/index.md) - Session state and the app:, user:, and temp: prefixes that decide what is shared and what is stored. ### Tools +* [MCP Tasks](tools/mcp_tool/tasks/index.md) - Accept task-augmented MCP tool calls, so an operation can outlive the request that started it. * [Node as tool](tools/node_tool/index.md) - Exposing workflows and deterministic nodes as agent tools with isolated runtime branching and resume support. * [to_mcp_server](tools/mcp_tool/agent_to_mcp/index.md) - Expose an ADK agent as an MCP server so any MCP host can drive it as a single tool (the MCP counterpart of to_a2a). diff --git a/docs/guides/tools/mcp_tool/tasks/index.md b/docs/guides/tools/mcp_tool/tasks/index.md new file mode 100644 index 00000000000..47a2fe94b0f --- /dev/null +++ b/docs/guides/tools/mcp_tool/tasks/index.md @@ -0,0 +1,134 @@ +# MCP Tasks + +Accepts task-augmented tool calls from an MCP server, so an operation can +outlive the request that started it. Enabled with `enable_tasks=True` on +`McpToolset`. + +## Introduction + +A tool call normally answers on the connection that made it. That breaks down +once the work takes minutes: an intermediary caps how long a request may stay +open, and a dropped connection loses the operation with no way to pick it back +up. The usual workaround is to split one operation into `start_x`, `x_status` +and `x_result` tools and let the model run the polling loop, which costs a +model turn per poll and only happens if the prompt says so. + +The MCP Tasks extension, `io.modelcontextprotocol/tasks`, is the protocol's +answer. The server decides per request that the work is long-running and +replies to `tools/call` with a durable task handle rather than a result. The +client polls `tasks/get` until the task reaches a terminal state and reads the +result from there. + +`McpToolset` does that polling for you. The agent declares one tool, calls it +once, and receives an ordinary tool result; only the wire changes. + +## Requirements + +MCP SDK 2.x. ADK's pin admits both majors, and the client extension seam the +Tasks path rides on exists only in 2.x -- 1.x has no such parameters on +`ClientSession` at all. Passing `enable_tasks=True` on a 1.x install raises at +toolset construction rather than going quiet, since an opt-in that silently +did nothing here would read as a tool call that never returns. Install +`mcp>=2,<3` to use it. + +## Get started + +Point a toolset at a server that supports the extension and opt in: + +```python +toolset = McpToolset( + connection_params=StreamableHTTPConnectionParams( + url="http://localhost:3000/mcp", + # The connection budget, not a cap on the operation. + timeout=10.0, + ), + enable_tasks=True, +) + +agent = LlmAgent( + name="tasks_agent", + instruction="Run long operations with the tools you have.", + tools=[toolset], +) +``` + +Nothing else changes. A tool the server chooses to serve as a task returns the +same result shape as one it answers inline, and a server that does not support +the extension behaves exactly as it does today. + +## How it works + +Enabling tasks registers two things on the client session together: the +extension identifier in the capabilities the client advertises, and a *result +claim* that teaches `tools/call` parsing to accept a `resultType: "task"` +response. They are registered as a pair because a claim whose extension is not +advertised is rejected when the session is built. + +Because the capability is advertised per request, the server can decide call +by call. When it answers with a task handle, `McpTool` resolves that handle +before returning: it polls `tasks/get` at the interval the server states in +`pollIntervalMs`, following the latest value on each poll, until the task +reports `completed`, `failed` or `cancelled`. A completed task carries the +tool result, which is validated against the tool's output schema exactly as an +inline result would be. + +Two details follow from the operation outliving the request: + +- The pooled session is held out of the idle sweep for the whole resolution, + not just the initial call, so a long poll cannot have its transport closed + underneath it. +- Cancelling the invocation sends `tasks/cancel` on a best-effort basis before + the cancellation propagates, so the server can stop work the caller no + longer wants. + +Extensions only bind on a `2026-07-28` connection. `initialize()` always +performs the older handshake, so enabling tasks also switches session bring-up +to `server/discover`, falling back to `initialize()` when the server does not +support it. That fallback logs a warning and leaves tasks inactive rather than +failing the session. + +## Configuration options + +| Option | Type | Default | Description | +| :--- | :--- | :--- | :--- | +| `enable_tasks` | `bool` | `False` | Accept task-augmented tool calls. | + +`enable_tasks` is off by default because it changes how sessions negotiate. +Turning it on costs nothing against a server without the extension — the +capability is advertised and ignored — but the `server/discover` probe is new +wire traffic, so it is not paid for by callers who did not ask for it. + +The option is also available in YAML tool configuration as `enable_tasks: +true`. + +For an extension you carry yourself, the underlying seam is exposed directly +as `extensions`, `result_claims` and `notification_bindings`, which are passed +to the MCP `ClientSession` unchanged. Combining `enable_tasks=True` with your +own `io.modelcontextprotocol/tasks` entry raises `ValueError`: two claims on +one result type is refused when the session is built, and failing at +construction says so while the cause is still in view. + +## Limitations + +- **A task that asks for input is not answered.** A task can reach + `input_required` to request data mid-flight. Answering it means bridging + `tasks/update` to an interaction channel, which is not implemented; the task + is cancelled and the call returns an error result. +- **The call still blocks.** The operation survives a dropped connection at + the protocol level, but `run_async` does not return until the task finishes. + There is no way to hand the agent a task handle and collect the result on a + later turn. +- **Notifications are not used.** `notifications/tasks` would remove the + polling, but it requires a subscription the toolset does not open, so + progress is discovered by polling only. +- **A failed task is reported, not raised.** `failed` and server-side + `cancelled` come back as a tool result with `isError` set, matching what a + tool that fails inline produces. +- **Servers are scarce.** The extension is recent, so few servers implement + it. The sample below includes one to test against. + +## Related samples + +- [MCP Tasks agent](../../../../../contributing/samples/mcp/mcp_tasks_agent/agent.py) - + an agent with `enable_tasks=True`, and a server that answers a slow tool + call with a task. diff --git a/src/google/adk/dependencies/_mcp.py b/src/google/adk/dependencies/_mcp.py index 6e49138016f..65ad6e61bf5 100644 --- a/src/google/adk/dependencies/_mcp.py +++ b/src/google/adk/dependencies/_mcp.py @@ -19,8 +19,8 @@ vendored there, which carries a different distribution name. Every MCP import in ADK goes through this module so that difference lives in one place. -The pin admits both 1.x and 2.x, so this module also spans them. Fourteen of -the seventeen names below sit at the same path in both. The three that moved +The pin admits both 1.x and 2.x, so this module also spans them. Fifteen of +the eighteen names below sit at the same path in both. The three that moved are resolved here, once, rather than at seventeen call sites: * `McpError` was renamed `MCPError` @@ -31,6 +31,11 @@ Try 2.x first. A stale 1.x path that still happens to exist in 2.x would otherwise win and bind the wrong object. +Three further names have no 1.x path at all rather than a moved one: the +client extension seam `ClientSession` grew in 2.x. They are bound to `Any` +on 1.x so that the modules annotating them stay importable there, and the +opt-in that would reach them is refused before it can (`McpToolset`). + This flavor must never name the internal copy. That name is already taken on PyPI by an unrelated project, so a released wheel importing it would bind to someone else's package on any machine that happened to have it installed. @@ -38,6 +43,8 @@ from __future__ import annotations +from typing import Any + from mcp import ClientSession as ClientSession from mcp import SamplingCapability as SamplingCapability from mcp import StdioServerParameters as StdioServerParameters @@ -49,6 +56,7 @@ from mcp.client.streamable_http import create_mcp_http_client as create_mcp_http_client from mcp.client.streamable_http import streamable_http_client as streamable_http_client from mcp.server.session import ServerSession as ServerSession +from mcp.types import CallToolResult as CallToolResult from mcp.types import ListResourcesResult as ListResourcesResult from mcp.types import ListToolsResult as ListToolsResult from mcp.types import Tool as Tool @@ -71,8 +79,24 @@ IS_MCP_SDK_V2 = False +if IS_MCP_SDK_V2: + from mcp.client.extension import ClaimContext as ClaimContext + from mcp.client.extension import NotificationBinding as NotificationBinding + from mcp.client.extension import ResultClaim as ResultClaim +else: + # `Any`, not a stand-in class: a stand-in would satisfy an `isinstance` or a + # constructor call that has no business succeeding here. `Any` is also why + # the annotations naming these stay unsubscripted downstream -- `Any[Any]` + # is a `TypeError` to anything that resolves an annotation, and the point of + # these three is to keep 1.x importable. + ClaimContext = Any + NotificationBinding = Any + ResultClaim = Any + __all__ = [ "IS_MCP_SDK_V2", + "CallToolResult", + "ClaimContext", "ClientSession", "Context", "ElicitationFnT", @@ -80,6 +104,8 @@ "ListResourcesResult", "ListToolsResult", "McpError", + "NotificationBinding", + "ResultClaim", "SamplingCapability", "SamplingFnT", "ServerSession", diff --git a/src/google/adk/tools/mcp_tool/_tasks.py b/src/google/adk/tools/mcp_tool/_tasks.py new file mode 100644 index 00000000000..d401b0c09e0 --- /dev/null +++ b/src/google/adk/tools/mcp_tool/_tasks.py @@ -0,0 +1,279 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Client support for the MCP Tasks extension. + +A server that expects an operation to be slow may answer `tools/call` with a +task handle instead of a result. The work then outlives the request: the +client polls `tasks/get` until the task reaches a terminal state and reads +the result from there. That is what makes an operation survive a dropped +connection, and what removes the need to split one tool into a start/status/ +result triplet driven by the model. + +The wire models live here because no Python implementation of the extension +exists yet -- neither the SDK nor a published package -- and the task types +the SDK does carry are the older, wire-incompatible design that predates it. +Everything in this module is private so it can be replaced by an upstream +implementation without changing ADK's own surface. + +MCP SDK 2.x only, and unimportable on 1.x: `ClaimContext` and `ResultClaim` +have no 1.x counterpart. `McpToolset` refuses `enable_tasks=True` there, and +imports this module only after that check. +""" + +from __future__ import annotations + +import asyncio +import logging +from typing import Any +from typing import Literal + +from ...dependencies._mcp import CallToolResult +from ...dependencies._mcp import ClaimContext +from ...dependencies._mcp import McpError +from ...dependencies._mcp import ResultClaim +from ...dependencies._mcp import types as mcp_types + +logger = logging.getLogger('google_adk.' + __name__) + +TASKS_EXTENSION_ID = 'io.modelcontextprotocol/tasks' + +# A task in one of these states will not change again, so polling stops. +_TERMINAL_STATUSES = frozenset({'completed', 'failed', 'cancelled'}) + +# What to wait between polls when the server states no preference. Tasks exist +# for minute-scale work, so a second costs nothing and keeps idle load low. +_DEFAULT_POLL_INTERVAL_SECONDS = 1.0 + +# The floor stops a server that reports 0 from turning the loop into a spin. +# The ceiling bounds how long a finished task can go unnoticed; it overrides +# the server's preference, but only in the direction of noticing sooner. +_MIN_POLL_INTERVAL_SECONDS = 0.1 +_MAX_POLL_INTERVAL_SECONDS = 30.0 + +# Cancellation is best-effort by specification, and it runs while the caller +# is already being torn down, so it gets one short round trip and no more. +_CANCEL_TIMEOUT_SECONDS = 2.0 + +# How often a still-running task says so, counted in polls. Without it a task +# that never finishes is indistinguishable from a client that stopped asking. +_PROGRESS_LOG_EVERY_N_POLLS = 30 + + +class _TaskResult(mcp_types.Result): + """The `resultType: "task"` answer to `tools/call`. + + Flat by specification: the task's fields sit alongside `resultType` rather + than under a nested object. Field names are snake_case with the wire + spelling supplied by the base model's alias generator. + """ + + result_type: Literal['task'] = 'task' + task_id: str + status: mcp_types.TaskStatus + status_message: str | None = None + created_at: str + last_updated_at: str + ttl_ms: int | None + poll_interval_ms: int | None = None + + +class _TaskParams(mcp_types.RequestParams): + """Parameters shared by every `tasks/*` request.""" + + task_id: str + + +class _GetTaskRequest(mcp_types.Request[_TaskParams, Literal['tasks/get']]): + """Reads the current state of a task. + + `name_param` is what puts the task id in the `Mcp-Name` header, which the + specification requires on `tasks/*` over streamable HTTP so an intermediary + can route the request to whoever holds the task. + """ + + method: Literal['tasks/get'] = 'tasks/get' + params: _TaskParams + name_param = 'taskId' + + +class _CancelTaskRequest( + mcp_types.Request[_TaskParams, Literal['tasks/cancel']] +): + """Asks the server to stop working on a task.""" + + method: Literal['tasks/cancel'] = 'tasks/cancel' + params: _TaskParams + name_param = 'taskId' + + +class _GetTaskResult(mcp_types.Result): + """The state of a task, plus whatever its current state carries. + + Only one of `result`, `error` and `input_requests` is ever populated, and + which one follows from `status`. + """ + + task_id: str + status: mcp_types.TaskStatus + status_message: str | None = None + created_at: str + last_updated_at: str + ttl_ms: int | None = None + poll_interval_ms: int | None = None + result: dict[str, Any] | None = None + error: dict[str, Any] | None = None + input_requests: dict[str, Any] | None = None + + +class _EmptyResult(mcp_types.Result): + """An acknowledgement with no payload, as `tasks/cancel` returns.""" + + +def _error_result(message: str) -> CallToolResult: + """Builds the failed tool result the agent will see.""" + return CallToolResult( + content=[mcp_types.TextContent(type='text', text=message)], is_error=True + ) + + +def _poll_interval_seconds(poll_interval_ms: int | None) -> float: + """Turns the server's stated interval into one worth sleeping for.""" + if poll_interval_ms is None: + return _DEFAULT_POLL_INTERVAL_SECONDS + return min( + max(poll_interval_ms / 1000.0, _MIN_POLL_INTERVAL_SECONDS), + _MAX_POLL_INTERVAL_SECONDS, + ) + + +def _describe_error(error: dict[str, Any] | None) -> str: + """Renders a JSON-RPC error object for a human reading a tool result.""" + if not error: + return 'no error detail was provided' + code, message = error.get('code'), error.get('message') + if code is None and message is None: + return str(error) + return f'{message or "unknown error"} (code {code})' + + +async def _cancel_task(ctx: ClaimContext, task_id: str) -> None: + """Tells the server to drop a task, without ever raising. + + Called from a cancellation path, so it must neither fail nor hang: whatever + goes wrong here, the caller's own cancellation is what matters. + """ + try: + await asyncio.wait_for( + asyncio.shield( + ctx.session.send_request( + _CancelTaskRequest(params=_TaskParams(task_id=task_id)), + _EmptyResult, + ) + ), + timeout=_CANCEL_TIMEOUT_SECONDS, + ) + except (Exception, asyncio.CancelledError) as e: # pylint: disable=broad-except + logger.debug('Best-effort tasks/cancel for %s did not land: %s', task_id, e) + + +def _completed_to_result(task: _GetTaskResult) -> CallToolResult: + """Reads the tool's result out of a completed task.""" + if task.result is None: + return _error_result(f'MCP task {task.task_id} completed without a result.') + try: + return CallToolResult.model_validate(task.result) + except Exception as e: # pylint: disable=broad-except + # The task did finish; it is the payload that is unusable. Report that as + # a failed tool call rather than letting a validation error escape a + # resolver the caller has no way to guard. + return _error_result( + f'MCP task {task.task_id} returned a result that is not a valid' + f' tool result: {e}' + ) + + +async def _resolve_task( + claimed: _TaskResult, ctx: ClaimContext +) -> CallToolResult: + """Polls a task to completion and returns what the tool call produced. + + Blocks until the task reaches a terminal state, so the agent sees the same + result it would have seen from a tool that answered inline. + + A task that fails, is cancelled by the server, or asks for input it cannot + be given comes back as a failed tool result rather than an exception: a + tool that fails reports `isError`, and going through a task should not + change that. + """ + task_id = claimed.task_id + status: str = claimed.status + interval = _poll_interval_seconds(claimed.poll_interval_ms) + polls = 0 + + try: + while True: + # A task that is already terminal still has to be read once: the result + # only ever travels on tasks/get. + if status not in _TERMINAL_STATUSES: + await asyncio.sleep(interval) + + polls += 1 + task = await ctx.session.send_request( + _GetTaskRequest(params=_TaskParams(task_id=task_id)), + _GetTaskResult, + request_read_timeout_seconds=ctx.read_timeout_seconds, + ) + status = task.status + interval = _poll_interval_seconds(task.poll_interval_ms) + + if status == 'completed': + return _completed_to_result(task) + if status == 'failed': + return _error_result( + f'MCP task {task_id} failed: {_describe_error(task.error)}' + ) + if status == 'cancelled': + return _error_result(f'MCP task {task_id} was cancelled by the server.') + if status == 'input_required': + # Answering would mean bridging tasks/update to an interaction + # channel, which this does not do yet. Release the task rather than + # leaving the server holding it until its ttl runs out. + await _cancel_task(ctx, task_id) + return _error_result( + f'MCP task {task_id} requires additional input, which is not' + ' supported yet.' + ) + + if polls % _PROGRESS_LOG_EVERY_N_POLLS == 0: + logger.debug( + 'MCP task %s still %s after %d polls', task_id, status, polls + ) + except asyncio.CancelledError: + await _cancel_task(ctx, task_id) + raise + except McpError as e: + # An expired or forgotten task answers with an error. Retrying would not + # help: the handle is what is gone. + return _error_result( + f'MCP task {task_id} could not be read (last known status' + f' {status!r}): {e}' + ) + + +def make_tasks_claim() -> ResultClaim[_TaskResult]: + """Builds the claim that turns a task handle into a tool result.""" + return ResultClaim( + result_type='task', model=_TaskResult, resolve=_resolve_task + ) diff --git a/src/google/adk/tools/mcp_tool/mcp_session_manager.py b/src/google/adk/tools/mcp_tool/mcp_session_manager.py index 1583308d6da..a18e91b04be 100644 --- a/src/google/adk/tools/mcp_tool/mcp_session_manager.py +++ b/src/google/adk/tools/mcp_tool/mcp_session_manager.py @@ -31,9 +31,11 @@ from typing import AsyncIterator from typing import Callable from typing import Dict +from typing import Mapping from typing import Optional from typing import Protocol from typing import runtime_checkable +from typing import Sequence from typing import TextIO import urllib.parse @@ -64,6 +66,8 @@ class AsyncAuthorizedSession: # pylint: disable=g-bad-classes from ...dependencies._mcp import create_mcp_http_client as _create_mcp_http_client from ...dependencies._mcp import ElicitationFnT from ...dependencies._mcp import IS_MCP_SDK_V2 +from ...dependencies._mcp import NotificationBinding +from ...dependencies._mcp import ResultClaim from ...dependencies._mcp import SamplingCapability from ...dependencies._mcp import SamplingFnT from ...dependencies._mcp import sse_client @@ -134,6 +138,50 @@ def create_mcp_http_client( ) +def _require_extension_support(what: str) -> None: + """Refuses an extension opt-in the installed MCP SDK cannot honor. + + The pin admits both majors and the extension seam exists only in 2.x, so a + caller can ask for something that has nowhere to go. Raising is a choice + over no-oping: what extensions are for is a tool call that runs for + minutes, and an opt-in that silently does nothing there is indistinguishable + from a hang -- the failure would surface as a call that never returns, on + the install least equipped to explain why. It fires at construction, before + a session exists, so the message arrives while the cause is still in view. + + Args: + what: The opt-in being refused, named as the caller wrote it. + + Raises: + ValueError: If the installed MCP SDK is 1.x. + """ + if IS_MCP_SDK_V2: + return + raise ValueError( + f'{what} requires MCP SDK 2.x; the installed SDK is 1.x, whose' + ' ClientSession has no extension seam to carry it. Install' + " `mcp>=2,<3`, or drop the option: a server's extensions stay unused" + ' either way, and every other MCP feature works on both majors.' + ) + + +def _index_claims_by_model( + result_claims: Mapping[str, Sequence[ResultClaim]] | None, +) -> dict[type[Any], ResultClaim]: + """Indexes result claims by the model each one parses into. + + A claimed `tools/call` response arrives already parsed into the claim's + model, and its type is the only thing tying it back to the claim that has + to resolve it. `ClientSession` validates and parses but never resolves, so + the lookup has to live on this side. + """ + by_model: dict[type[Any], ResultClaim] = {} + for claims in (result_claims or {}).values(): + for claim in claims: + by_model[claim.model] = claim + return by_model + + def _redact_headers(headers: dict[str, str]) -> dict[str, str]: sensitive_keys = { 'api-key', @@ -686,6 +734,9 @@ def __init__( sampling_callback: SamplingFnT | None = None, sampling_capabilities: SamplingCapability | None = None, elicitation_callback: ElicitationFnT | None = None, + extensions: dict[str, dict[str, Any]] | None = None, + result_claims: Mapping[str, Sequence[ResultClaim]] | None = None, + notification_bindings: Sequence[NotificationBinding] | None = None, ): """Initializes the MCP session manager. @@ -701,10 +752,23 @@ def __init__( elicitation_callback: Optional callback to handle elicitation requests from the MCP server (``elicitation/create``), including URL-mode elicitations used for out-of-band flows such as auth challenges. + extensions: MCP extensions to advertise, keyed by extension + identifier. + result_claims: Non-core ``tools/call`` result shapes to accept, keyed + by the identifier of the extension that defines them. + notification_bindings: Handlers for extension notifications. """ self._sampling_callback = sampling_callback self._sampling_capabilities = sampling_capabilities self._elicitation_callback = elicitation_callback + if extensions or result_claims or notification_bindings: + _require_extension_support( + 'extensions, result_claims and notification_bindings' + ) + self._extensions = extensions + self._result_claims = result_claims + self._notification_bindings = notification_bindings + self._claims_by_model = _index_claims_by_model(result_claims) if isinstance(connection_params, StdioServerParameters): # So far timeout is not configurable. Given MCP is still evolving, we @@ -948,6 +1012,10 @@ def _is_session_disconnected(self, session: ClientSession) -> bool: or getattr(write_stream, '_closed', False) ) + def _claim_for(self, result: Any) -> Optional[ResultClaim]: + """Returns the claim that owns `result`, if any is registered for it.""" + return self._claims_by_model.get(type(result)) + def _get_session_context( self, headers: Optional[Dict[str, str]] = None ) -> Optional[SessionContext]: @@ -1326,6 +1394,9 @@ async def create_session( sampling_callback=self._sampling_callback, sampling_capabilities=self._sampling_capabilities, elicitation_callback=self._elicitation_callback, + extensions=self._extensions, + result_claims=self._result_claims, + notification_bindings=self._notification_bindings, ) if is_feature_enabled(FeatureName._MCP_GRACEFUL_ERROR_HANDLING): # pylint: disable=protected-access diff --git a/src/google/adk/tools/mcp_tool/mcp_tool.py b/src/google/adk/tools/mcp_tool/mcp_tool.py index 09c4c9fe33f..25ff9ab20f4 100644 --- a/src/google/adk/tools/mcp_tool/mcp_tool.py +++ b/src/google/adk/tools/mcp_tool/mcp_tool.py @@ -35,6 +35,8 @@ from ...auth.auth_credential import AuthCredential from ...auth.auth_schemes import AuthScheme from ...auth.auth_tool import AuthConfig +from ...dependencies._mcp import CallToolResult +from ...dependencies._mcp import ClaimContext from ...dependencies._mcp import ClientSession from ...dependencies._mcp import IS_MCP_SDK_V2 from ...dependencies._mcp import McpError @@ -591,11 +593,20 @@ async def _run_async_impl( # Resolve progress callback (may be a factory that needs runtime context) resolved_callback = self._resolve_progress_callback(tool_context) + # allow_claimed lets a registered extension's result shape come back + # instead of raising. Only pass it when a claim is actually registered, so + # an unclaimed non-core result keeps failing validation as it does today. + # The isinstance check is here for the same reason as the one further + # down: a mock session manager answers every attribute with a truthy Mock, + # and that must not look like a registered claim. + claims = getattr(self._mcp_session_manager, "_claims_by_model", None) + claims_registered = isinstance(claims, dict) and bool(claims) call_coro = session.call_tool( self._mcp_tool.name, arguments=args, progress_callback=resolved_callback, meta=meta_trace_context, + **({"allow_claimed": True} if claims_registered else {}), ) # Hold the session out of the pool's idle sweep for as long as the call @@ -603,6 +614,7 @@ async def _run_async_impl( # only looks idle because its call has not come back yet must not have # its transport closed underneath it. self._mcp_session_manager._begin_session_use(final_headers) # pylint: disable=protected-access + session_context = None try: if is_feature_enabled(FeatureName._MCP_GRACEFUL_ERROR_HANDLING): # pylint: disable=protected-access # Race the tool call against the background session task so that @@ -629,6 +641,11 @@ async def _run_async_impl( # Pre-fix behavior: await the call directly. This is what causes the # ~300s hang when the underlying transport crashes. response = await call_coro + + if claims_registered and not isinstance(response, CallToolResult): + response = await self._resolve_claimed( + response, session, session_context + ) finally: self._mcp_session_manager._end_session_use(final_headers) # pylint: disable=protected-access @@ -664,6 +681,61 @@ async def _run_async_impl( ) return result + async def _resolve_claimed( + self, + response: Any, + session: ClientSession, + session_context: SessionContext | None, + ) -> CallToolResult: + """Turns an extension's claimed `tools/call` result into a normal one. + + `ClientSession` parses a claimed result but stops there; resolving it is + the caller's job, and for the tasks extension resolving means polling + until the task reaches a terminal state. Cancellation therefore has to + reach the resolver rather than orphaning it, which is what + `propagate_cancel` is for -- a resolver that never learns it was + cancelled never gets to tell the server. + + Args: + response: The claimed result, already parsed into its claim's model. + session: The session the call was made on. + session_context: The context guarding that session, when there is one. + + Returns: + The `CallToolResult` the claim resolved to. + + Raises: + RuntimeError: If no registered claim owns this result's type. + """ + claim = self._mcp_session_manager._claim_for(response) # pylint: disable=protected-access + if claim is None: + raise RuntimeError( + f"MCP server returned an unclaimed result of type {type(response)!r}" + ) + + read_timeout = getattr(session, "read_timeout_seconds", None) + resolve_coro = claim.resolve( + response, + ClaimContext( + session=session, + tool_name=self._mcp_tool.name, + read_timeout_seconds=read_timeout, + ), + ) + if isinstance(session_context, SessionContext): + resolved = await session_context._run_guarded( # pylint: disable=protected-access + resolve_coro, propagate_cancel=True + ) + else: + resolved = await resolve_coro + + # Mirror the schema enforcement the direct path gets from `call_tool`, so + # a result that arrived by way of an extension is held to the same + # contract as one that arrived inline. + if not resolved.is_error: + await session.validate_tool_result(self._mcp_tool.name, resolved) + return resolved + def _detect_error_in_response(self, response: Any) -> str | None: """Telemetry hook: returns an error type if the response indicates an error.""" # `response` is a dumped CallToolResult. `_run_async_impl` restores diff --git a/src/google/adk/tools/mcp_tool/mcp_toolset.py b/src/google/adk/tools/mcp_tool/mcp_toolset.py index dcf4a7e0c68..3b36e2eda3b 100644 --- a/src/google/adk/tools/mcp_tool/mcp_toolset.py +++ b/src/google/adk/tools/mcp_tool/mcp_toolset.py @@ -26,7 +26,9 @@ from typing import Callable from typing import Dict from typing import List +from typing import Mapping from typing import Optional +from typing import Sequence from typing import TextIO from typing import TypeVar from typing import Union @@ -43,6 +45,8 @@ from ...dependencies._mcp import ElicitationFnT from ...dependencies._mcp import ListResourcesResult from ...dependencies._mcp import ListToolsResult +from ...dependencies._mcp import NotificationBinding +from ...dependencies._mcp import ResultClaim from ...dependencies._mcp import SamplingCapability from ...dependencies._mcp import SamplingFnT from ...dependencies._mcp import StdioServerParameters @@ -55,6 +59,7 @@ from ..tool_configs import BaseToolConfig from ..tool_configs import ToolArgsConfig from .mcp_session_manager import _http_debug_var +from .mcp_session_manager import _require_extension_support from .mcp_session_manager import MCPSessionManager from .mcp_session_manager import retry_on_errors from .mcp_session_manager import SseConnectionParams @@ -112,6 +117,63 @@ class _CachedToolList: expires_at: float +def _compose_tasks_extension( + extensions: dict[str, dict[str, Any]] | None, + result_claims: Mapping[str, Sequence[ResultClaim]] | None, + *, + enable_tasks: bool, +) -> tuple[ + dict[str, dict[str, Any]] | None, + Mapping[str, Sequence[ResultClaim]] | None, +]: + """Folds the built-in tasks extension into caller-supplied extensions. + + The advertisement and the claim are added together, because a claim whose + extension is not advertised is rejected when the session is built. + + Args: + extensions: Extensions the caller asked for, if any. + result_claims: Result claims the caller asked for, if any. + enable_tasks: Whether to add the built-in tasks extension. + + Returns: + The extensions and result claims to hand to the session manager. + + Raises: + ValueError: If the installed MCP SDK is 1.x, or if the caller already + carries the tasks extension. Two claims on one result type is refused + when the session is built, deep inside a later tool call; refusing it + here says so while the cause is still in view. + """ + if not enable_tasks: + return extensions, result_claims + + _require_extension_support("enable_tasks=True") + + # Imported here rather than at module scope: the wire models in `_tasks` + # derive from `mcp_types`, which only a 2.x install carries. The check above + # is what guarantees nothing reaches this line on 1.x. + from ._tasks import make_tasks_claim # pylint: disable=g-import-not-at-top + from ._tasks import TASKS_EXTENSION_ID # pylint: disable=g-import-not-at-top + + carries_tasks = TASKS_EXTENSION_ID in (extensions or {}) or ( + TASKS_EXTENSION_ID in (result_claims or {}) + ) + if carries_tasks: + raise ValueError( + f"enable_tasks=True conflicts with the {TASKS_EXTENSION_ID!r} entry" + " already passed in extensions or result_claims. Use one or the" + " other: enable_tasks for the built-in implementation, or your own" + " entry to replace it." + ) + + composed_extensions = dict(extensions or {}) + composed_extensions[TASKS_EXTENSION_ID] = {} + composed_claims = dict(result_claims or {}) + composed_claims[TASKS_EXTENSION_ID] = [make_tasks_claim()] + return composed_extensions, composed_claims + + class McpToolset(BaseToolset): """Connects to a MCP Server, and retrieves MCP Tools into ADK Tools. @@ -170,6 +232,10 @@ def __init__( sampling_callback: SamplingFnT | None = None, sampling_capabilities: SamplingCapability | None = None, elicitation_callback: ElicitationFnT | None = None, + extensions: dict[str, dict[str, Any]] | None = None, + result_claims: Mapping[str, Sequence[ResultClaim]] | None = None, + notification_bindings: Sequence[NotificationBinding] | None = None, + enable_tasks: bool = False, credential_key: str | None = None, ): """Initializes the McpToolset. @@ -221,6 +287,26 @@ def __init__( elicitation_callback: Optional callback to handle elicitation requests from the MCP server (``elicitation/create``), including URL-mode elicitations used for out-of-band flows such as auth challenges. + extensions: MCP extensions this client advertises, keyed by extension + identifier (e.g. ``{"io.modelcontextprotocol/tasks": {}}``). Passing + any of the three extension arguments also switches session bring-up + to ``server/discover``, falling back to ``initialize()``, because an + extension capability is only live on a modern connection. Requires + MCP SDK 2.x, which is where ``ClientSession`` grew the seam; on 1.x + this raises rather than going quiet. + result_claims: Non-core ``tools/call`` result shapes to accept, keyed by + the identifier of the extension that defines them. The toolset + resolves a claimed result through its claim before handing it to the + agent, so the agent sees an ordinary tool result either way. + notification_bindings: Handlers for extension notifications. + enable_tasks: Whether to accept task-augmented tool calls. When enabled + and the server also supports the extension, a tool call the server + answers with a task handle is polled to completion and returns its + result, instead of failing. The agent sees the same tool and the same + result either way; what changes is that the operation is no longer + bounded by the lifetime of one request. Defaults to False. Servers + that do not support the extension are unaffected. Rides the extension + seam, so it requires MCP SDK 2.x and raises on 1.x. credential_key: A user specified key used to load and save this credential in a credential service. Used with auth_scheme. """ @@ -230,6 +316,11 @@ def __init__( self._sampling_callback = sampling_callback self._sampling_capabilities = sampling_capabilities self._elicitation_callback = elicitation_callback + self._enable_tasks = enable_tasks + self._extensions, self._result_claims = _compose_tasks_extension( + extensions, result_claims, enable_tasks=enable_tasks + ) + self._notification_bindings = notification_bindings if not connection_params: raise ValueError("Missing connection params in McpToolset.") @@ -260,6 +351,9 @@ def __init__( sampling_callback=self._sampling_callback, sampling_capabilities=self._sampling_capabilities, elicitation_callback=self._elicitation_callback, + extensions=self._extensions, + result_claims=self._result_claims, + notification_bindings=self._notification_bindings, ) self._auth_scheme = auth_scheme self._auth_credential = auth_credential @@ -670,6 +764,7 @@ def from_config( auth_credential=mcp_toolset_config.auth_credential, credential_key=mcp_toolset_config.credential_key, use_mcp_resources=mcp_toolset_config.use_mcp_resources, + enable_tasks=mcp_toolset_config.enable_tasks, ) def __getstate__(self): @@ -728,6 +823,7 @@ class McpToolsetConfig(BaseToolConfig): credential_key: str | None = None use_mcp_resources: bool = False + enable_tasks: bool = False @model_validator(mode="after") def _check_only_one_params_field(self): diff --git a/src/google/adk/tools/mcp_tool/session_context.py b/src/google/adk/tools/mcp_tool/session_context.py index e2e8475bd11..1963da390a5 100644 --- a/src/google/adk/tools/mcp_tool/session_context.py +++ b/src/google/adk/tools/mcp_tool/session_context.py @@ -22,12 +22,19 @@ from types import TracebackType from typing import Any from typing import Coroutine +from typing import Mapping from typing import Optional +from typing import Sequence from typing import TypeVar +import anyio + from ...dependencies._mcp import ClientSession from ...dependencies._mcp import ElicitationFnT from ...dependencies._mcp import IS_MCP_SDK_V2 +from ...dependencies._mcp import McpError +from ...dependencies._mcp import NotificationBinding +from ...dependencies._mcp import ResultClaim from ...dependencies._mcp import SamplingCapability from ...dependencies._mcp import SamplingFnT from ...features import FeatureName @@ -37,6 +44,11 @@ _T = TypeVar('_T') +# The SDK's own ceiling for one `server/discover` probe. Used here as the cap +# on the whole negotiation budget, so a slow probe cannot starve the +# `initialize()` fallback that follows it. +_DISCOVER_TIMEOUT_SECONDS = 10.0 + def _read_timeout(seconds: Optional[float]) -> Optional[float | timedelta]: """Converts a timeout in seconds to the type ``ClientSession`` expects. @@ -120,6 +132,9 @@ def __init__( sampling_callback: SamplingFnT | None = None, sampling_capabilities: SamplingCapability | None = None, elicitation_callback: ElicitationFnT | None = None, + extensions: dict[str, dict[str, Any]] | None = None, + result_claims: Mapping[str, Sequence[ResultClaim]] | None = None, + notification_bindings: Sequence[NotificationBinding] | None = None, ): """Initializes SessionContext. @@ -137,6 +152,13 @@ def __init__( sampling_capabilities: Optional capabilities for sampling. elicitation_callback: Optional callback to handle elicitation requests from the MCP server (``elicitation/create``). + extensions: MCP extensions this client advertises, keyed by extension + identifier. Supplying any extension argument also makes the session + negotiate with ``server/discover`` first, because extensions are only + live on a modern connection. + result_claims: Non-core ``tools/call`` result shapes to accept, keyed by + the identifier of the extension that defines them. + notification_bindings: Handlers for extension notifications. """ self._client = client self._timeout = timeout @@ -150,6 +172,9 @@ def __init__( self._sampling_callback = sampling_callback self._sampling_capabilities = sampling_capabilities self._elicitation_callback = elicitation_callback + self._extensions = extensions + self._result_claims = result_claims + self._notification_bindings = notification_bindings @property def session(self) -> Optional[ClientSession]: @@ -241,7 +266,12 @@ def _retrieve_exception(t: asyncio.Task[None]) -> None: return self._session # type: ignore[return-value] - async def _run_guarded(self, coro: Coroutine[Any, Any, _T]) -> _T: + async def _run_guarded( + self, + coro: Coroutine[Any, Any, _T], + *, + propagate_cancel: bool = False, + ) -> _T: """Run a coroutine while monitoring the background session task. Races the given coroutine against the background task. If the task @@ -251,6 +281,13 @@ async def _run_guarded(self, coro: Coroutine[Any, Any, _T]) -> _T: Args: coro: The coroutine to run (e.g. session.call_tool(...)). + propagate_cancel: Whether to cancel ``coro`` when this call is itself + cancelled. ``asyncio.wait`` does not cancel what it waits on, so + by default a cancelled caller leaves ``coro`` running detached. + That is tolerable for a single request, which the transport will + eventually fail, but not for a coroutine that has cleanup of its + own to do on the wire: it would never be told to run it. Off by + default so the existing call path keeps its current semantics. Returns: The result of the coroutine. @@ -273,10 +310,19 @@ async def _run_guarded(self, coro: Coroutine[Any, Any, _T]) -> _T: coro_task = asyncio.ensure_future(coro) - done, _ = await asyncio.wait( - [coro_task, self._task], - return_when=asyncio.FIRST_COMPLETED, - ) + try: + done, _ = await asyncio.wait( + [coro_task, self._task], + return_when=asyncio.FIRST_COMPLETED, + ) + except asyncio.CancelledError: + if propagate_cancel and not coro_task.done(): + coro_task.cancel() + try: + await coro_task + except BaseException: + pass + raise if coro_task in done: # If the coroutine itself raised, the exception propagates as-is @@ -334,6 +380,68 @@ async def __aexit__( ) -> None: await self.close() + @property + def _extension_kwargs(self) -> dict[str, Any]: + """The extension arguments to hand `ClientSession`, when there are any. + + Spread rather than passed as three `None`s: MCP SDK 1.x declares none of + these parameters, and naming one there is a `TypeError`. Nothing + configures an extension on 1.x -- the opt-in is refused before a session + manager exists -- so this is empty and the construction below is the one + 1.x has always made. + """ + if not self._wants_extensions: + return {} + return { + 'extensions': self._extensions, + 'result_claims': self._result_claims, + 'notification_bindings': self._notification_bindings, + } + + @property + def _wants_extensions(self) -> bool: + """Whether any MCP extension was configured on this session.""" + return bool( + self._extensions or self._result_claims or self._notification_bindings + ) + + async def _negotiate(self, session: ClientSession) -> None: + """Brings `session` up, preferring `server/discover` when it can matter. + + `initialize()` always performs the pre-2026 handshake, and an extension + capability has nowhere to ride on that wire -- the SDK drops + claim-bearing identifiers from the advertisement at legacy protocol + versions, so a claim registered on this session would never fire. Only + `server/discover` reaches a version where extensions are live. + + A server that predates `server/discover` answers it with an error, so the + probe is bounded and falls back rather than failing the session. The + bound matters: the probe and the fallback share one bring-up budget, and + an unbounded probe against a server that simply never answers would spend + all of it and leave nothing for `initialize()`. + """ + if not self._wants_extensions: + await session.initialize() + return + + budget = min( + self._timeout or _DISCOVER_TIMEOUT_SECONDS, _DISCOVER_TIMEOUT_SECONDS + ) + try: + with anyio.fail_after(budget / 2): + session.adopt(await session.discover()) + return + except (McpError, RuntimeError, TimeoutError) as e: + # RuntimeError is what `adopt` raises when the server and this client + # share no modern protocol version. + logger.warning( + 'MCP extensions were requested but server/discover is unavailable' + ' (%s); falling back to initialize(). Extensions, including tasks,' + ' stay inactive on this session.', + e, + ) + await session.initialize() + async def _run(self) -> None: """Run the complete session context within a single task.""" try: @@ -372,6 +480,7 @@ async def _run(self) -> None: sampling_callback=self._sampling_callback, sampling_capabilities=self._sampling_capabilities, elicitation_callback=self._elicitation_callback, + **self._extension_kwargs, ) ) else: @@ -384,6 +493,7 @@ async def _run(self) -> None: sampling_callback=self._sampling_callback, sampling_capabilities=self._sampling_capabilities, elicitation_callback=self._elicitation_callback, + **self._extension_kwargs, ) ) # pylint: disable-next=protected-access @@ -391,12 +501,12 @@ async def _run(self) -> None: # Use anyio.fail_after to keep session.initialize within the AnyIO # cancel scope instead of asyncio.wait_for which runs in a nested # task. - import anyio - with anyio.fail_after(self._timeout): - await session.initialize() + await self._negotiate(session) else: - await asyncio.wait_for(session.initialize(), timeout=self._timeout) + await asyncio.wait_for( + self._negotiate(session), timeout=self._timeout + ) logger.debug('Session has been successfully initialized') self._session = session diff --git a/tests/unittests/test_samples.py b/tests/unittests/test_samples.py index 186500ec4f7..3b32bdf5b87 100644 --- a/tests/unittests/test_samples.py +++ b/tests/unittests/test_samples.py @@ -29,6 +29,7 @@ from google.adk.apps.app import App from google.adk.cli.agent_test_runner import test_agent_replay as _test_agent_replay from google.adk.cli.utils.agent_loader import AgentLoader +from google.adk.dependencies._mcp import IS_MCP_SDK_V2 from google.adk.events import Event from google.genai import types import pytest @@ -143,6 +144,15 @@ def test_sample(sample_dir: Path, test_file: Path, monkeypatch): ), } +# Added conditionally rather than listed above, because the reason is the +# installed SDK and not the sample: it opts into the MCP Tasks extension, which +# rides a seam that exists only in MCP SDK 2.x, and `McpToolset` refuses the +# opt-in on 1.x at construction. On 2.x the sample has to load like any other. +if not IS_MCP_SDK_V2: + SKIP_LOAD["mcp/mcp_tasks_agent"] = ( + "opts into the MCP Tasks extension, which needs MCP SDK 2.x" + ) + # Samples whose own code is currently broken against the ADK API. Loading them # fails today; remove the entry once the sample is fixed. XFAIL_LOAD = { diff --git a/tests/unittests/tools/mcp_tool/_sdk_compat.py b/tests/unittests/tools/mcp_tool/_sdk_compat.py index ed4df773204..ba0f6fa422c 100644 --- a/tests/unittests/tools/mcp_tool/_sdk_compat.py +++ b/tests/unittests/tools/mcp_tool/_sdk_compat.py @@ -26,9 +26,23 @@ from google.adk.dependencies._mcp import IS_MCP_SDK_V2 from google.adk.dependencies._mcp import McpError +import pytest _CAMEL_BOUNDARY = re.compile(r'(? Any: """Reads a model field under whichever spelling the installed SDK uses. diff --git a/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py b/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py index d0282be0301..8753b13f4e5 100644 --- a/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py +++ b/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py @@ -19,6 +19,7 @@ import logging import sys import time +from types import SimpleNamespace from unittest.mock import ANY from unittest.mock import AsyncMock from unittest.mock import Mock @@ -51,6 +52,8 @@ from mcp import StdioServerParameters import pytest +from ._sdk_compat import requires_sdk_v2 + try: from google.auth.aio.transport.sessions import AsyncAuthorizedSession @@ -513,6 +516,68 @@ async def test_create_session_stdio_new(self): # Verify enter_async_context was called (which internally calls __aenter__) mock_exit_stack.enter_async_context.assert_called_once() + @pytest.mark.asyncio + @requires_sdk_v2 + async def test_create_session_passes_extension_arguments(self): + """The three extension arguments are forwarded to the SessionContext.""" + extensions = {"io.modelcontextprotocol/tasks": {}} + result_claims = { + "io.modelcontextprotocol/tasks": [SimpleNamespace(model=object)] + } + notification_bindings = ["sentinel-binding"] + + manager = MCPSessionManager( + self.mock_stdio_connection_params, + extensions=extensions, + result_claims=result_claims, + notification_bindings=notification_bindings, + ) + mock_exit_stack = MockAsyncExitStack() + with patch( + "google.adk.tools.mcp_tool.mcp_session_manager.stdio_client" + ) as mock_stdio: + with patch( + "google.adk.tools.mcp_tool.mcp_session_manager.AsyncExitStack" + ) as mock_exit_stack_class: + with patch( + "google.adk.tools.mcp_tool.mcp_session_manager.SessionContext" + ) as mock_session_context_class: + mock_exit_stack_class.return_value = mock_exit_stack + mock_stdio.return_value = AsyncMock() + mock_session = AsyncMock() + mock_session_context = MockSessionContext(session=mock_session) + mock_session_context_class.return_value = mock_session_context + mock_exit_stack.enter_async_context.return_value = mock_session + await manager.create_session() + _, kwargs = mock_session_context_class.call_args + assert kwargs["extensions"] is extensions + assert kwargs["result_claims"] is result_claims + assert kwargs["notification_bindings"] is notification_bindings + + @requires_sdk_v2 + def test_claim_for_finds_the_claim_that_owns_a_result(self): + """A claimed result is routed back to its claim by model type.""" + + class _Claimed: + pass + + claim = SimpleNamespace(model=_Claimed) + manager = MCPSessionManager( + self.mock_stdio_connection_params, + extensions={"vendor/ext": {}}, + result_claims={"vendor/ext": [claim]}, + ) + + assert manager._claim_for(_Claimed()) is claim + assert manager._claim_for(object()) is None + + def test_claim_index_is_empty_without_claims(self): + """No claims registered means nothing to resolve, and nothing claimed.""" + manager = MCPSessionManager(self.mock_stdio_connection_params) + + assert not manager._claims_by_model + assert manager._claim_for(object()) is None + @pytest.mark.asyncio async def test_create_session_passes_elicitation_callback(self): """Elicitation callback is forwarded to the SessionContext.""" diff --git a/tests/unittests/tools/mcp_tool/test_mcp_tasks.py b/tests/unittests/tools/mcp_tool/test_mcp_tasks.py new file mode 100644 index 00000000000..5d68ad53646 --- /dev/null +++ b/tests/unittests/tools/mcp_tool/test_mcp_tasks.py @@ -0,0 +1,326 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import asyncio + +import pytest + +# Skipped at collection, not per test: the module under test builds on the +# client extension seam, which is the one part of the SDK that 1.x does not +# have under any name, so this file cannot be imported there at all. +pytest.importorskip( + 'mcp.client.extension', + reason='the MCP Tasks extension requires MCP SDK 2.x', +) + +# pylint: disable=g-import-not-at-top +from google.adk.tools.mcp_tool._tasks import _CancelTaskRequest +from google.adk.tools.mcp_tool._tasks import _GetTaskRequest +from google.adk.tools.mcp_tool._tasks import _GetTaskResult +from google.adk.tools.mcp_tool._tasks import _poll_interval_seconds +from google.adk.tools.mcp_tool._tasks import _resolve_task +from google.adk.tools.mcp_tool._tasks import _TaskParams +from google.adk.tools.mcp_tool._tasks import _TaskResult +from google.adk.tools.mcp_tool._tasks import make_tasks_claim +from google.adk.tools.mcp_tool._tasks import TASKS_EXTENSION_ID +from mcp.client.extension import ClaimContext +from mcp.shared.exceptions import MCPError + +# pylint: enable=g-import-not-at-top + +_CREATED = '2026-08-21T10:30:00Z' + + +def _task_handle(**overrides) -> _TaskResult: + """A `resultType: "task"` answer, working unless told otherwise.""" + fields = { + 'task_id': 'task-1', + 'status': 'working', + 'created_at': _CREATED, + 'last_updated_at': _CREATED, + 'ttl_ms': 60000, + 'poll_interval_ms': 1, + } + fields.update(overrides) + return _TaskResult(**fields) + + +def _task_state(status: str, **overrides) -> _GetTaskResult: + """One `tasks/get` answer.""" + fields = { + 'task_id': 'task-1', + 'status': status, + 'created_at': _CREATED, + 'last_updated_at': _CREATED, + 'ttl_ms': 60000, + 'poll_interval_ms': 1, + } + fields.update(overrides) + return _GetTaskResult(**fields) + + +class _FakeSession: + """A session that replays a scripted sequence of `tasks/*` answers.""" + + def __init__(self, answers): + self._answers = list(answers) + self.requests = [] + + async def send_request(self, request, result_type, **kwargs): + del result_type, kwargs + self.requests.append(request) + if isinstance(request, _CancelTaskRequest): + return None + answer = self._answers.pop(0) + if isinstance(answer, Exception): + raise answer + return answer + + @property + def polls(self): + return [r for r in self.requests if isinstance(r, _GetTaskRequest)] + + @property + def cancels(self): + return [r for r in self.requests if isinstance(r, _CancelTaskRequest)] + + +def _ctx(session) -> ClaimContext: + return ClaimContext( + session=session, tool_name='slow_tool', read_timeout_seconds=None + ) + + +class TestWireModels: + """The models have to match SEP-2663 byte for byte on the wire.""" + + def test_task_handle_round_trips_the_specified_shape(self): + wire = { + 'resultType': 'task', + 'taskId': '786512e2-9e0d-44bd-8f29-789f320fe840', + 'status': 'working', + 'statusMessage': 'The operation is now in progress.', + 'createdAt': '2025-11-25T10:30:00Z', + 'lastUpdatedAt': '2025-11-25T10:40:00Z', + 'ttlMs': 60000, + 'pollIntervalMs': 5000, + } + + parsed = _TaskResult.model_validate(wire) + + assert parsed.task_id == '786512e2-9e0d-44bd-8f29-789f320fe840' + assert parsed.ttl_ms == 60000 + assert parsed.poll_interval_ms == 5000 + assert parsed.model_dump(by_alias=True, exclude_none=True) == wire + + def test_ttl_is_required_and_nullable(self): + """The specification makes ttlMs required, but allows a null value.""" + assert _task_handle(ttl_ms=None).ttl_ms is None + + with pytest.raises(Exception): + _TaskResult.model_validate({ + 'resultType': 'task', + 'taskId': 't', + 'status': 'working', + 'createdAt': _CREATED, + 'lastUpdatedAt': _CREATED, + }) + + def test_task_requests_carry_the_task_id_as_their_routing_name(self): + """`Mcp-Name` is built from name_param, and tasks/* must set it.""" + assert _GetTaskRequest.name_param == 'taskId' + assert _CancelTaskRequest.name_param == 'taskId' + + request = _GetTaskRequest(params=_TaskParams(task_id='task-1')) + + assert request.model_dump(by_alias=True, exclude_none=True) == { + 'method': 'tasks/get', + 'params': {'taskId': 'task-1'}, + } + + def test_claim_is_registered_for_the_task_result_type(self): + claim = make_tasks_claim() + + assert claim.result_type == 'task' + assert claim.model is _TaskResult + assert TASKS_EXTENSION_ID == 'io.modelcontextprotocol/tasks' + + +class TestPollInterval: + """The server states a preference; it is honored, within reason.""" + + def test_absent_interval_falls_back_to_a_default(self): + assert _poll_interval_seconds(None) == 1.0 + + def test_stated_interval_is_honored(self): + assert _poll_interval_seconds(2500) == 2.5 + + def test_zero_is_clamped_so_the_loop_cannot_spin(self): + assert _poll_interval_seconds(0) == 0.1 + + def test_an_absurd_interval_is_capped(self): + assert _poll_interval_seconds(600000) == 30.0 + + +class TestResolveTask: + """The state machine that turns a task handle into a tool result.""" + + @pytest.mark.asyncio + async def test_polls_until_completion_and_returns_the_result(self): + session = _FakeSession([ + _task_state('working'), + _task_state( + 'completed', + result={'content': [{'type': 'text', 'text': 'done'}]}, + ), + ]) + + result = await _resolve_task(_task_handle(), _ctx(session)) + + assert not result.is_error + assert result.content[0].text == 'done' + assert len(session.polls) == 2 + + @pytest.mark.asyncio + async def test_a_handle_that_is_already_terminal_is_still_read(self): + """The result only ever travels on tasks/get, never on the handle.""" + session = _FakeSession([ + _task_state('completed', result={'content': []}), + ]) + + result = await _resolve_task( + _task_handle(status='completed'), _ctx(session) + ) + + assert not result.is_error + assert len(session.polls) == 1 + + @pytest.mark.asyncio + async def test_a_failed_task_becomes_a_failed_tool_result(self): + """A tool that fails reports isError; a task should not change that.""" + session = _FakeSession([ + _task_state( + 'failed', error={'code': -32000, 'message': 'device offline'} + ), + ]) + + result = await _resolve_task(_task_handle(), _ctx(session)) + + assert result.is_error + assert 'device offline' in result.content[0].text + assert '-32000' in result.content[0].text + + @pytest.mark.asyncio + async def test_a_server_cancelled_task_becomes_a_failed_tool_result(self): + session = _FakeSession([_task_state('cancelled')]) + + result = await _resolve_task(_task_handle(), _ctx(session)) + + assert result.is_error + assert 'cancelled by the server' in result.content[0].text + + @pytest.mark.asyncio + async def test_input_required_is_reported_and_the_task_released(self): + """Nothing here can answer, so the server should stop holding the task.""" + session = _FakeSession([ + _task_state('input_required', input_requests={'q1': {}}), + ]) + + result = await _resolve_task(_task_handle(), _ctx(session)) + + assert result.is_error + assert 'requires additional input' in result.content[0].text + assert len(session.cancels) == 1 + + @pytest.mark.asyncio + async def test_an_unreadable_task_becomes_a_failed_tool_result(self): + """An expired handle answers with an error, and retrying cannot help.""" + session = _FakeSession([MCPError(code=-32602, message='unknown taskId')]) + + result = await _resolve_task(_task_handle(), _ctx(session)) + + assert result.is_error + assert 'unknown taskId' in result.content[0].text + assert len(session.polls) == 1 + + @pytest.mark.asyncio + async def test_a_completed_task_without_a_result_is_reported(self): + session = _FakeSession([_task_state('completed')]) + + result = await _resolve_task(_task_handle(), _ctx(session)) + + assert result.is_error + assert 'without a result' in result.content[0].text + + @pytest.mark.asyncio + async def test_an_unusable_payload_does_not_escape_as_an_exception(self): + """The caller has no way to guard a resolver, so nothing may leak out.""" + session = _FakeSession([ + _task_state('completed', result={'content': 'not a content list'}), + ]) + + result = await _resolve_task(_task_handle(), _ctx(session)) + + assert result.is_error + assert 'not a valid tool result' in result.content[0].text + + @pytest.mark.asyncio + async def test_cancelling_the_caller_cancels_the_task_on_the_server(self): + """The whole point of propagating cancellation into the resolver.""" + started = asyncio.Event() + + class _SlowSession(_FakeSession): + + async def send_request(self, request, result_type, **kwargs): + if isinstance(request, _GetTaskRequest): + started.set() + await asyncio.sleep(30) + return await super().send_request(request, result_type, **kwargs) + + session = _SlowSession([]) + task = asyncio.create_task(_resolve_task(_task_handle(), _ctx(session))) + await asyncio.wait_for(started.wait(), timeout=5) + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + assert len(session.cancels) == 1 + assert session.cancels[0].params.task_id == 'task-1' + + @pytest.mark.asyncio + async def test_the_interval_follows_the_latest_value_the_server_sends(self): + """A server may change its mind between polls, and is obeyed.""" + slept = [] + session = _FakeSession([ + _task_state('working', poll_interval_ms=2000), + _task_state('completed', result={'content': []}), + ]) + + real_sleep = asyncio.sleep + + async def _record(delay): + slept.append(delay) + await real_sleep(0) + + asyncio.sleep = _record + try: + await _resolve_task(_task_handle(poll_interval_ms=50), _ctx(session)) + finally: + asyncio.sleep = real_sleep + + # First from the handle, then from what the first poll reported. + assert slept == [0.1, 2.0] diff --git a/tests/unittests/tools/mcp_tool/test_mcp_tool.py b/tests/unittests/tools/mcp_tool/test_mcp_tool.py index 57856043cac..ab22e88579b 100644 --- a/tests/unittests/tools/mcp_tool/test_mcp_tool.py +++ b/tests/unittests/tools/mcp_tool/test_mcp_tool.py @@ -50,6 +50,7 @@ from ._sdk_compat import expected_tool_result from ._sdk_compat import make_mcp_error +from ._sdk_compat import requires_sdk_v2 from ._sdk_compat import sdk_progress_fn_t @@ -1827,6 +1828,183 @@ async def mock_call_tool(*args, **kwargs): assert debug_info[0]["status_code"] == 403 +@requires_sdk_v2 +class TestMCPToolClaimedResults: + """Tests for resolving an extension's claimed `tools/call` result. + + `ClientSession` parses a claimed result but never resolves it, so the + toolset has to do that itself or the extension seam carries nothing. + """ + + def setup_method(self): + self.mock_mcp_tool = MockMCPTool(name="test_tool") + self.mock_session = AsyncMock() + self.mock_session.validate_tool_result = AsyncMock() + self.mock_session_manager = Mock(spec=MCPSessionManager) + self.mock_session_manager.create_session = AsyncMock( + return_value=self.mock_session + ) + self.mock_session_manager._get_session_context = Mock(return_value=None) + + def _tool_context(self): + tool_context = ToolContext(invocation_context=Mock()) + tool_context.function_call_id = "test-call-id" + return tool_context + + def _register_claim(self, resolve, model): + claim = SimpleNamespace(model=model, resolve=resolve) + self.mock_session_manager._claims_by_model = {model: claim} + self.mock_session_manager._claim_for = Mock( + side_effect=lambda r: {model: claim}.get(type(r)) + ) + return claim + + @pytest.mark.asyncio + async def test_claimed_result_is_resolved_before_the_agent_sees_it(self): + """The agent gets the resolved CallToolResult, not the claim's model.""" + + class _Claimed: + pass + + resolved = CallToolResult( + content=[TextContent(type="text", text="resolved")] + ) + resolve = AsyncMock(return_value=resolved) + self._register_claim(resolve, _Claimed) + claimed = _Claimed() + self.mock_session.call_tool = AsyncMock(return_value=claimed) + + tool = MCPTool( + mcp_tool=self.mock_mcp_tool, + mcp_session_manager=self.mock_session_manager, + ) + result = await tool._run_async_impl( + args={}, tool_context=self._tool_context(), credential=None + ) + + # The same dict a result that arrived inline would have produced: going + # through an extension changes how it travelled, not what the agent gets. + assert result == expected_tool_result(resolved) + # The claim gets the parsed result plus the context it needs to poll. + passed_result, ctx = resolve.await_args.args + assert passed_result is claimed + assert ctx.tool_name == "test_tool" + assert ctx.session is self.mock_session + # A resolved result is held to the same output schema as a direct one. + self.mock_session.validate_tool_result.assert_awaited_once_with( + "test_tool", resolved + ) + + @pytest.mark.asyncio + async def test_allow_claimed_is_passed_only_when_a_claim_exists(self): + """Without claims the call is unchanged, so unclaimed results still fail.""" + self.mock_session_manager._claims_by_model = {} + self.mock_session.call_tool = AsyncMock( + return_value=CallToolResult(content=[]) + ) + + tool = MCPTool( + mcp_tool=self.mock_mcp_tool, + mcp_session_manager=self.mock_session_manager, + ) + await tool._run_async_impl( + args={}, tool_context=self._tool_context(), credential=None + ) + + assert "allow_claimed" not in self.mock_session.call_tool.await_args.kwargs + + @pytest.mark.asyncio + async def test_allow_claimed_is_passed_when_a_claim_exists(self): + """A registered claim opts the call into non-core result shapes.""" + + class _Claimed: + pass + + self._register_claim(AsyncMock(), _Claimed) + self.mock_session.call_tool = AsyncMock( + return_value=CallToolResult(content=[]) + ) + + tool = MCPTool( + mcp_tool=self.mock_mcp_tool, + mcp_session_manager=self.mock_session_manager, + ) + await tool._run_async_impl( + args={}, tool_context=self._tool_context(), credential=None + ) + + assert self.mock_session.call_tool.await_args.kwargs["allow_claimed"] + + @pytest.mark.asyncio + async def test_an_ordinary_result_is_not_sent_through_a_claim(self): + """A core result short-circuits resolution even when claims exist.""" + + class _Claimed: + pass + + resolve = AsyncMock() + self._register_claim(resolve, _Claimed) + plain = CallToolResult(content=[TextContent(type="text", text="plain")]) + self.mock_session.call_tool = AsyncMock(return_value=plain) + + tool = MCPTool( + mcp_tool=self.mock_mcp_tool, + mcp_session_manager=self.mock_session_manager, + ) + result = await tool._run_async_impl( + args={}, tool_context=self._tool_context(), credential=None + ) + + resolve.assert_not_awaited() + assert result == expected_tool_result(plain) + + @pytest.mark.asyncio + async def test_unclaimed_non_core_result_raises(self): + """A result no claim owns is a server bug, and must not pass silently.""" + + class _Claimed: + pass + + class _Unknown: + pass + + self._register_claim(AsyncMock(), _Claimed) + self.mock_session.call_tool = AsyncMock(return_value=_Unknown()) + + tool = MCPTool( + mcp_tool=self.mock_mcp_tool, + mcp_session_manager=self.mock_session_manager, + ) + with pytest.raises(RuntimeError, match="unclaimed result"): + await tool._run_async_impl( + args={}, tool_context=self._tool_context(), credential=None + ) + + @pytest.mark.asyncio + async def test_failed_resolution_is_not_schema_checked(self): + """An error result has no structured output to validate.""" + + class _Claimed: + pass + + failed = CallToolResult( + content=[TextContent(type="text", text="boom")], is_error=True + ) + self._register_claim(AsyncMock(return_value=failed), _Claimed) + self.mock_session.call_tool = AsyncMock(return_value=_Claimed()) + + tool = MCPTool( + mcp_tool=self.mock_mcp_tool, + mcp_session_manager=self.mock_session_manager, + ) + result = await tool._run_async_impl( + args={}, tool_context=self._tool_context(), credential=None + ) + + self.mock_session.validate_tool_result.assert_not_awaited() + assert result["isError"] is True + + class TestMCPToolGracefulErrorHandling: """Tests for the _MCP_GRACEFUL_ERROR_HANDLING feature flag. diff --git a/tests/unittests/tools/mcp_tool/test_mcp_toolset.py b/tests/unittests/tools/mcp_tool/test_mcp_toolset.py index cbce212f6e6..55a5f147ee2 100644 --- a/tests/unittests/tools/mcp_tool/test_mcp_toolset.py +++ b/tests/unittests/tools/mcp_tool/test_mcp_toolset.py @@ -19,6 +19,7 @@ import pickle import sys import time +from types import SimpleNamespace from unittest.mock import AsyncMock from unittest.mock import MagicMock from unittest.mock import Mock @@ -56,6 +57,14 @@ from mcp.types import TextResourceContents import pytest +from ._sdk_compat import requires_sdk_v1 +from ._sdk_compat import requires_sdk_v2 + +# The extension's wire identifier. Spelled out rather than imported from +# `_tasks`, which needs MCP SDK 2.x while this file runs on both -- and an +# identifier the protocol fixes is worth pinning independently anyway. +TASKS_EXTENSION_ID = "io.modelcontextprotocol/tasks" + class MockMCPTool: """Mock MCP Tool for testing.""" @@ -929,6 +938,125 @@ async def mock_sampling_handler(messages, params=None, context=None): assert result["role"] == "assistant" assert result["content"]["text"] == "sampling response" + @requires_sdk_v2 + def test_extension_arguments_plumbed_to_session_manager(self): + """The three extension arguments reach the session manager.""" + extensions = {"io.modelcontextprotocol/tasks": {}} + result_claims = { + "io.modelcontextprotocol/tasks": [SimpleNamespace(model=object)] + } + notification_bindings = ["sentinel-binding"] + + toolset = McpToolset( + connection_params=self.mock_stdio_params, + extensions=extensions, + result_claims=result_claims, + notification_bindings=notification_bindings, + ) + + # pylint: disable=protected-access + assert toolset._extensions is extensions + assert toolset._mcp_session_manager._extensions is extensions + assert toolset._mcp_session_manager._result_claims is result_claims + assert ( + toolset._mcp_session_manager._notification_bindings + is notification_bindings + ) + # pylint: enable=protected-access + + @requires_sdk_v1 + def test_extension_arguments_are_refused_on_sdk_v1(self): + """An extension opt-in raises where the seam does not exist. + + The alternative would be to accept it and do nothing, which on a tool + call the extension exists to make long-running reads as a hang. The + message has to name the SDK, because nothing else about the install + explains why an argument the signature accepts had no effect. + """ + with pytest.raises(ValueError, match="requires MCP SDK 2.x"): + McpToolset( + connection_params=self.mock_stdio_params, + extensions={"io.modelcontextprotocol/tasks": {}}, + ) + + @requires_sdk_v2 + def test_enable_tasks_registers_the_extension_and_its_claim(self): + """The advertisement and the claim have to be added together.""" + toolset = McpToolset( + connection_params=self.mock_stdio_params, enable_tasks=True + ) + + # pylint: disable=protected-access + assert toolset._extensions == {TASKS_EXTENSION_ID: {}} + claims = toolset._result_claims[TASKS_EXTENSION_ID] + assert [c.result_type for c in claims] == ["task"] + assert toolset._mcp_session_manager._claims_by_model + # pylint: enable=protected-access + + @requires_sdk_v2 + def test_enable_tasks_preserves_caller_supplied_extensions(self): + """Opting into tasks must not drop an extension the caller carried.""" + toolset = McpToolset( + connection_params=self.mock_stdio_params, + enable_tasks=True, + extensions={"vendor/ext": {"setting": True}}, + ) + + # pylint: disable=protected-access + assert toolset._extensions == { + "vendor/ext": {"setting": True}, + TASKS_EXTENSION_ID: {}, + } + # pylint: enable=protected-access + + @requires_sdk_v1 + def test_enable_tasks_is_refused_on_sdk_v1(self): + """The built-in tasks path rides the seam, so it is refused with it. + + Separate from the generic extension refusal because `enable_tasks` is the + argument a caller is most likely to reach for without knowing there is a + seam underneath, and because the composition it triggers imports models + a 1.x install does not carry. + """ + with pytest.raises(ValueError, match="requires MCP SDK 2.x"): + McpToolset(connection_params=self.mock_stdio_params, enable_tasks=True) + + def test_enable_tasks_is_off_by_default(self): + toolset = McpToolset(connection_params=self.mock_stdio_params) + + # pylint: disable=protected-access + assert toolset._enable_tasks is False + assert toolset._extensions is None + # pylint: enable=protected-access + + @requires_sdk_v2 + def test_enable_tasks_conflicting_with_a_caller_entry_is_refused(self): + """Two claims on one result type only fails much later otherwise.""" + with pytest.raises(ValueError, match="conflicts with"): + McpToolset( + connection_params=self.mock_stdio_params, + enable_tasks=True, + extensions={TASKS_EXTENSION_ID: {}}, + ) + + with pytest.raises(ValueError, match="conflicts with"): + McpToolset( + connection_params=self.mock_stdio_params, + enable_tasks=True, + result_claims={TASKS_EXTENSION_ID: [SimpleNamespace(model=object)]}, + ) + + def test_extension_arguments_default_to_none(self): + """A toolset without extensions advertises none.""" + toolset = McpToolset(connection_params=self.mock_stdio_params) + + # pylint: disable=protected-access + assert toolset._extensions is None + assert toolset._result_claims is None + assert toolset._notification_bindings is None + assert not toolset._mcp_session_manager._claims_by_model + # pylint: enable=protected-access + @pytest.mark.asyncio async def test_elicitation_callback_plumbed_to_session_manager(self): """Elicitation callback reaches the session manager unchanged.""" diff --git a/tests/unittests/tools/mcp_tool/test_session_context.py b/tests/unittests/tools/mcp_tool/test_session_context.py index 892c2413e44..6efdfb63dd7 100644 --- a/tests/unittests/tools/mcp_tool/test_session_context.py +++ b/tests/unittests/tools/mcp_tool/test_session_context.py @@ -22,6 +22,7 @@ from unittest.mock import Mock from unittest.mock import patch +from google.adk.dependencies._mcp import McpError from google.adk.features import FeatureName from google.adk.features._feature_registry import temporary_feature_override from google.adk.tools.mcp_tool.session_context import _format_exception @@ -31,6 +32,8 @@ from mcp import ClientSession import pytest +from ._sdk_compat import requires_sdk_v2 + class MockClientSession: """Mock ClientSession for testing.""" @@ -725,6 +728,147 @@ async def elicitation_callback(context, params): assert kwargs['elicitation_callback'] is elicitation_callback +class TestSessionContextExtensions: + """Tests for the mcp 2.x extension seam and its protocol negotiation.""" + + def _patch_client_session(self): + return patch( + 'google.adk.tools.mcp_tool.session_context.ClientSession', + autospec=True, + ) + + @staticmethod + def _make_session(mock_client_session_class): + session = mock_client_session_class.return_value + # The session the context actually uses is what the exit stack enters, + # not the constructor's return value. + session.__aenter__ = AsyncMock(return_value=session) + session.__aexit__ = AsyncMock(return_value=None) + session.initialize = AsyncMock() + session.discover = AsyncMock(return_value='discover-result') + session.adopt = Mock() + return session + + @pytest.mark.asyncio + @requires_sdk_v2 + async def test_passes_extension_arguments_to_client_session(self): + """The three extension arguments reach ClientSession verbatim.""" + extensions = {'io.modelcontextprotocol/tasks': {}} + result_claims = {'io.modelcontextprotocol/tasks': ['sentinel-claim']} + notification_bindings = ['sentinel-binding'] + + context = SessionContext( + client=MockClient(), + timeout=5.0, + sse_read_timeout=None, + extensions=extensions, + result_claims=result_claims, + notification_bindings=notification_bindings, + ) + with self._patch_client_session() as mock_client_session_class: + self._make_session(mock_client_session_class) + async with context: + pass + + _, kwargs = mock_client_session_class.call_args + assert kwargs['extensions'] is extensions + assert kwargs['result_claims'] is result_claims + assert kwargs['notification_bindings'] is notification_bindings + + @pytest.mark.asyncio + async def test_defaults_to_no_extensions(self): + """Without extension arguments the three kwargs are not passed at all. + + Absent rather than `None`: MCP SDK 1.x declares none of them, so naming + one would be a `TypeError` on the major the default path still runs on. + """ + context = SessionContext( + client=MockClient(), timeout=5.0, sse_read_timeout=None + ) + with self._patch_client_session() as mock_client_session_class: + self._make_session(mock_client_session_class) + async with context: + pass + + _, kwargs = mock_client_session_class.call_args + assert 'extensions' not in kwargs + assert 'result_claims' not in kwargs + assert 'notification_bindings' not in kwargs + + @pytest.mark.asyncio + async def test_negotiates_with_initialize_when_no_extensions(self): + """The default path is unchanged: initialize(), and no discover probe.""" + context = SessionContext( + client=MockClient(), timeout=5.0, sse_read_timeout=None + ) + with self._patch_client_session() as mock_client_session_class: + session = self._make_session(mock_client_session_class) + async with context: + pass + + session.initialize.assert_awaited_once() + session.discover.assert_not_awaited() + + @pytest.mark.asyncio + @requires_sdk_v2 + async def test_negotiates_with_discover_when_extensions_requested(self): + """Extensions only bind on a modern connection, so discover comes first.""" + context = SessionContext( + client=MockClient(), + timeout=5.0, + sse_read_timeout=None, + extensions={'io.modelcontextprotocol/tasks': {}}, + ) + with self._patch_client_session() as mock_client_session_class: + session = self._make_session(mock_client_session_class) + async with context: + pass + + session.discover.assert_awaited_once() + session.adopt.assert_called_once_with('discover-result') + session.initialize.assert_not_awaited() + + @pytest.mark.asyncio + @requires_sdk_v2 + async def test_falls_back_to_initialize_when_discover_unsupported(self): + """A server too old for server/discover still gets a working session.""" + context = SessionContext( + client=MockClient(), + timeout=5.0, + sse_read_timeout=None, + extensions={'io.modelcontextprotocol/tasks': {}}, + ) + with self._patch_client_session() as mock_client_session_class: + session = self._make_session(mock_client_session_class) + session.discover = AsyncMock( + side_effect=McpError(code=-32601, message='Method not found') + ) + async with context: + pass + + session.discover.assert_awaited_once() + session.initialize.assert_awaited_once() + session.adopt.assert_not_called() + + @pytest.mark.asyncio + @requires_sdk_v2 + async def test_falls_back_to_initialize_when_no_mutual_version(self): + """`adopt` raising RuntimeError is a negotiation failure, not a crash.""" + context = SessionContext( + client=MockClient(), + timeout=5.0, + sse_read_timeout=None, + result_claims={'io.modelcontextprotocol/tasks': ['sentinel-claim']}, + ) + with self._patch_client_session() as mock_client_session_class: + session = self._make_session(mock_client_session_class) + session.adopt = Mock(side_effect=RuntimeError('no mutual version')) + async with context: + pass + + session.initialize.assert_awaited_once() + + class TestSessionContextIsTaskAlive: """Tests for the SessionContext._is_task_alive property.""" @@ -764,6 +908,71 @@ class TestSessionContextRunGuarded: crashes immediately. """ + @staticmethod + async def _started_context(): + context = SessionContext(MockClient(), timeout=5.0, sse_read_timeout=None) + with patch( + 'google.adk.tools.mcp_tool.session_context.ClientSession', + autospec=True, + ) as mock_client_session_class: + session = mock_client_session_class.return_value + session.__aenter__ = AsyncMock(return_value=session) + session.__aexit__ = AsyncMock(return_value=None) + session.initialize = AsyncMock() + await context.start() + return context + + @pytest.mark.asyncio + async def test_run_guarded_orphans_the_coroutine_on_cancel_by_default(self): + """Documents the default: asyncio.wait does not cancel what it waits on.""" + context = await self._started_context() + cancelled = asyncio.Event() + + async def coro(): + try: + await asyncio.sleep(30) + except asyncio.CancelledError: + cancelled.set() + raise + + task = asyncio.create_task(context._run_guarded(coro())) + await asyncio.sleep(0) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + await asyncio.sleep(0) + assert not cancelled.is_set() + await context.close() + + @pytest.mark.asyncio + async def test_run_guarded_propagates_cancel_when_asked(self): + """With propagate_cancel the coroutine is told, so it can clean up. + + A resolver that polls a remote task has to send a cancellation of its own; + it can only do that if the cancellation actually reaches it. + """ + context = await self._started_context() + cancelled = asyncio.Event() + + async def coro(): + try: + await asyncio.sleep(30) + except asyncio.CancelledError: + cancelled.set() + raise + + task = asyncio.create_task( + context._run_guarded(coro(), propagate_cancel=True) + ) + await asyncio.sleep(0) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert cancelled.is_set() + await context.close() + @pytest.mark.asyncio async def test_run_guarded_raises_when_task_not_started(self): """If start() was never called, _run_guarded refuses to run the coro."""