refactor: flatten the methods in Client (#1914)

This commit is contained in:
Marcelo Trylesinski 2026-01-20 13:12:07 +01:00 • committed by GitHub
commit bcb07c2b25
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 213 additions and 315 deletions

View file

@ -36,12 +36,7 @@ class InMemoryTransport:
result = await client.call_tool("my_tool", {...})
"""
def __init__(
self,
server: Server[Any] | FastMCP,
*,
raise_exceptions: bool = False,
) -> None:
def __init__(self, server: Server[Any] | FastMCP, *, raise_exceptions: bool = False) -> None:
"""Initialize the in-memory transport.
Args:

View file

@ -58,6 +58,7 @@ class Client:
self,
server: Server[Any] | FastMCP,
*,
# TODO(Marcelo): When do `raise_exceptions=True` actually raises?
raise_exceptions: bool = False,
read_timeout_seconds: float | None = None,
sampling_callback: SamplingFnT | None = None,
@ -125,14 +126,9 @@ class Client:
self._exit_stack = exit_stack.pop_all()
return self
async def __aexit__(
self,
exc_type: type[BaseException] | None,
exc_val: BaseException | None,
exc_tb: Any,
) -> None:
async def __aexit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any) -> None:
"""Exit the async context manager."""
if self._exit_stack:
if self._exit_stack: # pragma: no branch
await self._exit_stack.__aexit__(exc_type, exc_val, exc_tb)
self._session = None
@ -177,28 +173,22 @@ class Client:
"""Set the logging level on the server."""
return await self.session.set_logging_level(level)
async def list_resources(
self,
params: types.PaginatedRequestParams | None = None,
) -> types.ListResourcesResult:
async def list_resources(self, *, cursor: str | None = None) -> types.ListResourcesResult:
"""List available resources from the server."""
return await self.session.list_resources(params=params)
return await self.session.list_resources(params=types.PaginatedRequestParams(cursor=cursor))
async def list_resource_templates(
self,
params: types.PaginatedRequestParams | None = None,
) -> types.ListResourceTemplatesResult:
async def list_resource_templates(self, *, cursor: str | None = None) -> types.ListResourceTemplatesResult:
"""List available resource templates from the server."""
return await self.session.list_resource_templates(params=params)
return await self.session.list_resource_templates(params=types.PaginatedRequestParams(cursor=cursor))
async def read_resource(self, uri: str | AnyUrl) -> types.ReadResourceResult:
"""Read a resource from the server.
Args:
uri: The URI of the resource to read
uri: The URI of the resource to read.
Returns:
The resource content
The resource content.
"""
return await self.session.read_resource(uri)
@ -239,18 +229,11 @@ class Client:
meta=meta,
)
async def list_prompts(
self,
params: types.PaginatedRequestParams | None = None,
) -> types.ListPromptsResult:
async def list_prompts(self, *, cursor: str | None = None) -> types.ListPromptsResult:
"""List available prompts from the server."""
return await self.session.list_prompts(params=params)
return await self.session.list_prompts(params=types.PaginatedRequestParams(cursor=cursor))
async def get_prompt(
self,
name: str,
arguments: dict[str, str] | None = None,
) -> types.GetPromptResult:
async def get_prompt(self, name: str, arguments: dict[str, str] | None = None) -> types.GetPromptResult:
"""Get a prompt from the server.
Args:
@ -258,7 +241,7 @@ class Client:
arguments: Arguments to pass to the prompt
Returns:
The prompt content
The prompt content.
"""
return await self.session.get_prompt(name=name, arguments=arguments)
@ -276,21 +259,15 @@ class Client:
context_arguments: Additional context arguments
Returns:
Completion suggestions
Completion suggestions.
"""
return await self.session.complete(
ref=ref,
argument=argument,
context_arguments=context_arguments,
)
return await self.session.complete(ref=ref, argument=argument, context_arguments=context_arguments)
async def list_tools(
self,
params: types.PaginatedRequestParams | None = None,
) -> types.ListToolsResult:
async def list_tools(self, *, cursor: str | None = None) -> types.ListToolsResult:
"""List available tools from the server."""
return await self.session.list_tools(params=params)
return await self.session.list_tools(params=types.PaginatedRequestParams(cursor=cursor))
async def send_roots_list_changed(self) -> None:
"""Send a notification that the roots list has changed."""
await self.session.send_roots_list_changed()
# TODO(Marcelo): Currently, there is no way for the server to handle this. We should add support.
await self.session.send_roots_list_changed() # pragma: no cover

View file

@ -349,10 +349,7 @@ class ClientSession(
Args:
params: Full pagination parameters including cursor and any future fields
"""
return await self.send_request(
types.ListPromptsRequest(params=params),
types.ListPromptsResult,
)
return await self.send_request(types.ListPromptsRequest(params=params), types.ListPromptsResult)
async def get_prompt(self, name: str, arguments: dict[str, str] | None = None) -> types.GetPromptResult:
"""Send a prompts/get request."""

