feat: Add mcp transport protocol

packages/toolbox-core/src/toolbox_core/client.py

@@ -20,7 +20,12 @@
 from deprecated import deprecated
 
 from .itransport import ITransport
-from .protocol import ToolSchema
+from .mcp_transport import (
+    McpHttpTransportV20241105,
+    McpHttpTransportV20250326,
+    McpHttpTransportV20250618,
+)
+from .protocol import Protocol, ToolSchema
 from .tool import ToolboxTool
 from .toolbox_transport import ToolboxTransport
 from .utils import identify_auth_requirements, resolve_value
@@ -44,6 +49,7 @@ def __init__(
         client_headers: Optional[
             Mapping[str, Union[Callable[[], str], Callable[[], Awaitable[str]], str]]
         ] = None,
+        protocol: Protocol = Protocol.TOOLBOX,
     ):
         """
         Initializes the ToolboxClient.
@@ -54,8 +60,21 @@ def __init__(
                 If None (default), a new session is created internally. Note that
                 if a session is provided, its lifecycle (including closing)
                 should typically be managed externally.
-            client_headers: Headers to include in each request sent through this client.
+            client_headers: Headers to include in each request sent through this
+            client.
+            protocol: The communication protocol to use.
         """
+        if protocol == Protocol.TOOLBOX:
+            self.__transport = ToolboxTransport(url, session)
+        elif protocol in Protocol.get_supported_mcp_versions():
+            if protocol == Protocol.MCP_v20250618:
+                self.__transport = McpHttpTransportV20250618(url, session, protocol)
+            elif protocol == Protocol.MCP_v20250326:
+                self.__transport = McpHttpTransportV20250326(url, session, protocol)
+            elif protocol == Protocol.MCP_v20241105:
+                self.__transport = McpHttpTransportV20241105(url, session, protocol)
+        else:
+            raise ValueError(f"Unsupported MCP protocol version: {protocol}")
 
         self.__transport = ToolboxTransport(url, session)
         self.__client_headers = client_headers if client_headers is not None else {}

packages/toolbox-core/src/toolbox_core/mcp_transport/__init__.py

@@ -0,0 +1,23 @@
+# Copyright 2025 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 .v20241105.mcp import McpHttpTransportV20241105
+from .v20250326.mcp import McpHttpTransportV20250326
+from .v20250618.mcp import McpHttpTransportV20250618
+
+__all__ = [
+    "McpHttpTransportV20241105",
+    "McpHttpTransportV20250326",
+    "McpHttpTransportV20250618",
+]

packages/toolbox-core/src/toolbox_core/mcp_transport/transport_base.py

@@ -0,0 +1,102 @@
+# Copyright 2025 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.
+
+import asyncio
+from abc import ABC, abstractmethod
+from typing import Optional
+
+from aiohttp import ClientSession
+
+from ..itransport import ITransport
+from ..protocol import (
+    AdditionalPropertiesSchema,
+    ParameterSchema,
+    Protocol,
+    ToolSchema,
+)
+
+
+class _McpHttpTransportBase(ITransport, ABC):
+    """Base transport for MCP protocols."""
+
+    def __init__(
+        self,
+        base_url: str,
+        session: Optional[ClientSession] = None,
+        protocol: Protocol = Protocol.MCP,
+    ):
+        self._mcp_base_url = f"{base_url}/mcp/"
+        self._protocol_version = protocol.value
+        self._server_version: Optional[str] = None
+
+        self._manage_session = session is None
+        self._session = session or ClientSession()
+        self._init_lock = asyncio.Lock()
+        self._init_task: Optional[asyncio.Task] = None
+
+    async def _ensure_initialized(self):
+        """Ensures the session is initialized before making requests."""
+        async with self._init_lock:
+            if self._init_task is None:
+                self._init_task = asyncio.create_task(self._initialize_session())
+        await self._init_task
+
+    @property
+    def base_url(self) -> str:
+        return self._mcp_base_url
+
+    def _convert_tool_schema(self, tool_data: dict) -> ToolSchema:
+        """Converts a raw MCP tool dictionary into the Toolbox ToolSchema."""
+        parameters = []
+        input_schema = tool_data.get("inputSchema", {})
+        properties = input_schema.get("properties", {})
+        required = input_schema.get("required", [])
+
+        for name, schema in properties.items():
+            additional_props = schema.get("additionalProperties")
+            if isinstance(additional_props, dict):
+                additional_props = AdditionalPropertiesSchema(
+                    type=additional_props["type"]
+                )
+            else:
+                additional_props = True
+            parameters.append(
+                ParameterSchema(
+                    name=name,
+                    type=schema["type"],
+                    description=schema.get("description", ""),
+                    required=name in required,
+                    additionalProperties=additional_props,
+                )
+            )
+
+        return ToolSchema(
+            description=tool_data.get("description") or "", parameters=parameters
+        )
+
+    async def close(self):
+        async with self._init_lock:
+            if self._init_task:
+                try:
+                    await self._init_task
+                except Exception:
+                    # If initialization failed, we can still try to close.
+                    pass
+        if self._manage_session and self._session and not self._session.closed:
+            await self._session.close()
+
+    @abstractmethod
+    async def _initialize_session(self):
+        """Initializes the MCP session."""
+        pass

packages/toolbox-core/src/toolbox_core/mcp_transport/v20241105/mcp.py

