Spaces:
Running
Running
File size: 6,887 Bytes
17763eb 0c29eb8 17763eb 0c29eb8 17763eb 0c29eb8 17763eb 0c29eb8 09438a8 0c29eb8 aa46c09 0c29eb8 aa46c09 0c29eb8 17763eb 0c29eb8 aa46c09 0c29eb8 aa46c09 0c29eb8 aa46c09 0c29eb8 aa46c09 0c29eb8 aa46c09 0c29eb8 aa46c09 0c29eb8 aa46c09 0c29eb8 9e4d18c 0c29eb8 445f86a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 | import asyncio
import gc
import inspect
import os
import weakref
import psutil
import pytest
from fastmcp import Client
from fastmcp.client.transports import PythonStdioTransport, StdioTransport
def running_under_debugger():
return os.environ.get("DEBUGPY_RUNNING") == "true"
def gc_collect_harder():
gc.collect()
gc.collect()
gc.collect()
gc.collect()
gc.collect()
gc.collect()
class TestKeepAlive:
# https://github.com/jlowin/fastmcp/issues/581
@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_keep_alive_default_true(self):
client = Client(transport=StdioTransport(command="python", args=[""]))
assert client.transport.keep_alive is True
async def test_keep_alive_set_false(self):
client = Client(
transport=StdioTransport(command="python", args=[""], keep_alive=False)
)
assert client.transport.keep_alive is False
async def test_keep_alive_maintains_session_across_multiple_calls(
self, stdio_script
):
client = Client(transport=PythonStdioTransport(script_path=stdio_script))
assert client.transport.keep_alive is True
async with client:
result1 = await client.call_tool("pid")
pid1: int = result1.data
async with client:
result2 = await client.call_tool("pid")
pid2: int = result2.data
assert pid1 == pid2
@pytest.mark.skipif(
running_under_debugger(), reason="Debugger holds a reference to the transport"
)
async def test_keep_alive_true_exit_scope_kills_transport(self, stdio_script):
transport_weak_ref: weakref.ref[PythonStdioTransport] | None = None
async def test_server():
transport = PythonStdioTransport(script_path=stdio_script, keep_alive=True)
nonlocal transport_weak_ref
transport_weak_ref = weakref.ref(transport)
async with transport.connect_session():
pass
await test_server()
gc_collect_harder()
# This test will fail while debugging because the debugger holds a reference to the underlying transport
assert transport_weak_ref
transport = transport_weak_ref()
assert transport is None
@pytest.mark.skipif(
running_under_debugger(), reason="Debugger holds a reference to the transport"
)
async def test_keep_alive_true_exit_scope_kills_client(self, stdio_script):
pid: int | None = None
async def test_server():
transport = PythonStdioTransport(script_path=stdio_script, keep_alive=True)
client = Client(transport=transport)
assert client.transport.keep_alive is True
async with client:
result1 = await client.call_tool("pid")
nonlocal pid
pid = result1.data
await test_server()
gc_collect_harder()
# This test may fail/hang while debugging because the debugger holds a reference to the underlying transport
with pytest.raises(psutil.NoSuchProcess):
while True:
psutil.Process(pid)
await asyncio.sleep(0.1)
async def test_keep_alive_false_exit_scope_kills_server(self, stdio_script):
pid: int | None = None
async def test_server():
transport = PythonStdioTransport(script_path=stdio_script, keep_alive=False)
client = Client(transport=transport)
assert client.transport.keep_alive is False
async with client:
result1 = await client.call_tool("pid")
nonlocal pid
pid = result1.data
del client
await test_server()
with pytest.raises(psutil.NoSuchProcess):
while True:
psutil.Process(pid)
await asyncio.sleep(0.1)
async def test_keep_alive_false_starts_new_session_across_multiple_calls(
self, stdio_script
):
client = Client(
transport=PythonStdioTransport(script_path=stdio_script, keep_alive=False)
)
assert client.transport.keep_alive is False
async with client:
result1 = await client.call_tool("pid")
pid1: int = result1.data
async with client:
result2 = await client.call_tool("pid")
pid2: int = result2.data
assert pid1 != pid2
async def test_keep_alive_starts_new_session_if_manually_closed(self, stdio_script):
client = Client(transport=PythonStdioTransport(script_path=stdio_script))
assert client.transport.keep_alive is True
async with client:
result1 = await client.call_tool("pid")
pid1: int = result1.data
await client.close()
async with client:
result2 = await client.call_tool("pid")
pid2: int = result2.data
assert pid1 != pid2
async def test_keep_alive_maintains_session_if_reentered(self, stdio_script):
client = Client(transport=PythonStdioTransport(script_path=stdio_script))
assert client.transport.keep_alive is True
async with client:
result1 = await client.call_tool("pid")
pid1: int = result1.data
async with client:
result2 = await client.call_tool("pid")
pid2: int = result2.data
result3 = await client.call_tool("pid")
pid3: int = result3.data
assert pid1 == pid2 == pid3
async def test_close_session_and_try_to_use_client_raises_error(self, stdio_script):
client = Client(transport=PythonStdioTransport(script_path=stdio_script))
assert client.transport.keep_alive is True
async with client:
await client.close()
with pytest.raises(RuntimeError, match="Client is not connected"):
await client.call_tool("pid")
async def test_session_task_failure_raises_immediately_on_enter(self):
# Use a command that will fail to start
client = Client(
transport=StdioTransport(command="nonexistent_command", args=[])
)
# Should raise RuntimeError immediately, not defer until first use
with pytest.raises(RuntimeError, match="Client failed to connect"):
async with client:
pass
|