feat(sdk): add mcp extension

Signed-off-by: Tomas Pilar <thomas7pilar@gmail.com>

Author: pilartomas

apps/agentstack-sdk-py/examples/mcp_agent.py

@@ -5,13 +5,14 @@
 from typing import Annotated
 
 from a2a.types import Message
-from mcp import ClientSession
 
 from agentstack_sdk.a2a.extensions.auth.oauth import OAuthExtensionServer, OAuthExtensionSpec
+from agentstack_sdk.a2a.extensions.mcp import MCPExtensionParams, MCPExtensionServer, MCPExtensionSpec
 from agentstack_sdk.a2a.extensions.services.mcp import MCPServiceExtensionServer, MCPServiceExtensionSpec
 from agentstack_sdk.a2a.types import RunYield
 from agentstack_sdk.server import Server
 from agentstack_sdk.server.context import RunContext
+from agentstack_sdk.server.mcp.session import MCPClientSession
 
 server = Server()
 
@@ -21,24 +22,34 @@ async def mcp_agent(
     message: Message,
     context: RunContext,
     oauth: Annotated[OAuthExtensionServer, OAuthExtensionSpec.single_demand()],
-    mcp: Annotated[
+    mcp: Annotated[MCPExtensionServer, MCPExtensionSpec(params=MCPExtensionParams())],
+    mcp_service: Annotated[
         MCPServiceExtensionServer,
         MCPServiceExtensionSpec.single_demand(),
     ],
 ) -> AsyncGenerator[RunYield, Message]:
     """Lists tools"""
 
-    if not mcp:
+    if not mcp_service:
         yield "MCP extension hasn't been activated, no tools are available"
         return
 
-    async with mcp.create_client() as (read, write), ClientSession(read, write) as session:
+    async with (
+        mcp_service.create_client() as (read, write),
+        MCPClientSession(read, write, context=context).apply(mcp) as session,
+    ):
         await session.initialize()
 
-        tools = await session.list_tools()
+        result = await session.list_tools()
 
         yield "Available tools: \n"
-        yield "\n".join([t.name for t in tools.tools])
+        yield "\n".join([t.name for t in result.tools])
+
+        if result.tools:
+            tool = result.tools[0]
+            yield f"Calling tool {tool.name}"
+            await session.call_tool(tool.name, None)
+            yield "Tool call finished"
 
 
 if __name__ == "__main__":

apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/mcp/__init__.py

@@ -0,0 +1,4 @@
+# Copyright 2025 © BeeAI a Series of LF Projects, LLC
+# SPDX-License-Identifier: Apache-2.0
+
+from .mcp import *

apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/mcp/mcp.py

@@ -0,0 +1,76 @@
+# Copyright 2025 © BeeAI a Series of LF Projects, LLC
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import uuid
+from types import NoneType
+from typing import TYPE_CHECKING, Any, Literal
+
+import a2a.types
+import pydantic
+
+from agentstack_sdk.a2a.extensions.base import BaseExtensionClient, BaseExtensionServer, BaseExtensionSpec
+from agentstack_sdk.a2a.types import AgentMessage
+
+if TYPE_CHECKING:
+    from agentstack_sdk.server.context import RunContext
+
+
+class ToolCallApprovalRequest(pydantic.BaseModel):
+    name: str
+    arguments: dict[str, Any] | None = None
+
+
+class ToolCallApprovalResponse(pydantic.BaseModel):
+    action: Literal["accept", "reject"]
+
+
+class MCPExtensionParams(pydantic.BaseModel):
+    pass
+
+
+class MCPExtensionSpec(BaseExtensionSpec[MCPExtensionParams]):
+    URI: str = "https://a2a-extensions.agentstack.beeai.dev/mcp/v1"
+
+
+class MCPExtensionMetadata(pydantic.BaseModel):
+    pass
+
+
+class MCPExtensionServer(BaseExtensionServer[MCPExtensionSpec, MCPExtensionMetadata]):
+    def create_message(self, *, request: ToolCallApprovalRequest):
+        return AgentMessage(
+            text="Tool call approval requested", metadata={self.spec.URI: request.model_dump(mode="json")}
+        )
+
+    def parse_message(self, *, message: a2a.types.Message):
+        if not message or not message.metadata or not (data := message.metadata.get(self.spec.URI)):
+            raise RuntimeError("Invalid mcp response")
+        return ToolCallApprovalResponse.model_validate(data)
+
+    async def raise_tool_approval(self, request: ToolCallApprovalRequest, context: RunContext):
+        message = self.create_message(request=request)
+        message = await context.yield_async(message)
+        if message:
+            result = self.parse_message(message=message)
+            if result.action != "accept":
+                raise RuntimeError("User has rejected the tool call")
+        else:
+            raise RuntimeError("Tool call approval response is missing")
+
+
+class MCPExtensionClient(BaseExtensionClient[MCPExtensionSpec, NoneType]):
+    def create_message(self, *, response: ToolCallApprovalResponse, task_id: str | None):
+        return a2a.types.Message(
+            message_id=str(uuid.uuid4()),
+            role=a2a.types.Role.user,
+            parts=[],
+            task_id=task_id,
+            metadata={self.spec.URI: response.model_dump(mode="json")},
+        )
+
+    def parse_message(self, *, message: a2a.types.Message):
+        if not message or not message.metadata or not (data := message.metadata.get(self.spec.URI)):
+            raise ValueError("Invalid mcp request")
+        return ToolCallApprovalRequest.model_validate(data)

apps/agentstack-sdk-py/src/agentstack_sdk/server/mcp/__init__.py

@@ -0,0 +1,2 @@
+# Copyright 2025 © BeeAI a Series of LF Projects, LLC
+# SPDX-License-Identifier: Apache-2.0

apps/agentstack-sdk-py/src/agentstack_sdk/server/mcp/session.py

@@ -0,0 +1,67 @@
+# Copyright 2025 © BeeAI a Series of LF Projects, LLC
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from datetime import timedelta
+from typing import Any, Self
+
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+from mcp import ClientSession
+from mcp.client.session import ElicitationFnT, ListRootsFnT, LoggingFnT, MessageHandlerFnT, SamplingFnT
+from mcp.shared.message import SessionMessage
+from mcp.shared.session import ProgressFnT
+from mcp.types import CallToolResult, Implementation
+
+from agentstack_sdk.a2a.extensions.mcp import MCPExtensionServer, ToolCallApprovalRequest
+from agentstack_sdk.server.context import RunContext
+
+
+class MCPClientSession(ClientSession):
+    def __init__(
+        self,
+        read_stream: MemoryObjectReceiveStream[SessionMessage | Exception],
+        write_stream: MemoryObjectSendStream[SessionMessage],
+        read_timeout_seconds: timedelta | None = None,
+        sampling_callback: SamplingFnT | None = None,
+        elicitation_callback: ElicitationFnT | None = None,
+        list_roots_callback: ListRootsFnT | None = None,
+        logging_callback: LoggingFnT | None = None,
+        message_handler: MessageHandlerFnT | None = None,
+        client_info: Implementation | None = None,
+        *,
+        context: RunContext,
+    ) -> None:
+        super().__init__(
+            read_stream,
+            write_stream,
+            read_timeout_seconds,
+            sampling_callback,
+            elicitation_callback,
+            list_roots_callback,
+            logging_callback,
+            message_handler,
+            client_info,
+        )
+        self._context = context
+        self._mcp_extension = None
+
+    def apply(self, extension: MCPExtensionServer) -> Self:
+        self._mcp_extension = extension
+        return self
+
+    async def call_tool(
+        self,
+        name: str,
+        arguments: dict[str, Any] | None = None,
+        read_timeout_seconds: timedelta | None = None,
+        progress_callback: ProgressFnT | None = None,
+    ) -> CallToolResult:
+        """Send a tools/call request with optional progress callback support."""
+
+        if self._mcp_extension:
+            await self._mcp_extension.raise_tool_approval(
+                request=ToolCallApprovalRequest(name=name, arguments=arguments), context=self._context
+            )
+
+        return await super().call_tool(name, arguments, read_timeout_seconds, progress_callback)