@@ -0,0 +1,168 @@
+# Copyright 2025 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 typing import Mapping, Optional, TypeVar
+
+from pydantic import BaseModel
+
+from ... import version
+from ...protocol import ManifestSchema
+from ..transport_base import _McpHttpTransportBase
+from . import types
+
+ReceiveResultT = TypeVar("ReceiveResultT", bound=BaseModel)
+
+
+class McpHttpTransportV20241105(_McpHttpTransportBase):
+    """Transport for the MCP v2024-11-05 protocol."""
+
+    async def _send_request(
+        self,
+        url: str,
+        request: types.MCPRequest[ReceiveResultT] | types.MCPNotification,
+        headers: Optional[Mapping[str, str]] = None,
+    ) -> ReceiveResultT | None:
+        """Sends a JSON-RPC request to the MCP server."""
+        params = (
+            request.params.model_dump(mode="json", exclude_none=True)
+            if isinstance(request.params, BaseModel)
+            else request.params
+        )
+        rpc_msg: BaseModel
+        if isinstance(request, types.MCPNotification):
+            rpc_msg = types.JSONRPCNotification(method=request.method, params=params)
+        else:
+            rpc_msg = types.JSONRPCRequest(method=request.method, params=params)
+
+        payload = rpc_msg.model_dump(mode="json", exclude_none=True)
+
+        async with self._session.post(
+            url, json=payload, headers=dict(headers or {})
+        ) as response:
+            if not response.ok:
+                error_text = await response.text()
+                raise RuntimeError(
+                    f"API request failed with status {response.status} "
+                    f"({response.reason}). Server response: {error_text}"
+                )
+
+            if response.status == 204 or response.content.at_eof():
+                return None
+
+            json_resp = await response.json()
+
+            # Check for JSON-RPC Error
+            if "error" in json_resp:
+                try:
+                    err = types.JSONRPCError.model_validate(json_resp).error
+                    raise RuntimeError(
+                        f"MCP request failed with code {err.code}: {err.message}"
+                    )
+                except Exception:
+                    raise RuntimeError(f"MCP request failed: {json_resp.get('error')}")
+
+            # Parse Result
+            if isinstance(request, types.MCPRequest):
+                try:
+                    rpc_resp = types.JSONRPCResponse.model_validate(json_resp)
+                    return request.get_result_model().model_validate(rpc_resp.result)
+                except Exception as e:
+                    raise RuntimeError(f"Failed to parse JSON-RPC response: {e}")
+            return None
+
+    async def _initialize_session(self):
+        """Initializes the MCP session."""
+        params = types.InitializeRequestParams(
+            protocolVersion=self._protocol_version,
+            capabilities=types.ClientCapabilities(),
+            clientInfo=types.Implementation(
+                name="toolbox-python-sdk", version=version.__version__
+            ),
+        )
+
+        result = await self._send_request(
+            url=self._mcp_base_url, request=types.InitializeRequest(params=params)
+        )
+
+        self._server_version = result.serverInfo.version
+        if result.protocolVersion != self._protocol_version:
+            raise RuntimeError(
+                f"MCP version mismatch: client does not support server version {result.protocolVersion}"
+            )
+        if not result.capabilities.tools:
+            if self._manage_session:
+                await self.close()
+            raise RuntimeError("Server does not support the 'tools' capability.")
+
+        await self._send_request(
+            url=self._mcp_base_url, request=types.InitializedNotification()
+        )
+
+    async def tools_list(
+        self,
+        toolset_name: Optional[str] = None,
+        headers: Optional[Mapping[str, str]] = None,
+    ) -> ManifestSchema:
+        """Lists available tools from the server using the MCP protocol."""
+        await self._ensure_initialized()
+
+        url = self._mcp_base_url + (toolset_name if toolset_name else "")
+        result = await self._send_request(
+            url=url, request=types.ListToolsRequest(), headers=headers
+        )
+        if result is None:
+            raise RuntimeError("Failed to list tools: No response from server.")
+
+        tools_map = {
+            t.name: self._convert_tool_schema(t.model_dump(mode="json", by_alias=True))
+            for t in result.tools
+        }
+        if self._server_version is None:
+            raise RuntimeError("Server version not available.")
+
+        return ManifestSchema(serverVersion=self._server_version, tools=tools_map)
+
+    async def tool_get(
+        self, tool_name: str, headers: Optional[Mapping[str, str]] = None
+    ) -> ManifestSchema:
+        """Gets a single tool from the server by listing all and filtering."""
+        manifest = await self.tools_list(headers=headers)
+
+        if tool_name not in manifest.tools:
+            raise ValueError(f"Tool '{tool_name}' not found.")
+
+        return ManifestSchema(
+            serverVersion=manifest.serverVersion,
+            tools={tool_name: manifest.tools[tool_name]},
+        )
+
+    async def tool_invoke(
+        self, tool_name: str, arguments: dict, headers: Optional[Mapping[str, str]]
+    ) -> str:
+        """Invokes a specific tool on the server using the MCP protocol."""
+        await self._ensure_initialized()
+
+        result = await self._send_request(
+            url=self._mcp_base_url,
+            request=types.CallToolRequest(
+                params=types.CallToolRequestParams(name=tool_name, arguments=arguments)
+            ),
+            headers=headers,
+        )
+        if result is None:
+            raise RuntimeError(
+                f"Failed to invoke tool '{tool_name}': No response from server."
+            )
+
+        return "".join(c.text for c in result.content if c.type == "text") or "null"