From c5f06c345725fd1fb54a4d28c93f45a8b1737fe8 Mon Sep 17 00:00:00 2001 From: zeishr <67806187+theo-abel@users.noreply.github.com> Date: Fri, 21 Aug 2026 12:44:01 +0200 Subject: [PATCH 1/2] feat(mcp): pass the mcp 2.x extension seam through McpToolset An MCP extension needs three things from the client session: its identifier advertised in per-request capabilities, a claim that folds its non-core result shape into tools/call parsing, and bindings for its notifications. MCP SDK 2.x exposes all three on ClientSession. McpToolset owns that session and did not expose them, so an extension could not be carried without forking the toolset. They now follow the same path sampling and elicitation already take, down to SessionContext and both ClientSession construction sites, with the three names resolved in dependencies/_mcp like every other MCP import. Passing them through is necessary but not sufficient, for two reasons that are easy to miss. ClientSession parses a claimed result and stops there. It raises UnexpectedClaimedResult unless the call opted in with allow_claimed, and it never calls claim.resolve -- the high-level mcp Client does that. ADK uses the low-level session, so McpTool now opts in when a claim is registered, looks the claim up by the model its result parsed into, resolves it, and holds the outcome to the same output schema a direct result gets. The agent sees an ordinary tool result either way. Claims are also inert unless the connection negotiated 2026-07-28, and initialize() always performs the pre-2026 handshake, which is what SessionContext called. An extension registered on such a session could never fire. Bring-up now prefers server/discover, but only when an extension was actually configured: the default path keeps calling initialize() with no new wire traffic and no new failure modes. The probe is bounded to half the bring-up budget so that a server which never answers it still leaves room for the initialize() fallback, and a server too old for server/discover logs why extensions are inactive rather than failing the session. Since 856acf21 the pin admits both majors, and 1.x has no extension parameters on ClientSession at all -- not a rename, an absence. An opt-in there is refused at construction rather than accepted and ignored. A silent no-op would surface as a tool call that never returns, on the install least able to explain why, so the message names the SDK and the fix. Nothing else changes on 1.x: with no extension configured the three arguments are not passed to ClientSession, which is the call 1.x has always received. One fix comes along because the tasks extension needs it. _run_guarded races a coroutine against the session task with asyncio.wait, which does not cancel what it waits on, so a cancelled caller leaves the coroutine running detached. That is survivable for a single request the transport will eventually fail, but not for a resolver that owes the server a cancellation: it would never learn it should send one. _run_guarded takes an opt-in propagate_cancel for that, and the existing call path keeps its semantics. Bug: #6826 --- src/google/adk/dependencies/_mcp.py | 30 ++- .../adk/tools/mcp_tool/mcp_session_manager.py | 71 ++++++ src/google/adk/tools/mcp_tool/mcp_tool.py | 72 ++++++ src/google/adk/tools/mcp_tool/mcp_toolset.py | 26 +++ .../adk/tools/mcp_tool/session_context.py | 128 ++++++++++- tests/unittests/tools/mcp_tool/_sdk_compat.py | 14 ++ .../mcp_tool/test_mcp_session_manager.py | 65 ++++++ .../unittests/tools/mcp_tool/test_mcp_tool.py | 178 +++++++++++++++ .../tools/mcp_tool/test_mcp_toolset.py | 56 +++++ .../tools/mcp_tool/test_session_context.py | 209 ++++++++++++++++++ 10 files changed, 838 insertions(+), 11 deletions(-) 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/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..dc5dc9a5d7e 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 @@ -41,8 +43,11 @@ from ...auth.auth_schemes import AuthScheme from ...auth.auth_tool import AuthConfig from ...dependencies._mcp import ElicitationFnT +from ...dependencies._mcp import IS_MCP_SDK_V2 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 @@ -170,6 +175,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, credential_key: str | None = None, ): """Initializes the McpToolset. @@ -221,6 +229,18 @@ 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. credential_key: A user specified key used to load and save this credential in a credential service. Used with auth_scheme. """ @@ -230,6 +250,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 if not connection_params: raise ValueError("Missing connection params in McpToolset.") @@ -260,6 +283,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 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/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_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..fd64b1233da 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,9 @@ from mcp.types import TextResourceContents import pytest +from ._sdk_compat import requires_sdk_v1 +from ._sdk_compat import requires_sdk_v2 + class MockMCPTool: """Mock MCP Tool for testing.""" @@ -929,6 +933,58 @@ 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": {}}, + ) + + 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.""" From ee31781488c6fae9eba8f658edbcb571aef9917f Mon Sep 17 00:00:00 2001 From: zeishr <67806187+theo-abel@users.noreply.github.com> Date: Fri, 21 Aug 2026 13:02:06 +0200 Subject: [PATCH 2/2] feat(mcp): support the MCP Tasks extension (SEP-2663) in McpToolset A tool call normally answers on the connection that made it, which stops working once the work takes minutes: an intermediary caps how long a request may stay open, and a dropped connection loses the operation with no handle to resume from. Servers work around it by splitting one operation into start, status and result tools and letting the model drive the polling loop -- a model turn per poll, and only if the prompt says so. The Tasks extension is the protocol's answer. The server decides per request that the work is long-running and replies with a durable task handle; the client polls tasks/get until a terminal state and reads the result from there. With enable_tasks=True the toolset does that itself, so the agent declares one tool, calls it once, and gets an ordinary tool result. Off by default, and a server without the extension is unaffected. The wire models are ours because there is nothing to depend on: ext-tasks publishes TypeScript only, the SDK's client-side implementation is still an open issue, and the task types the SDK does carry are the older in-core design that SEP-2663 replaced -- nested task object, ttl rather than ttlMs, a different notification method. They stay private in _tasks.py so an upstream implementation can replace them without touching ADK's surface. No raw JSON-RPC is involved: send_request takes an arbitrary request, and the Mcp-Name header the specification requires on tasks/* comes from the request's name_param. _tasks.py is MCP SDK 2.x only and cannot be imported on 1.x, so the toolset imports it after refusing enable_tasks there, not at module scope. That refusal is the same one the extension arguments get, for the same reason: an opt-in that quietly did nothing on a call this feature exists to make long-running is indistinguishable from a hang. The resolver reports rather than raises. A task that fails or that the server cancels comes back as a result with isError set, because that is what a tool failing inline produces and going through a task should not change it. A task that reaches input_required is cancelled and reported as unsupported, rather than left for the server to hold until its ttl expires; bridging tasks/update to an interaction channel is a follow-up, and upstream has not settled that design either. Poll intervals follow the server, clamped to [100ms, 30s] so that a zero cannot spin and an absurd value cannot hide a finished task. Verified end to end against the sample server: with enable_tasks=True the call is served as a task and polled to completion over a dozen-odd requests -- the count follows the server's stated interval -- and the result dict is byte-identical to the one the same tool returns with the option off. Closes: #6826 --- .../samples/mcp/mcp_tasks_agent/__init__.py | 15 + .../samples/mcp/mcp_tasks_agent/agent.py | 70 ++++ .../mcp/mcp_tasks_agent/task_server.py | 210 +++++++++++ docs/guides/README.md | 1 + docs/guides/tools/mcp_tool/tasks/index.md | 134 +++++++ src/google/adk/tools/mcp_tool/_tasks.py | 279 +++++++++++++++ src/google/adk/tools/mcp_tool/mcp_toolset.py | 76 +++- tests/unittests/test_samples.py | 10 + .../tools/mcp_tool/test_mcp_tasks.py | 326 ++++++++++++++++++ .../tools/mcp_tool/test_mcp_toolset.py | 72 ++++ 10 files changed, 1190 insertions(+), 3 deletions(-) create mode 100644 contributing/samples/mcp/mcp_tasks_agent/__init__.py create mode 100644 contributing/samples/mcp/mcp_tasks_agent/agent.py create mode 100644 contributing/samples/mcp/mcp_tasks_agent/task_server.py create mode 100644 docs/guides/tools/mcp_tool/tasks/index.md create mode 100644 src/google/adk/tools/mcp_tool/_tasks.py create mode 100644 tests/unittests/tools/mcp_tool/test_mcp_tasks.py 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/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_toolset.py b/src/google/adk/tools/mcp_tool/mcp_toolset.py index dc5dc9a5d7e..3b36e2eda3b 100644 --- a/src/google/adk/tools/mcp_tool/mcp_toolset.py +++ b/src/google/adk/tools/mcp_tool/mcp_toolset.py @@ -43,7 +43,6 @@ from ...auth.auth_schemes import AuthScheme from ...auth.auth_tool import AuthConfig from ...dependencies._mcp import ElicitationFnT -from ...dependencies._mcp import IS_MCP_SDK_V2 from ...dependencies._mcp import ListResourcesResult from ...dependencies._mcp import ListToolsResult from ...dependencies._mcp import NotificationBinding @@ -60,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 @@ -117,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. @@ -178,6 +235,7 @@ def __init__( 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. @@ -241,6 +299,14 @@ def __init__( 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. """ @@ -250,8 +316,10 @@ 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._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: @@ -696,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): @@ -754,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/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/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_toolset.py b/tests/unittests/tools/mcp_tool/test_mcp_toolset.py index fd64b1233da..55a5f147ee2 100644 --- a/tests/unittests/tools/mcp_tool/test_mcp_toolset.py +++ b/tests/unittests/tools/mcp_tool/test_mcp_toolset.py @@ -60,6 +60,11 @@ 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.""" @@ -974,6 +979,73 @@ def test_extension_arguments_are_refused_on_sdk_v1(self): 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)