Skip to content

Commit 396abb6

Browse files
feat: add MCP Resources support
1 parent 73e7843 commit 396abb6

File tree

2 files changed

+52
-6
lines changed

2 files changed

+52
-6
lines changed

pyproject.toml

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,10 @@ dev = [
7777
]
7878

7979
[tool.uv.workspace]
80-
members = ["agents"]
80+
members = [
81+
"agents",
82+
"checking",
83+
]
8184

8285
[tool.uv.sources]
8386
agents = { workspace = true }

src/agents/mcp/server.py

Lines changed: 48 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -7,15 +7,16 @@
77
from contextlib import AbstractAsyncContextManager, AsyncExitStack
88
from datetime import timedelta
99
from pathlib import Path
10-
from typing import TYPE_CHECKING, Any, Callable, Literal, TypeVar
10+
from typing import TYPE_CHECKING, Any, Callable, Literal, TypeVar, AsyncIterator
11+
1112

1213
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
1314
from mcp import ClientSession, StdioServerParameters, Tool as MCPTool, stdio_client
1415
from mcp.client.session import MessageHandlerFnT
1516
from mcp.client.sse import sse_client
1617
from mcp.client.streamable_http import GetSessionIdCallback, streamablehttp_client
1718
from mcp.shared.message import SessionMessage
18-
from mcp.types import CallToolResult, GetPromptResult, InitializeResult, ListPromptsResult
19+
from mcp.types import CallToolResult, GetPromptResult, InitializeResult, ListPromptsResult, ListResourcesResult, ReadResourceResult
1920
from typing_extensions import NotRequired, TypedDict
2021

2122
from ..exceptions import UserError
@@ -78,6 +79,19 @@ async def call_tool(self, tool_name: str, arguments: dict[str, Any] | None) -> C
7879
"""Invoke a tool on the server."""
7980
pass
8081

82+
@abc.abstractmethod
83+
async def list_resources(self)-> ListResourcesResult:
84+
"""List resources available on the server."""
85+
pass
86+
@abc.abstractmethod
87+
async def read_resource(self, uri:str)-> ReadResourceResult:
88+
"""Read the content of a resource by URI."""
89+
pass
90+
@abc.abstractmethod
91+
async def subscribe_resource(self, uri:str, **kwargs) -> AsyncIterator[SessionMessage]:
92+
"""Subscribe to resource updates."""
93+
pass
94+
8195
@abc.abstractmethod
8296
async def list_prompts(
8397
self,
@@ -91,7 +105,7 @@ async def get_prompt(
91105
) -> GetPromptResult:
92106
"""Get a specific prompt from the server."""
93107
pass
94-
108+
95109

96110
class _MCPServerWithClientSession(MCPServer, abc.ABC):
97111
"""Base class for MCP servers that use a `ClientSession` to communicate with the server."""
@@ -140,7 +154,6 @@ def __init__(
140154
self.max_retry_attempts = max_retry_attempts
141155
self.retry_backoff_seconds_base = retry_backoff_seconds_base
142156
self.message_handler = message_handler
143-
144157
# The cache is always dirty at startup, so that we fetch tools at least once
145158
self._cache_dirty = True
146159
self._tools_list: list[MCPTool] | None = None
@@ -288,6 +301,10 @@ async def connect(self):
288301
await self.cleanup()
289302
raise
290303

304+
305+
306+
307+
291308
async def list_tools(
292309
self,
293310
run_context: RunContextWrapper[Any] | None = None,
@@ -326,6 +343,30 @@ async def call_tool(self, tool_name: str, arguments: dict[str, Any] | None) -> C
326343

327344
return await self._run_with_retries(lambda: session.call_tool(tool_name, arguments))
328345

346+
async def list_resources(self)-> ListResourcesResult:
347+
if not self.session:
348+
raise UserError("Server not initialized. Make sure you call `connect()` first.")
349+
session = self.session
350+
assert session is not None
351+
return await session.list_resources()
352+
353+
async def read_resource(self,uri:str)-> ReadResourceResult:
354+
if not self.session:
355+
raise UserError("Server not initialized. Make sure you call `connect()` first.")
356+
session = self.session
357+
assert session is not None
358+
return await session.read_resource(uri)
359+
360+
async def subscribe_resource(self, uri:str, **kwargs)-> AsyncIterator[SessionMessage]:
361+
if not self.session:
362+
raise UserError("Server not initialized. Make sure you call `connect()` first.")
363+
session = self.session
364+
assert session is not None
365+
async for msg in session.subscribe_resource(uri, **kwargs):
366+
yield msg
367+
368+
369+
329370
async def list_prompts(
330371
self,
331372
) -> ListPromptsResult:
@@ -593,6 +634,8 @@ class MCPServerStreamableHttpParams(TypedDict):
593634
"""Custom HTTP client factory for configuring httpx.AsyncClient behavior."""
594635

595636

637+
638+
596639
class MCPServerStreamableHttp(_MCPServerWithClientSession):
597640
"""MCP server implementation that uses the Streamable HTTP transport. See the [spec]
598641
(https://modelcontextprotocol.io/specification/2025-03-26/basic/transports#streamable-http)
@@ -688,4 +731,4 @@ def create_streams(
688731
@property
689732
def name(self) -> str:
690733
"""A readable name for the server."""
691-
return self._name
734+
return self._name

0 commit comments

Comments
 (0)