View file

@ -234,12 +234,12 @@ class BaseSession(
metadata: MessageMetadata = None,
progress_callback: ProgressFnT | None = None,
) -> ReceiveResultT:
"""Sends a request and wait for a response. Raises an McpError if the
response contains an error. If a request read timeout is provided, it
will take precedence over the session read timeout.
"""Sends a request and wait for a response.
Do not use this method to emit notifications! Use send_notification()
instead.
Raises an McpError if the response contains an error. If a request read timeout is provided, it will take
precedence over the session read timeout.
Do not use this method to emit notifications! Use send_notification() instead.
"""
request_id = self._request_id
self._request_id = request_id + 1
@ -261,15 +261,10 @@ class BaseSession(
try:
jsonrpc_request = JSONRPCRequest(jsonrpc="2.0", id=request_id, **request_data)
await self._write_stream.send(SessionMessage(message=jsonrpc_request, metadata=metadata))
# request read timeout takes precedence over session read timeout
timeout = None
if request_read_timeout_seconds is not None: # pragma: no cover
timeout = request_read_timeout_seconds
elif self._session_read_timeout_seconds is not None: # pragma: no cover
timeout = self._session_read_timeout_seconds
timeout = request_read_timeout_seconds or self._session_read_timeout_seconds
try:
with anyio.fail_after(timeout):
@ -279,9 +274,8 @@ class BaseSession(
ErrorData(
code=httpx.codes.REQUEST_TIMEOUT,
message=(
f"Timed out while waiting for response to "
f"{request.__class__.__name__}. Waited "
f"{timeout} seconds."
f"Timed out while waiting for response to {request.__class__.__name__}. "
f"Waited {timeout} seconds."
),
)
)
@ -302,9 +296,7 @@ class BaseSession(
notification: SendNotificationT,
related_request_id: RequestId | None = None,
) -> None:
"""Emits a notification, which is a one-way message that does not expect
a response.
"""
"""Emits a notification, which is a one-way message that does not expect a response."""
# Some transport implementations may need to set the related_request_id
# to attribute to the notifications to the request that triggered them.
jsonrpc_notification = JSONRPCNotification(
@ -373,11 +365,7 @@ class BaseSession(
error_response = JSONRPCError(
jsonrpc="2.0",
id=message.message.id,
error=ErrorData(
code=INVALID_PARAMS,
message="Invalid request parameters",
data="",
),
error=ErrorData(code=INVALID_PARAMS, message="Invalid request parameters", data=""),
)
session_message = SessionMessage(message=error_response)
await self._write_stream.send(session_message)
@ -518,13 +506,9 @@ class BaseSession(
total: float | None = None,
message: str | None = None,
) -> None:
"""Sends a progress notification for a request that is currently being
processed.
"""
"""Sends a progress notification for a request that is currently being processed."""
async def _handle_incoming(
self,
req: RequestResponder[ReceiveRequestT, SendResultT] | ReceiveNotificationT | Exception,
self, req: RequestResponder[ReceiveRequestT, SendResultT] | ReceiveNotificationT | Exception
) -> None:
"""A generic handler for incoming messages. Overwritten by subclasses."""
pass # pragma: no cover