feat(sdk): add mcp extension

Signed-off-by: Tomas Pilar <[email protected]>

Author: pilartomas

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

@@ -5,13 +5,18 @@
 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.tool_call import (
+    ToolCallExtensionParams,
+    ToolCallExtensionServer,
+    ToolCallExtensionSpec,
+)
 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 +26,34 @@ async def mcp_agent(
     message: Message,
     context: RunContext,
     oauth: Annotated[OAuthExtensionServer, OAuthExtensionSpec.single_demand()],
-    mcp: Annotated[
+    mcp_tool_call: Annotated[ToolCallExtensionServer, ToolCallExtensionSpec(params=ToolCallExtensionParams())],
+    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_tool_call) 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/examples/mcp_client.py

@@ -12,6 +12,7 @@
 from pydantic import AnyHttpUrl, AnyUrl
 
 import agentstack_sdk.a2a.extensions
+from agentstack_sdk.a2a.extensions.mcp.tool_call import ToolCallResponse
 
 
 class OAuthHandler:
@@ -67,67 +68,83 @@ async def handler(request: web.Request) -> web.Response:
 async def run(base_url: str = "http://127.0.0.1:10000"):
     async with httpx.AsyncClient(timeout=30) as httpx_client:
         card = await a2a.client.A2ACardResolver(httpx_client, base_url=base_url).get_agent_card()
-        mcp_spec = agentstack_sdk.a2a.extensions.MCPServiceExtensionSpec.from_agent_card(card)
+        mcp_service_spec = agentstack_sdk.a2a.extensions.MCPServiceExtensionSpec.from_agent_card(card)
         oauth_spec = agentstack_sdk.a2a.extensions.OAuthExtensionSpec.from_agent_card(card)
+        tool_call_spec = agentstack_sdk.a2a.extensions.ToolCallExtensionSpec.from_agent_card(card)
 
-        if not mcp_spec:
+        if not mcp_service_spec:
             raise ValueError(f"Agent at {base_url} does not support MCP service injection")
         if not oauth_spec:
             raise ValueError(f"Agent at {base_url} does not support oAuth")
+        if not tool_call_spec:
+            raise ValueError(f"Agent at {base_url} does not support MCP")
 
-        mcp_extension_client = agentstack_sdk.a2a.extensions.MCPServiceExtensionClient(mcp_spec)
+        mcp_service_extension_client = agentstack_sdk.a2a.extensions.MCPServiceExtensionClient(mcp_service_spec)
         oauth_extension_client = agentstack_sdk.a2a.extensions.OAuthExtensionClient(oauth_spec)
+        tool_call_extension_client = agentstack_sdk.a2a.extensions.ToolCallExtensionClient(tool_call_spec)
 
         oauth = OAuthHandler()
         message = a2a.types.Message(
             message_id=str(uuid.uuid4()),
             role=a2a.types.Role.user,
             parts=[a2a.types.Part(root=a2a.types.TextPart(text="Howdy!"))],
-            metadata=mcp_extension_client.fulfillment_metadata(
+            metadata=mcp_service_extension_client.fulfillment_metadata(
                 mcp_fulfillments={
                     key: agentstack_sdk.a2a.extensions.services.mcp.MCPFulfillment(
                         transport=agentstack_sdk.a2a.extensions.services.mcp.StreamableHTTPTransport(
                             url=AnyHttpUrl("https://mcp.stripe.com")
                         ),
                     )
-                    for key in mcp_spec.params.mcp_demands
+                    for key in mcp_service_spec.params.mcp_demands
                 }
             )
             | oauth_extension_client.fulfillment_metadata(
                 oauth_fulfillments={
                     key: agentstack_sdk.a2a.extensions.OAuthFulfillment(redirect_uri=AnyUrl(oauth.redirect_uri))
                     for key in oauth_spec.params.oauth_demands
                 }
-            ),
+            )
+            | tool_call_extension_client.metadata(),
         )
 
         client = a2a.client.ClientFactory(a2a.client.ClientConfig(httpx_client=httpx_client, polling=True)).create(
             card=card
         )
 
         task = None
-        async for event in client.send_message(message):
-            if isinstance(event, a2a.types.Message):
-                print(event)
-                return
-            task, _update = event
+        while True:
+            async for event in client.send_message(message):
+                if isinstance(event, a2a.types.Message):
+                    print(event)
+                    return
+                task, _update = event
 
-        if task and task.status.state == a2a.types.TaskState.auth_required:
-            if not task.status.message:
-                raise RuntimeError("Missing message")
+            if task and task.status.state == a2a.types.TaskState.auth_required:
+                if not task.status.message:
+                    raise RuntimeError("Missing message")
 
-            auth_request = oauth_extension_client.parse_auth_request(message=task.status.message)
+                auth_request = oauth_extension_client.parse_auth_request(message=task.status.message)
 
-            print("Agent has requested authorization")
-            oauth.open_browser(str(auth_request.authorization_endpoint_url))
-            request = await oauth.handle_redirect()
+                print("Agent has requested authorization")
+                oauth.open_browser(str(auth_request.authorization_endpoint_url))
+                request = await oauth.handle_redirect()
 
-            async for event in client.send_message(
-                oauth_extension_client.create_auth_response(task_id=task.id, redirect_uri=AnyUrl(str(request.url)))
-            ):
-                if isinstance(event, a2a.types.Message):
-                    raise RuntimeError("Agent responded with message to a task")
-                task, _update = event
+                message = oauth_extension_client.create_auth_response(
+                    task_id=task.id, redirect_uri=AnyUrl(str(request.url))
+                )
+            elif task and task.status.state == a2a.types.TaskState.input_required:
+                if not task.status.message:
+                    raise RuntimeError("Missing message")
+
+                approval_request = tool_call_extension_client.parse_message(message=task.status.message)
+
+                print("Agent has requested a tool call")
+                print(approval_request)
+                choice = input("Approve (Y/n): ")
+                response = ToolCallResponse(action="accept" if choice.lower() == "y" else "reject")
+                message = tool_call_extension_client.create_message(task_id=task.id, response=response)
+            else:
+                break
 
         print(task)
 

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

@@ -2,5 +2,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from .auth import *
+from .mcp import *
 from .services import *
 from .ui import *

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 .tool_call import *

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

@@ -0,0 +1,83 @@
+# 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 mcp.types import Implementation, ToolAnnotations
+
+from agentstack_sdk.a2a.extensions.base import BaseExtensionClient, BaseExtensionServer, BaseExtensionSpec
+from agentstack_sdk.a2a.types import AgentMessage, InputRequired
+
+if TYPE_CHECKING:
+    from agentstack_sdk.server.context import RunContext
+
+
+class ToolCallRequest(pydantic.BaseModel):
+    server: Implementation | None
+    name: str
+    arguments: dict[str, Any] | None
+    annotations: ToolAnnotations | None
+
+
+class ToolCallResponse(pydantic.BaseModel):
+    action: Literal["accept", "reject"]
+
+
+class ToolCallExtensionParams(pydantic.BaseModel):
+    pass
+
+
+class ToolCallExtensionSpec(BaseExtensionSpec[ToolCallExtensionParams]):
+    URI: str = "https://a2a-extensions.agentstack.beeai.dev/mcp/tool-call/v1"
+
+
+class ToolCallExtensionMetadata(pydantic.BaseModel):
+    pass
+
+
+class ToolCallExtensionServer(BaseExtensionServer[ToolCallExtensionSpec, ToolCallExtensionMetadata]):
+    def create_message(self, *, request: ToolCallRequest):
+        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 ToolCallResponse.model_validate(data)
+
+    async def request_tool_call(self, request: ToolCallRequest, context: RunContext) -> ToolCallResponse:
+        message = self.create_message(request=request)
+        message = await context.yield_async(InputRequired(message=message))
+        if message:
+            result = self.parse_message(message=message)
+            if result.action != "accept":
+                raise RuntimeError("User has rejected the tool call")
+            return result
+        else:
+            raise RuntimeError("Tool call approval response is missing")
+
+
+class ToolCallExtensionClient(BaseExtensionClient[ToolCallExtensionSpec, NoneType]):
+    def create_message(self, *, response: ToolCallResponse, 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 tool call request")
+        return ToolCallRequest.model_validate(data)
+
+    def metadata(self) -> dict[str, Any]:
+        return {self.spec.URI: ToolCallExtensionMetadata().model_dump(mode="json")}

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,93 @@
+# 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, InitializeResult, ListToolsResult, Tool
+
+from agentstack_sdk.a2a.extensions.mcp.tool_call import ToolCallExtensionServer, ToolCallRequest
+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
+        self._server_info: Implementation | None = None
+        self._observed_tools: dict[str, Tool] = {}
+
+    def apply(self, extension: ToolCallExtensionServer) -> Self:
+        self._mcp_extension = extension
+        return self
+
+    async def initialize(self) -> InitializeResult:
+        result = await super().initialize()
+        self._server_info = result.serverInfo
+        return result
+
+    async def list_tools(self, cursor: str | None = None) -> ListToolsResult:
+        result = await super().list_tools(cursor=cursor)
+        for tool in result.tools:
+            self._observed_tools[tool.name] = tool
+        return result
+
+    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:
+            tool = self._observed_tools.get(name)
+            await self._mcp_extension.request_tool_call(
+                request=ToolCallRequest(
+                    server=self._server_info,
+                    name=name,
+                    arguments=arguments,
+                    annotations=tool.annotations if tool else None,
+                ),
+                context=self._context,
+            )
+            return await super().call_tool(
+                name,
+                arguments,
+                read_timeout_seconds,
+                progress_callback,
+            )
+        else:
+            return await super().call_tool(name, arguments, read_timeout_seconds, progress_callback)