fix: harden MCPDriver based on human review + MCP 1.26 SDK pass

Blocker fixes:
- tests/test_mcp_driver.py: strengthen integration test assertion for
  real add(3,4) result; assert structured_content.result==7 and
  content[0].text=='7' rather than bare 'is not None'
- mcp.py/_run_with_retry: handle McpError (protocol-level rejection)
  immediately as DriverError; McpError is not retryable — the server
  processed and rejected the request

Major fixes:
- mcp.py/discover(): paginate tools/list via _fetch_all_tools; loop on
  nextCursor until exhausted to avoid silent capability truncation on
  large MCP servers
- mcp_support.py/ToolSpec + discover(): forward ToolAnnotations hints
  (readOnlyHint, destructiveHint, idempotentHint) to ToolSpec; derive
  SafetyClass from them via _infer_safety_class(); safety_class_map still
  overrides. Eliminates always-READ misclassification of destructive tools
- mcp_support.py/call_tool(): forward read_timeout_seconds from constraints
  to ClientSession.call_tool(); prevents indefinite hangs over HTTP
- mcp.py/_run_with_retry: document at-least-once delivery semantics for
  HTTP transport; advise callers to set max_retries=0 for WRITE/DESTRUCTIVE
- pyproject.toml: tighten mcp lower bound to >=1.6 in both [mcp] extra
  and [dev] extras; ToolAnnotations, outputSchema, and nextCursor require
  this release line

Minor fixes:
- mcp_support.py/ToolSpec: forward Tool.outputSchema for downstream use
  by firewall budget/redaction rules
- docs/integrations.md: expand HTTP section with full async def main()
  discover() + register + route example; add at-least-once warning inline

Author: dgenio

docs/integrations.md

@@ -47,13 +47,40 @@ asyncio.run(main())
 ### Streamable HTTP transport
 
 ```python
+import asyncio
+
+from agent_kernel import CapabilityRegistry, Kernel, StaticRouter
 from agent_kernel.drivers.mcp import MCPDriver
 
-http_driver = MCPDriver.from_http(
-    url="https://example.com/mcp",
-    server_name="remote-tools",
-    max_retries=1,
-)
+
+async def main() -> None:
+    registry = CapabilityRegistry()
+    router = StaticRouter(fallback=[])
+    kernel = Kernel(registry=registry, router=router)
+
+    # Connect to a remote Streamable HTTP MCP server.
+    # Note: max_retries > 0 creates at-least-once delivery semantics for
+    # tools/call — if a connection drops after the server processes the
+    # request but before the response arrives, the call will be repeated.
+    # Ensure target tools are idempotent, or set max_retries=0 for
+    # WRITE/DESTRUCTIVE capabilities.
+    driver = MCPDriver.from_http(
+        url="https://example.com/mcp",
+        server_name="remote-tools",
+        max_retries=1,
+    )
+    kernel.register_driver(driver)
+
+    # Discover tools and register them as capabilities.
+    capabilities = await driver.discover(namespace="remote")
+    registry.register_many(capabilities)
+
+    # Route each discovered capability to this MCP driver.
+    for capability in capabilities:
+        router.add_route(capability.capability_id, [driver.driver_id])
+
+
+asyncio.run(main())
 ```
 
 ### Notes

pyproject.toml

@@ -38,9 +38,9 @@ dev = [
     "ruff>=0.4",
     "mypy>=1.10",
     "httpx>=0.27",
-    "mcp>=1.0",
+    "mcp>=1.6",
 ]
-mcp = ["mcp>=1.0"]
+mcp = ["mcp>=1.6"]
 otel = ["opentelemetry-api>=1.20"]
 
 [tool.hatch.build.targets.wheel]

src/agent_kernel/drivers/mcp.py

@@ -11,13 +11,35 @@
 from .base import ExecutionContext
 from .mcp_support import (
     SessionFactory,
+    ToolSpec,
     build_http_session_factory,
     build_stdio_session_factory,
     call_tool,
     extract_tool_specs,
     normalize_call_result,
 )
 
