Spaces:
Running
Running
Jeremiah Lowin commited on
Commit ·
c04707c
1
Parent(s): d4ff4a7
Ensure client is only initialized once
Browse files
src/fastmcp/client/client.py
CHANGED
|
@@ -145,6 +145,7 @@ class Client(Generic[ClientTransportT]):
|
|
| 145 |
progress_handler: ProgressHandler | None = None,
|
| 146 |
timeout: datetime.timedelta | float | int | None = None,
|
| 147 |
init_timeout: datetime.timedelta | float | int | None = None,
|
|
|
|
| 148 |
auth: httpx.Auth | Literal["oauth"] | str | None = None,
|
| 149 |
):
|
| 150 |
self.transport = cast(ClientTransportT, infer_transport(transport))
|
|
@@ -180,6 +181,7 @@ class Client(Generic[ClientTransportT]):
|
|
| 180 |
"logging_callback": create_log_callback(log_handler),
|
| 181 |
"message_handler": message_handler,
|
| 182 |
"read_timeout_seconds": timeout,
|
|
|
|
| 183 |
}
|
| 184 |
|
| 185 |
if roots is not None:
|
|
|
|
| 145 |
progress_handler: ProgressHandler | None = None,
|
| 146 |
timeout: datetime.timedelta | float | int | None = None,
|
| 147 |
init_timeout: datetime.timedelta | float | int | None = None,
|
| 148 |
+
client_info: mcp.types.Implementation | None = None,
|
| 149 |
auth: httpx.Auth | Literal["oauth"] | str | None = None,
|
| 150 |
):
|
| 151 |
self.transport = cast(ClientTransportT, infer_transport(transport))
|
|
|
|
| 181 |
"logging_callback": create_log_callback(log_handler),
|
| 182 |
"message_handler": message_handler,
|
| 183 |
"read_timeout_seconds": timeout,
|
| 184 |
+
"client_info": client_info,
|
| 185 |
}
|
| 186 |
|
| 187 |
if roots is not None:
|
src/fastmcp/client/transports.py
CHANGED
|
@@ -8,27 +8,15 @@ import sys
|
|
| 8 |
import warnings
|
| 9 |
from collections.abc import AsyncIterator, Callable
|
| 10 |
from pathlib import Path
|
| 11 |
-
from typing import
|
| 12 |
-
TYPE_CHECKING,
|
| 13 |
-
Any,
|
| 14 |
-
Literal,
|
| 15 |
-
TypedDict,
|
| 16 |
-
TypeVar,
|
| 17 |
-
cast,
|
| 18 |
-
overload,
|
| 19 |
-
)
|
| 20 |
|
| 21 |
import anyio
|
| 22 |
import httpx
|
|
|
|
| 23 |
from mcp import ClientSession, StdioServerParameters
|
| 24 |
-
from mcp.client.session import
|
| 25 |
-
ListRootsFnT,
|
| 26 |
-
LoggingFnT,
|
| 27 |
-
MessageHandlerFnT,
|
| 28 |
-
SamplingFnT,
|
| 29 |
-
)
|
| 30 |
from mcp.server.fastmcp import FastMCP as FastMCP1Server
|
| 31 |
-
from mcp.shared.memory import
|
| 32 |
from pydantic import AnyUrl
|
| 33 |
from typing_extensions import Unpack
|
| 34 |
|
|
@@ -65,11 +53,12 @@ __all__ = [
|
|
| 65 |
class SessionKwargs(TypedDict, total=False):
|
| 66 |
"""Keyword arguments for the MCP ClientSession constructor."""
|
| 67 |
|
|
|
|
| 68 |
sampling_callback: SamplingFnT | None
|
| 69 |
list_roots_callback: ListRootsFnT | None
|
| 70 |
logging_callback: LoggingFnT | None
|
| 71 |
message_handler: MessageHandlerFnT | None
|
| 72 |
-
|
| 73 |
|
| 74 |
|
| 75 |
class ClientTransport(abc.ABC):
|
|
@@ -662,24 +651,47 @@ class FastMCPTransport(ClientTransport):
|
|
| 662 |
tests or scenarios where client and server run in the same runtime.
|
| 663 |
"""
|
| 664 |
|
| 665 |
-
def __init__(self, mcp: FastMCP | FastMCP1Server):
|
| 666 |
"""Initialize a FastMCPTransport from a FastMCP server instance."""
|
| 667 |
|
| 668 |
# Accept both FastMCP 2.x and FastMCP 1.0 servers. Both expose a
|
| 669 |
# ``_mcp_server`` attribute pointing to the underlying MCP server
|
| 670 |
# implementation, so we can treat them identically.
|
| 671 |
self.server = mcp
|
|
|
|
| 672 |
|
| 673 |
@contextlib.asynccontextmanager
|
| 674 |
async def connect_session(
|
| 675 |
self, **session_kwargs: Unpack[SessionKwargs]
|
| 676 |
) -> AsyncIterator[ClientSession]:
|
| 677 |
-
|
| 678 |
-
|
| 679 |
-
|
| 680 |
-
|
| 681 |
-
|
| 682 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 683 |
|
| 684 |
def __repr__(self) -> str:
|
| 685 |
return f"<FastMCPTransport(server='{self.server.name}')>"
|
|
|
|
| 8 |
import warnings
|
| 9 |
from collections.abc import AsyncIterator, Callable
|
| 10 |
from pathlib import Path
|
| 11 |
+
from typing import TYPE_CHECKING, Any, Literal, TypedDict, TypeVar, cast, overload
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 12 |
|
| 13 |
import anyio
|
| 14 |
import httpx
|
| 15 |
+
import mcp.types
|
| 16 |
from mcp import ClientSession, StdioServerParameters
|
| 17 |
+
from mcp.client.session import ListRootsFnT, LoggingFnT, MessageHandlerFnT, SamplingFnT
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 18 |
from mcp.server.fastmcp import FastMCP as FastMCP1Server
|
| 19 |
+
from mcp.shared.memory import create_client_server_memory_streams
|
| 20 |
from pydantic import AnyUrl
|
| 21 |
from typing_extensions import Unpack
|
| 22 |
|
|
|
|
| 53 |
class SessionKwargs(TypedDict, total=False):
|
| 54 |
"""Keyword arguments for the MCP ClientSession constructor."""
|
| 55 |
|
| 56 |
+
read_timeout_seconds: datetime.timedelta | None
|
| 57 |
sampling_callback: SamplingFnT | None
|
| 58 |
list_roots_callback: ListRootsFnT | None
|
| 59 |
logging_callback: LoggingFnT | None
|
| 60 |
message_handler: MessageHandlerFnT | None
|
| 61 |
+
client_info: mcp.types.Implementation | None
|
| 62 |
|
| 63 |
|
| 64 |
class ClientTransport(abc.ABC):
|
|
|
|
| 651 |
tests or scenarios where client and server run in the same runtime.
|
| 652 |
"""
|
| 653 |
|
| 654 |
+
def __init__(self, mcp: FastMCP | FastMCP1Server, raise_exceptions: bool = False):
|
| 655 |
"""Initialize a FastMCPTransport from a FastMCP server instance."""
|
| 656 |
|
| 657 |
# Accept both FastMCP 2.x and FastMCP 1.0 servers. Both expose a
|
| 658 |
# ``_mcp_server`` attribute pointing to the underlying MCP server
|
| 659 |
# implementation, so we can treat them identically.
|
| 660 |
self.server = mcp
|
| 661 |
+
self.raise_exceptions = raise_exceptions
|
| 662 |
|
| 663 |
@contextlib.asynccontextmanager
|
| 664 |
async def connect_session(
|
| 665 |
self, **session_kwargs: Unpack[SessionKwargs]
|
| 666 |
) -> AsyncIterator[ClientSession]:
|
| 667 |
+
async with create_client_server_memory_streams() as (
|
| 668 |
+
client_streams,
|
| 669 |
+
server_streams,
|
| 670 |
+
):
|
| 671 |
+
client_read, client_write = client_streams
|
| 672 |
+
server_read, server_write = server_streams
|
| 673 |
+
|
| 674 |
+
# Create a cancel scope for the server task
|
| 675 |
+
async with anyio.create_task_group() as tg:
|
| 676 |
+
tg.start_soon(
|
| 677 |
+
lambda: self.server._mcp_server.run(
|
| 678 |
+
server_read,
|
| 679 |
+
server_write,
|
| 680 |
+
self.server._mcp_server.create_initialization_options(),
|
| 681 |
+
raise_exceptions=self.raise_exceptions,
|
| 682 |
+
)
|
| 683 |
+
)
|
| 684 |
+
|
| 685 |
+
try:
|
| 686 |
+
async with ClientSession(
|
| 687 |
+
read_stream=client_read,
|
| 688 |
+
write_stream=client_write,
|
| 689 |
+
**session_kwargs,
|
| 690 |
+
client_info=session_kwargs.get("client_info"),
|
| 691 |
+
) as client_session:
|
| 692 |
+
yield client_session
|
| 693 |
+
finally:
|
| 694 |
+
tg.cancel_scope.cancel()
|
| 695 |
|
| 696 |
def __repr__(self) -> str:
|
| 697 |
return f"<FastMCPTransport(server='{self.server.name}')>"
|
tests/client/test_client.py
CHANGED
|
@@ -1,7 +1,9 @@
|
|
| 1 |
import asyncio
|
| 2 |
import sys
|
| 3 |
from typing import cast
|
|
|
|
| 4 |
|
|
|
|
| 5 |
import pytest
|
| 6 |
from mcp import McpError
|
| 7 |
from mcp.client.auth import OAuthClientProvider
|
|
@@ -275,6 +277,14 @@ async def test_client_connection(fastmcp_server):
|
|
| 275 |
assert not client.is_connected()
|
| 276 |
|
| 277 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 278 |
async def test_initialize_result_connected(fastmcp_server):
|
| 279 |
"""Test that initialize_result returns the correct result when connected."""
|
| 280 |
client = Client(transport=FastMCPTransport(fastmcp_server))
|
|
|
|
| 1 |
import asyncio
|
| 2 |
import sys
|
| 3 |
from typing import cast
|
| 4 |
+
from unittest.mock import AsyncMock
|
| 5 |
|
| 6 |
+
import mcp
|
| 7 |
import pytest
|
| 8 |
from mcp import McpError
|
| 9 |
from mcp.client.auth import OAuthClientProvider
|
|
|
|
| 277 |
assert not client.is_connected()
|
| 278 |
|
| 279 |
|
| 280 |
+
async def test_initialize_called_once(fastmcp_server, monkeypatch):
|
| 281 |
+
mock_initialize = AsyncMock()
|
| 282 |
+
monkeypatch.setattr(mcp.ClientSession, "initialize", mock_initialize)
|
| 283 |
+
client = Client(transport=FastMCPTransport(fastmcp_server))
|
| 284 |
+
async with client:
|
| 285 |
+
assert mock_initialize.call_count == 1
|
| 286 |
+
|
| 287 |
+
|
| 288 |
async def test_initialize_result_connected(fastmcp_server):
|
| 289 |
"""Test that initialize_result returns the correct result when connected."""
|
| 290 |
client = Client(transport=FastMCPTransport(fastmcp_server))
|
tests/server/test_server_interactions.py
CHANGED
|
@@ -617,7 +617,7 @@ class TestToolContextInjection:
|
|
| 617 |
result = await client.call_tool("tool_with_context", {"x": 42})
|
| 618 |
assert len(result) == 1
|
| 619 |
content = result[0]
|
| 620 |
-
assert content.text == "
|
| 621 |
|
| 622 |
async def test_async_context(self):
|
| 623 |
"""Test that context works in async functions."""
|
|
@@ -632,7 +632,7 @@ class TestToolContextInjection:
|
|
| 632 |
result = await client.call_tool("async_tool", {"x": 42})
|
| 633 |
assert len(result) == 1
|
| 634 |
content = result[0]
|
| 635 |
-
assert content.text == "Async request
|
| 636 |
|
| 637 |
async def test_optional_context(self):
|
| 638 |
"""Test that context is optional."""
|
|
@@ -696,7 +696,7 @@ class TestToolContextInjection:
|
|
| 696 |
|
| 697 |
async with Client(mcp) as client:
|
| 698 |
result = await client.call_tool("MyTool", {"x": 2})
|
| 699 |
-
assert result[0].text == "
|
| 700 |
|
| 701 |
|
| 702 |
class TestResource:
|
|
@@ -780,7 +780,7 @@ class TestResourceContext:
|
|
| 780 |
|
| 781 |
async with Client(mcp) as client:
|
| 782 |
result = await client.read_resource(AnyUrl("resource://test"))
|
| 783 |
-
assert result[0].text == "
|
| 784 |
|
| 785 |
|
| 786 |
class TestResourceTemplates:
|
|
@@ -1015,7 +1015,7 @@ class TestResourceTemplateContext:
|
|
| 1015 |
|
| 1016 |
async with Client(mcp) as client:
|
| 1017 |
result = await client.read_resource(AnyUrl("resource://test"))
|
| 1018 |
-
assert result[0].text.startswith("Resource template: test
|
| 1019 |
|
| 1020 |
async def test_resource_template_context_with_callable_object(self):
|
| 1021 |
mcp = FastMCP()
|
|
@@ -1031,7 +1031,7 @@ class TestResourceTemplateContext:
|
|
| 1031 |
|
| 1032 |
async with Client(mcp) as client:
|
| 1033 |
result = await client.read_resource(AnyUrl("resource://test"))
|
| 1034 |
-
assert result[0].text.startswith("Resource template: test
|
| 1035 |
|
| 1036 |
|
| 1037 |
class TestPrompts:
|
|
@@ -1249,4 +1249,4 @@ class TestPromptContext:
|
|
| 1249 |
assert len(result.messages) == 1
|
| 1250 |
message = result.messages[0]
|
| 1251 |
assert message.role == "user"
|
| 1252 |
-
assert message.content.text == "Hello, World!
|
|
|
|
| 617 |
result = await client.call_tool("tool_with_context", {"x": 42})
|
| 618 |
assert len(result) == 1
|
| 619 |
content = result[0]
|
| 620 |
+
assert content.text == "1" # type: ignore[attr-defined]
|
| 621 |
|
| 622 |
async def test_async_context(self):
|
| 623 |
"""Test that context works in async functions."""
|
|
|
|
| 632 |
result = await client.call_tool("async_tool", {"x": 42})
|
| 633 |
assert len(result) == 1
|
| 634 |
content = result[0]
|
| 635 |
+
assert content.text == "Async request 1: 42" # type: ignore[attr-defined]
|
| 636 |
|
| 637 |
async def test_optional_context(self):
|
| 638 |
"""Test that context is optional."""
|
|
|
|
| 696 |
|
| 697 |
async with Client(mcp) as client:
|
| 698 |
result = await client.call_tool("MyTool", {"x": 2})
|
| 699 |
+
assert result[0].text == "3" # type: ignore[attr-defined]
|
| 700 |
|
| 701 |
|
| 702 |
class TestResource:
|
|
|
|
| 780 |
|
| 781 |
async with Client(mcp) as client:
|
| 782 |
result = await client.read_resource(AnyUrl("resource://test"))
|
| 783 |
+
assert result[0].text == "1" # type: ignore[attr-defined]
|
| 784 |
|
| 785 |
|
| 786 |
class TestResourceTemplates:
|
|
|
|
| 1015 |
|
| 1016 |
async with Client(mcp) as client:
|
| 1017 |
result = await client.read_resource(AnyUrl("resource://test"))
|
| 1018 |
+
assert result[0].text.startswith("Resource template: test 1") # type: ignore[attr-defined]
|
| 1019 |
|
| 1020 |
async def test_resource_template_context_with_callable_object(self):
|
| 1021 |
mcp = FastMCP()
|
|
|
|
| 1031 |
|
| 1032 |
async with Client(mcp) as client:
|
| 1033 |
result = await client.read_resource(AnyUrl("resource://test"))
|
| 1034 |
+
assert result[0].text.startswith("Resource template: test 1") # type: ignore[attr-defined]
|
| 1035 |
|
| 1036 |
|
| 1037 |
class TestPrompts:
|
|
|
|
| 1249 |
assert len(result.messages) == 1
|
| 1250 |
message = result.messages[0]
|
| 1251 |
assert message.role == "user"
|
| 1252 |
+
assert message.content.text == "Hello, World! 1" # type: ignore[attr-defined]
|