mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-21 21:14:17 +02:00
stdio client updates
This commit is contained in:
parent
20e20e7eda
commit
f848b4d14c
5 changed files with 91 additions and 111 deletions
|
|
@ -337,22 +337,15 @@ class Client(Generic[ClientTransportT]):
|
|||
await fresh_client.call_tool("some_tool", {})
|
||||
```
|
||||
"""
|
||||
# Warn if using stdio transport - it may cause race conditions
|
||||
from .transports import StdioTransport
|
||||
|
||||
if isinstance(self.transport, StdioTransport):
|
||||
logger.warning(
|
||||
"Creating new client instance for stdio transport. This may cause "
|
||||
"race conditions as stdio servers have only one session. Consider "
|
||||
"reusing the same client instance instead."
|
||||
)
|
||||
|
||||
new_client = copy.copy(self)
|
||||
|
||||
# Reset session state to fresh state
|
||||
new_client._session_state = ClientSessionState()
|
||||
if not isinstance(self.transport, StdioTransport):
|
||||
# Reset session state to fresh state
|
||||
new_client._session_state = ClientSessionState()
|
||||
|
||||
new_client.name += f":{secrets.token_hex(2)}"
|
||||
new_client.name += f":{secrets.token_hex(nbytes=2)}"
|
||||
|
||||
return new_client
|
||||
|
||||
|
|
@ -403,6 +396,7 @@ class Client(Generic[ClientTransportT]):
|
|||
self._session_state.session_task is None
|
||||
or self._session_state.session_task.done()
|
||||
)
|
||||
|
||||
if need_to_start:
|
||||
if self._session_state.nesting_counter != 0:
|
||||
raise RuntimeError(
|
||||
|
|
@ -428,6 +422,7 @@ class Client(Generic[ClientTransportT]):
|
|||
) from exception
|
||||
|
||||
self._session_state.nesting_counter += 1
|
||||
|
||||
return self
|
||||
|
||||
async def _disconnect(self, force: bool = False):
|
||||
|
|
@ -802,7 +797,7 @@ class Client(Generic[ClientTransportT]):
|
|||
Raises:
|
||||
RuntimeError: If called while the client is not connected.
|
||||
"""
|
||||
logger.debug(f"[{self.name}] called list_tools")
|
||||
logger.debug(f"[{self.name}] called list_tools: {self.session._client_info}")
|
||||
|
||||
result = await self.session.list_tools()
|
||||
return result
|
||||
|
|
|
|||
|
|
@ -52,15 +52,6 @@ def mcp_server_type_to_servers_and_transports(
|
|||
ProxyClient(transport=transport, name=client_name)
|
||||
)
|
||||
|
||||
# For stdio transports, create a client factory that reuses the same client
|
||||
# to prevent race conditions from multiple initialization attempts
|
||||
if isinstance(transport, StdioTransport):
|
||||
|
||||
def stdio_client_factory():
|
||||
return client # Reuse the same client instance
|
||||
|
||||
server = FastMCPProxy(name=server_name, client_factory=stdio_client_factory)
|
||||
else:
|
||||
server = FastMCP.as_proxy(name=server_name, backend=client)
|
||||
server = FastMCP.as_proxy(name=server_name, backend=client)
|
||||
|
||||
return name, server, transport
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ import weakref
|
|||
import psutil
|
||||
import pytest
|
||||
|
||||
from fastmcp import Client
|
||||
from fastmcp import Client, FastMCP
|
||||
from fastmcp.client.transports import PythonStdioTransport, StdioTransport
|
||||
|
||||
|
||||
|
|
@ -24,6 +24,44 @@ def gc_collect_harder():
|
|||
gc.collect()
|
||||
|
||||
|
||||
class TestParallelCalls:
|
||||
@pytest.fixture
|
||||
def stdio_script(self, tmp_path):
|
||||
script = inspect.cleandoc('''
|
||||
import os
|
||||
from fastmcp import FastMCP
|
||||
|
||||
mcp = FastMCP()
|
||||
|
||||
@mcp.tool
|
||||
def pid() -> int:
|
||||
"""Gets PID of server"""
|
||||
return os.getpid()
|
||||
|
||||
if __name__ == "__main__":
|
||||
mcp.run()
|
||||
''')
|
||||
script_file = tmp_path / "stdio.py"
|
||||
script_file.write_text(script)
|
||||
return script_file
|
||||
|
||||
async def test_parallel_calls(self, stdio_script):
|
||||
backend_transport = PythonStdioTransport(script_path=stdio_script)
|
||||
backend_client = Client(transport=backend_transport)
|
||||
|
||||
proxy = FastMCP.as_proxy(backend=backend_client, name="PROXY")
|
||||
|
||||
count = 10
|
||||
|
||||
tasks = [proxy.get_tools() for _ in range(count)]
|
||||
|
||||
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
assert len(results) == count
|
||||
errors = [result for result in results if isinstance(result, Exception)]
|
||||
assert len(errors) == 0
|
||||
|
||||
|
||||
class TestKeepAlive:
|
||||
# https://github.com/jlowin/fastmcp/issues/581
|
||||
|
||||
|
|
|
|||
|
|
@ -243,6 +243,50 @@ async def test_multi_client(tmp_path: Path):
|
|||
assert result_2.data == 3
|
||||
|
||||
|
||||
async def test_multi_client_parallel_calls(tmp_path: Path):
|
||||
server_script = inspect.cleandoc("""
|
||||
from fastmcp import FastMCP
|
||||
|
||||
mcp = FastMCP()
|
||||
|
||||
@mcp.tool
|
||||
def add(a: int, b: int) -> int:
|
||||
return a + b
|
||||
|
||||
if __name__ == '__main__':
|
||||
mcp.run()
|
||||
""")
|
||||
|
||||
script_path = tmp_path / "test.py"
|
||||
script_path.write_text(server_script)
|
||||
|
||||
config = {
|
||||
"mcpServers": {
|
||||
"test_1": {
|
||||
"command": "python",
|
||||
"args": [str(script_path)],
|
||||
},
|
||||
"test_2": {
|
||||
"command": "python",
|
||||
"args": [str(script_path)],
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
client = Client(config)
|
||||
|
||||
async with client:
|
||||
_ = await client.list_tools()
|
||||
|
||||
tasks = [client.list_tools() for _ in range(40)]
|
||||
|
||||
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
exceptions = [result for result in results if isinstance(result, Exception)]
|
||||
assert len(exceptions) == 0
|
||||
assert len(results) == 40
|
||||
assert all(len(result) == 2 for result in results)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
running_under_debugger() or sys.platform.startswith("win32"),
|
||||
reason="Debugger holds a reference to the transport; Windows has process lifecycle issues",
|
||||
|
|
|
|||
|
|
@ -1,88 +0,0 @@
|
|||
"""Tests for stdio client reuse in MCPConfig"""
|
||||
|
||||
import asyncio
|
||||
import inspect
|
||||
import tempfile
|
||||
import warnings
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from fastmcp.client import Client
|
||||
from fastmcp.client.transports import PythonStdioTransport
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_config_stdio_client_reuse(tmp_path: Path):
|
||||
"""Test that MCPConfig reuses clients for stdio servers to prevent race conditions."""
|
||||
server_script = inspect.cleandoc("""
|
||||
from fastmcp import FastMCP
|
||||
|
||||
mcp = FastMCP()
|
||||
|
||||
@mcp.tool
|
||||
def add(a: int, b: int) -> int:
|
||||
return a + b
|
||||
|
||||
if __name__ == '__main__':
|
||||
mcp.run()
|
||||
""")
|
||||
|
||||
script_path = tmp_path / "test_server.py"
|
||||
script_path.write_text(server_script)
|
||||
|
||||
config = {
|
||||
"mcpServers": {
|
||||
"test_1": {
|
||||
"command": "python",
|
||||
"args": [str(script_path)],
|
||||
},
|
||||
"test_2": {
|
||||
"command": "python",
|
||||
"args": [str(script_path)],
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
client = Client(config)
|
||||
|
||||
async with client:
|
||||
# Make parallel calls that would previously fail
|
||||
tasks = [client.call_tool("test_1_add", {"a": 1, "b": 2}) for _ in range(20)]
|
||||
|
||||
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
exceptions = [result for result in results if isinstance(result, Exception)]
|
||||
|
||||
# Should have no exceptions and all successful results
|
||||
assert len(exceptions) == 0
|
||||
assert len(results) == 20
|
||||
assert all(result.data == 3 for result in results)
|
||||
|
||||
|
||||
def test_stdio_client_new_warning():
|
||||
"""Test that new() method warns when called on stdio transport clients."""
|
||||
# Create a temporary script file
|
||||
with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
|
||||
f.write(
|
||||
"from fastmcp import FastMCP\nmcp = FastMCP()\nif __name__ == '__main__': mcp.run()"
|
||||
)
|
||||
script_path = Path(f.name)
|
||||
|
||||
try:
|
||||
transport = PythonStdioTransport(script_path=script_path)
|
||||
client = Client(transport=transport)
|
||||
|
||||
# Capture warnings
|
||||
with warnings.catch_warnings(record=True):
|
||||
warnings.simplefilter("always")
|
||||
# This should trigger a warning
|
||||
new_client = client.new()
|
||||
|
||||
# Check that warning was issued
|
||||
# Note: warnings from logger are not captured by warnings module
|
||||
# So we'll check that the method completes without error
|
||||
assert new_client is not None
|
||||
assert new_client != client # Should be a different instance
|
||||
|
||||
finally:
|
||||
script_path.unlink() # Clean up temp file
|
||||
Loading…
Add table
Add a link
Reference in a new issue