+# Lazy import of McpError — only available when the mcp optional dep is installed.
+# If mcp is absent, factory methods raise ImportError before any session is created,
+# so _McpError will never be None on a live driver instance.
+try:
+    from mcp.shared.exceptions import McpError as _McpError
+except ImportError:  # pragma: no cover
+    _McpError = None  # type: ignore[assignment,misc]
+
+
+def _infer_safety_class(spec: ToolSpec) -> SafetyClass:
+    """Infer a SafetyClass from MCP ToolAnnotations hints.
+
+    Uses a conservative default of READ when annotations are absent.
+    The caller's safety_class_map takes precedence over the inferred value.
+    """
+    if spec.destructive_hint:
+        return SafetyClass.DESTRUCTIVE
+    if spec.read_only_hint:
+        return SafetyClass.READ
+    return SafetyClass.READ
+
 
 class MCPDriver:
     """A driver that invokes capabilities via MCP tools/call."""
@@ -92,19 +114,20 @@ async def discover(
         namespace: str | None = None,
         safety_class_map: dict[str, SafetyClass] | None = None,
     ) -> list[Capability]:
-        """Discover MCP tools and convert them to capabilities."""
-        tool_list = await self._run_with_retry(
+        """Discover MCP tools across all pages and convert them to capabilities."""
+        tools = await self._run_with_retry(
             operation_name="tools/list",
-            action=lambda session: session.list_tools(),
+            action=self._fetch_all_tools,
         )
 
         capabilities: list[Capability] = []
-        for spec in extract_tool_specs(tool_list):
+        for spec in extract_tool_specs(tools):
             capability_id = f"{namespace}.{spec.name}" if namespace else spec.name
+            inferred = _infer_safety_class(spec)
             safety_class = (
-                safety_class_map.get(spec.name, SafetyClass.READ)
+                safety_class_map.get(spec.name, inferred)
                 if safety_class_map is not None
-                else SafetyClass.READ
+                else inferred
             )
             capabilities.append(
                 Capability(
@@ -121,21 +144,42 @@ async def discover(
             )
         return capabilities
 
+    async def _fetch_all_tools(self, session: Any) -> list[Any]:
+        """Paginate tools/list to exhaustion and return a flat list of Tool objects."""
+        all_tools: list[Any] = []
+        cursor: str | None = None
+        while True:
+            result = await session.list_tools(cursor=cursor)
+            all_tools.extend(getattr(result, "tools", []) or [])
+            cursor = getattr(result, "nextCursor", None)
+            if not cursor:
+                break
+        return all_tools
+
     async def execute(self, ctx: ExecutionContext) -> RawResult:
         """Execute an MCP tool call for the given capability context."""
         operation = str(ctx.args.get("operation", ctx.capability_id))
         params = {k: v for k, v in ctx.args.items() if k != "operation"}
 
         # Apply policy constraints as default arguments, without overriding explicit args.
+        # read_timeout_seconds is an SDK control parameter — applied to the session call
+        # directly rather than forwarded to the tool as an argument.
+        read_timeout_seconds_raw = ctx.constraints.get("read_timeout_seconds")
         for key, value in ctx.constraints.items():
-            params.setdefault(key, value)
+            if key != "read_timeout_seconds":
+                params.setdefault(key, value)
+
+        read_timeout_seconds: float | None = (
+            float(read_timeout_seconds_raw) if read_timeout_seconds_raw is not None else None
+        )
 
         result = await self._run_with_retry(
             operation_name=f"tools/call:{operation}",
             action=lambda session: call_tool(
                 session,
                 operation=operation,
                 params=params,
+                read_timeout_seconds=read_timeout_seconds,
             ),
         )
 
@@ -170,11 +214,19 @@ async def _run_with_retry(
             except DriverError:
                 raise
             except Exception as exc:
-                # Broad catch is intentional: exceptions at this level are
-                # session/transport failures (connection refused, EOF, timeout).
-                # MCP tool-level application errors are returned as isError=True
-                # responses and converted to DriverError before reaching this
-                # handler — they never appear as Python exceptions here.
+                # McpError is a protocol-level rejection (tool not found, auth
+                # failure, invalid params) — the server processed and rejected the
+                # request. It is not retryable; surface it immediately as DriverError.
+                if _McpError is not None and isinstance(exc, _McpError):
+                    raise DriverError(
+                        f"MCPDriver '{self._driver_id}' received a protocol error "
+                        f"during {operation_name}: {exc}"
+                    ) from exc
+                # All other exceptions are session/transport failures (connection
+                # refused, EOF, timeout) and are retryable for HTTP transport.
+                # Note: HTTP retries create at-least-once delivery semantics for
+                # tools/call. Callers using WRITE/DESTRUCTIVE capabilities over HTTP
+                # should ensure the target tool is idempotent, or set max_retries=0.
                 last_exc = exc
 
         reason = str(last_exc) if last_exc is not None else "unknown transport failure"

src/agent_kernel/drivers/mcp_support.py

@@ -6,6 +6,7 @@
 from collections.abc import AsyncIterator, Callable
 from contextlib import AbstractAsyncContextManager, asynccontextmanager
 from dataclasses import dataclass
+from datetime import timedelta
 from typing import Any
 
 from ..errors import DriverError
@@ -19,27 +20,42 @@ class ToolSpec:
 
     name: str
     description: str
-
-
-async def call_tool(session: Any, *, operation: str, params: dict[str, Any]) -> Any:
+    read_only_hint: bool = False
+    destructive_hint: bool = False
+    idempotent_hint: bool = False
+    output_schema: dict[str, Any] | None = None
+
+
+async def call_tool(
+    session: Any,
+    *,
+    operation: str,
+    params: dict[str, Any],
+    read_timeout_seconds: float | None = None,
+) -> Any:
     """Call an MCP tool via tools/call."""
-    return await session.call_tool(operation, arguments=params)
+    timeout = timedelta(seconds=read_timeout_seconds) if read_timeout_seconds is not None else None
+    return await session.call_tool(operation, arguments=params, read_timeout_seconds=timeout)
 
 
-def extract_tool_specs(tool_list_response: Any) -> list[ToolSpec]:
-    """Extract tool metadata from a tools/list response payload."""
-    tools = getattr(tool_list_response, "tools", [])
+def extract_tool_specs(tools: list[Any]) -> list[ToolSpec]:
+    """Extract tool metadata from a flat list of MCP Tool objects."""
     if not isinstance(tools, list):
         return []
     specs: list[ToolSpec] = []
     for tool in tools:
         name = getattr(tool, "name", None)
         if not isinstance(name, str) or not name:
             continue
+        ann = getattr(tool, "annotations", None)
         specs.append(
             ToolSpec(
                 name=name,
                 description=str(getattr(tool, "description", "") or ""),
+                read_only_hint=bool(getattr(ann, "readOnlyHint", False)),
+                destructive_hint=bool(getattr(ann, "destructiveHint", False)),
+                idempotent_hint=bool(getattr(ann, "idempotentHint", False)),
+                output_schema=getattr(tool, "outputSchema", None),
             )
         )
     return specs

tests/test_mcp_driver.py

@@ -37,13 +37,14 @@ def __init__(
         self._call_error = call_error
         self.calls: list[tuple[str, dict[str, Any]]] = []
 
-    async def list_tools(self) -> ListToolsResult:
+    async def list_tools(self, cursor: str | None = None) -> ListToolsResult:
         return ListToolsResult(tools=self._tools)
 
     async def call_tool(
         self,
         operation: str,
         arguments: dict[str, Any],
+        read_timeout_seconds: Any = None,
     ) -> CallToolResult:
         self.calls.append((operation, arguments))
         if self._call_error is not None:
@@ -292,4 +293,6 @@ async def in_memory_factory() -> AsyncIterator[ClientSession]:
         args={"operation": "add", "a": 3, "b": 4},
     )
     result = await driver.execute(ctx)
-    assert result.data is not None
+    assert isinstance(result.data, dict)
+    assert result.data["structured_content"]["result"] == 7
+    assert result.data["content"][0]["text"] == "7"