hud-python 0.4.35__py3-none-any.whl → 0.4.37__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Potentially problematic release.
This version of hud-python might be problematic. Click here for more details.
- hud/agents/__init__.py +2 -0
- hud/agents/lite_llm.py +72 -0
- hud/agents/openai_chat_generic.py +21 -7
- hud/agents/tests/test_claude.py +32 -7
- hud/agents/tests/test_openai.py +29 -6
- hud/cli/__init__.py +228 -79
- hud/cli/build.py +26 -6
- hud/cli/dev.py +21 -40
- hud/cli/eval.py +96 -15
- hud/cli/flows/tasks.py +198 -65
- hud/cli/init.py +222 -629
- hud/cli/pull.py +6 -0
- hud/cli/push.py +11 -1
- hud/cli/rl/__init__.py +14 -4
- hud/cli/rl/celebrate.py +187 -0
- hud/cli/rl/config.py +15 -8
- hud/cli/rl/local_runner.py +44 -20
- hud/cli/rl/remote_runner.py +166 -87
- hud/cli/rl/viewer.py +141 -0
- hud/cli/rl/wait_utils.py +89 -0
- hud/cli/tests/test_build.py +3 -27
- hud/cli/tests/test_mcp_server.py +1 -12
- hud/cli/utils/config.py +85 -0
- hud/cli/utils/docker.py +21 -39
- hud/cli/utils/env_check.py +196 -0
- hud/cli/utils/environment.py +4 -3
- hud/cli/utils/interactive.py +2 -1
- hud/cli/utils/local_runner.py +204 -0
- hud/cli/utils/metadata.py +3 -1
- hud/cli/utils/package_runner.py +292 -0
- hud/cli/utils/remote_runner.py +4 -1
- hud/cli/utils/source_hash.py +108 -0
- hud/clients/base.py +1 -1
- hud/clients/fastmcp.py +1 -1
- hud/clients/mcp_use.py +30 -7
- hud/datasets/parallel.py +3 -1
- hud/datasets/runner.py +4 -1
- hud/otel/config.py +1 -1
- hud/otel/context.py +40 -6
- hud/rl/buffer.py +3 -0
- hud/rl/tests/test_learner.py +1 -1
- hud/rl/vllm_adapter.py +1 -1
- hud/server/server.py +234 -7
- hud/server/tests/test_add_tool.py +60 -0
- hud/server/tests/test_context.py +128 -0
- hud/server/tests/test_mcp_server_handlers.py +44 -0
- hud/server/tests/test_mcp_server_integration.py +405 -0
- hud/server/tests/test_mcp_server_more.py +247 -0
- hud/server/tests/test_run_wrapper.py +53 -0
- hud/server/tests/test_server_extra.py +166 -0
- hud/server/tests/test_sigterm_runner.py +78 -0
- hud/settings.py +38 -0
- hud/shared/hints.py +2 -2
- hud/telemetry/job.py +2 -2
- hud/types.py +9 -2
- hud/utils/tasks.py +32 -24
- hud/utils/tests/test_version.py +1 -1
- hud/version.py +1 -1
- {hud_python-0.4.35.dist-info → hud_python-0.4.37.dist-info}/METADATA +43 -23
- {hud_python-0.4.35.dist-info → hud_python-0.4.37.dist-info}/RECORD +63 -46
- {hud_python-0.4.35.dist-info → hud_python-0.4.37.dist-info}/WHEEL +0 -0
- {hud_python-0.4.35.dist-info → hud_python-0.4.37.dist-info}/entry_points.txt +0 -0
- {hud_python-0.4.35.dist-info → hud_python-0.4.37.dist-info}/licenses/LICENSE +0 -0
|
@@ -0,0 +1,405 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import asyncio
|
|
4
|
+
import socket
|
|
5
|
+
from contextlib import suppress
|
|
6
|
+
from typing import Any, cast
|
|
7
|
+
|
|
8
|
+
import pytest
|
|
9
|
+
|
|
10
|
+
from hud.clients import MCPClient
|
|
11
|
+
from hud.server import MCPServer
|
|
12
|
+
from hud.server import server as server_mod # for toggling _sigterm_received
|
|
13
|
+
from hud.server.low_level import LowLevelServerWithInit
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
def _free_port() -> int:
|
|
17
|
+
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
|
18
|
+
s.bind(("127.0.0.1", 0))
|
|
19
|
+
return s.getsockname()[1]
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
async def _start_http_server(mcp: MCPServer, port: int) -> asyncio.Task:
|
|
23
|
+
# run the server in the background; cancel to stop
|
|
24
|
+
task = asyncio.create_task(
|
|
25
|
+
mcp.run_async(
|
|
26
|
+
transport="http",
|
|
27
|
+
host="127.0.0.1",
|
|
28
|
+
port=port,
|
|
29
|
+
path="/mcp",
|
|
30
|
+
log_level="ERROR",
|
|
31
|
+
show_banner=False,
|
|
32
|
+
)
|
|
33
|
+
)
|
|
34
|
+
# brief yield so uvicorn can boot
|
|
35
|
+
await asyncio.sleep(0.05)
|
|
36
|
+
return task
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
def _first_text(result) -> str | None:
|
|
40
|
+
# Result.content is usually a list of TextContent
|
|
41
|
+
c = getattr(result, "content", None)
|
|
42
|
+
if isinstance(c, list) and c and hasattr(c[0], "text"):
|
|
43
|
+
return c[0].text
|
|
44
|
+
if isinstance(c, str):
|
|
45
|
+
return c
|
|
46
|
+
return None
|
|
47
|
+
|
|
48
|
+
|
|
49
|
+
@pytest.mark.asyncio
|
|
50
|
+
async def test_low_level_injection_happens_when_initialize_used() -> None:
|
|
51
|
+
mcp = MCPServer(name="InitInject")
|
|
52
|
+
assert not isinstance(mcp._mcp_server, LowLevelServerWithInit)
|
|
53
|
+
|
|
54
|
+
@mcp.initialize
|
|
55
|
+
async def _init(_ctx) -> None:
|
|
56
|
+
return None
|
|
57
|
+
|
|
58
|
+
assert isinstance(mcp._mcp_server, LowLevelServerWithInit)
|
|
59
|
+
|
|
60
|
+
|
|
61
|
+
@pytest.mark.asyncio
|
|
62
|
+
async def test_initialize_runs_once_and_tools_work() -> None:
|
|
63
|
+
port = _free_port()
|
|
64
|
+
|
|
65
|
+
mcp = MCPServer(name="ServerInitOnce")
|
|
66
|
+
state = {"init_calls": 0, "initialized": False}
|
|
67
|
+
|
|
68
|
+
@mcp.initialize
|
|
69
|
+
async def _init(_ctx) -> None:
|
|
70
|
+
# this would corrupt stdout if not redirected; we rely on stderr redirection
|
|
71
|
+
print("hello from init") # noqa: T201
|
|
72
|
+
state["init_calls"] += 1
|
|
73
|
+
state["initialized"] = True
|
|
74
|
+
|
|
75
|
+
@mcp.tool()
|
|
76
|
+
async def initialized() -> bool: # type: ignore[override]
|
|
77
|
+
return state["initialized"]
|
|
78
|
+
|
|
79
|
+
@mcp.tool()
|
|
80
|
+
async def echo(text: str = "ok") -> str: # type: ignore[override]
|
|
81
|
+
return f"echo:{text}"
|
|
82
|
+
|
|
83
|
+
server_task = await _start_http_server(mcp, port)
|
|
84
|
+
|
|
85
|
+
async def connect_and_check() -> None:
|
|
86
|
+
cfg = {"srv": {"url": f"http://127.0.0.1:{port}/mcp"}}
|
|
87
|
+
client = MCPClient(mcp_config=cfg, auto_trace=False, verbose=False)
|
|
88
|
+
await client.initialize()
|
|
89
|
+
tools = await client.list_tools()
|
|
90
|
+
names = sorted(t.name for t in tools)
|
|
91
|
+
assert {"echo", "initialized"} <= set(names)
|
|
92
|
+
res = await client.call_tool(name="initialized", arguments={})
|
|
93
|
+
# boolean return is exposed via structuredContent["result"]
|
|
94
|
+
assert getattr(res, "structuredContent", {}).get("result") is True
|
|
95
|
+
res2 = await client.call_tool(name="echo", arguments={"text": "ping"})
|
|
96
|
+
assert _first_text(res2) == "echo:ping"
|
|
97
|
+
await client.shutdown()
|
|
98
|
+
|
|
99
|
+
try:
|
|
100
|
+
await connect_and_check()
|
|
101
|
+
await connect_and_check()
|
|
102
|
+
# initializer should have executed only once across multiple clients
|
|
103
|
+
assert state["init_calls"] == 1
|
|
104
|
+
finally:
|
|
105
|
+
with suppress(asyncio.CancelledError):
|
|
106
|
+
server_task.cancel()
|
|
107
|
+
await server_task
|
|
108
|
+
|
|
109
|
+
|
|
110
|
+
@pytest.mark.asyncio
|
|
111
|
+
async def test_shutdown_handler_only_on_sigterm_flag() -> None:
|
|
112
|
+
port = _free_port()
|
|
113
|
+
|
|
114
|
+
mcp = MCPServer(name="ShutdownTest")
|
|
115
|
+
called = asyncio.Event()
|
|
116
|
+
|
|
117
|
+
@mcp.shutdown
|
|
118
|
+
async def _on_shutdown() -> None:
|
|
119
|
+
called.set()
|
|
120
|
+
|
|
121
|
+
# no SIGTERM flag: should NOT call shutdown on normal cancel
|
|
122
|
+
server_task = await _start_http_server(mcp, port)
|
|
123
|
+
try:
|
|
124
|
+
# sanity connect so lifespan actually ran
|
|
125
|
+
cfg = {"srv": {"url": f"http://127.0.0.1:{port}/mcp"}}
|
|
126
|
+
c = MCPClient(mcp_config=cfg, auto_trace=False, verbose=False)
|
|
127
|
+
await c.initialize()
|
|
128
|
+
await c.shutdown()
|
|
129
|
+
finally:
|
|
130
|
+
with suppress(asyncio.CancelledError):
|
|
131
|
+
server_task.cancel()
|
|
132
|
+
await server_task
|
|
133
|
+
# give a tick to let lifespan finally run
|
|
134
|
+
await asyncio.sleep(0.05)
|
|
135
|
+
assert not called.is_set()
|
|
136
|
+
|
|
137
|
+
# now start again and simulate SIGTERM so shutdown handler fires
|
|
138
|
+
called.clear()
|
|
139
|
+
port2 = _free_port()
|
|
140
|
+
server_task2 = await _start_http_server(mcp, port=port2)
|
|
141
|
+
try:
|
|
142
|
+
cfg = {"srv": {"url": f"http://127.0.0.1:{port2}/mcp"}}
|
|
143
|
+
c = MCPClient(mcp_config=cfg, auto_trace=False, verbose=False)
|
|
144
|
+
await c.initialize()
|
|
145
|
+
await c.shutdown()
|
|
146
|
+
|
|
147
|
+
# flip the module-level flag the lifespan checks
|
|
148
|
+
server_mod._sigterm_received = True # type: ignore[attr-defined]
|
|
149
|
+
finally:
|
|
150
|
+
with suppress(asyncio.CancelledError):
|
|
151
|
+
server_task2.cancel()
|
|
152
|
+
await server_task2
|
|
153
|
+
|
|
154
|
+
# shutdown coroutine should have run because flag was set when lifespan exited
|
|
155
|
+
assert called.is_set()
|
|
156
|
+
# reset the flag for any other tests
|
|
157
|
+
server_mod._sigterm_received = False # type: ignore[attr-defined]
|
|
158
|
+
|
|
159
|
+
|
|
160
|
+
@pytest.mark.asyncio
|
|
161
|
+
async def test_initializer_exception_propagates_to_client() -> None:
|
|
162
|
+
port = _free_port()
|
|
163
|
+
|
|
164
|
+
mcp = MCPServer(name="InitError")
|
|
165
|
+
|
|
166
|
+
@mcp.initialize
|
|
167
|
+
async def _init(_ctx) -> None:
|
|
168
|
+
raise RuntimeError("boom during init")
|
|
169
|
+
|
|
170
|
+
server_task = await _start_http_server(mcp, port)
|
|
171
|
+
|
|
172
|
+
cfg = {"srv": {"url": f"http://127.0.0.1:{port}/mcp"}}
|
|
173
|
+
client = MCPClient(mcp_config=cfg, auto_trace=False, verbose=False)
|
|
174
|
+
|
|
175
|
+
try:
|
|
176
|
+
with pytest.raises(Exception):
|
|
177
|
+
await client.initialize()
|
|
178
|
+
finally:
|
|
179
|
+
with suppress(asyncio.CancelledError):
|
|
180
|
+
server_task.cancel()
|
|
181
|
+
await server_task
|
|
182
|
+
# defensive: client may or may not be fully created
|
|
183
|
+
with suppress(Exception):
|
|
184
|
+
await client.shutdown()
|
|
185
|
+
|
|
186
|
+
|
|
187
|
+
# --- additional tests for MCPServer coverage ---
|
|
188
|
+
|
|
189
|
+
|
|
190
|
+
@pytest.mark.asyncio
|
|
191
|
+
async def test_init_after_tools_preserves_handlers_and_runs_once() -> None:
|
|
192
|
+
"""If tools are added BEFORE @mcp.initialize, the handler copy during
|
|
193
|
+
low-level server replacement must keep them; init should still run once total.
|
|
194
|
+
"""
|
|
195
|
+
port = _free_port()
|
|
196
|
+
|
|
197
|
+
mcp = MCPServer(name="InitAfterTools")
|
|
198
|
+
state = {"init_calls": 0}
|
|
199
|
+
|
|
200
|
+
# Register tools first
|
|
201
|
+
@mcp.tool()
|
|
202
|
+
async def foo() -> str: # type: ignore[override]
|
|
203
|
+
return "bar"
|
|
204
|
+
|
|
205
|
+
# Now register initializer (this triggers server replacement)
|
|
206
|
+
@mcp.initialize
|
|
207
|
+
async def _init(_ctx) -> None:
|
|
208
|
+
state["init_calls"] += 1
|
|
209
|
+
|
|
210
|
+
server_task = await _start_http_server(mcp, port)
|
|
211
|
+
|
|
212
|
+
async def connect_and_check() -> None:
|
|
213
|
+
cfg = {"srv": {"url": f"http://127.0.0.1:{port}/mcp"}}
|
|
214
|
+
c = MCPClient(mcp_config=cfg, auto_trace=False, verbose=False)
|
|
215
|
+
await c.initialize()
|
|
216
|
+
tools = await c.list_tools()
|
|
217
|
+
names = sorted(t.name for t in tools)
|
|
218
|
+
assert "foo" in names, "tool registered before @initialize must survive replacement"
|
|
219
|
+
res = await c.call_tool(name="foo", arguments={})
|
|
220
|
+
assert _first_text(res) == "bar"
|
|
221
|
+
await c.shutdown()
|
|
222
|
+
|
|
223
|
+
try:
|
|
224
|
+
await connect_and_check()
|
|
225
|
+
await connect_and_check()
|
|
226
|
+
assert state["init_calls"] == 1, "initializer should execute exactly once"
|
|
227
|
+
finally:
|
|
228
|
+
with suppress(asyncio.CancelledError):
|
|
229
|
+
server_task.cancel()
|
|
230
|
+
await server_task
|
|
231
|
+
|
|
232
|
+
|
|
233
|
+
@pytest.mark.asyncio
|
|
234
|
+
async def test_tool_default_argument_used_when_omitted() -> None:
|
|
235
|
+
"""Echo tool should use its default when argument is omitted."""
|
|
236
|
+
port = _free_port()
|
|
237
|
+
|
|
238
|
+
mcp = MCPServer(name="EchoDefault")
|
|
239
|
+
|
|
240
|
+
@mcp.tool()
|
|
241
|
+
async def echo(text: str = "ok") -> str: # type: ignore[override]
|
|
242
|
+
return f"echo:{text}"
|
|
243
|
+
|
|
244
|
+
server_task = await _start_http_server(mcp, port)
|
|
245
|
+
try:
|
|
246
|
+
cfg = {"srv": {"url": f"http://127.0.0.1:{port}/mcp"}}
|
|
247
|
+
c = MCPClient(mcp_config=cfg, auto_trace=False, verbose=False)
|
|
248
|
+
await c.initialize()
|
|
249
|
+
# Call with no args → default should kick in
|
|
250
|
+
res = await c.call_tool(name="echo", arguments={})
|
|
251
|
+
assert _first_text(res) == "echo:ok"
|
|
252
|
+
await c.shutdown()
|
|
253
|
+
finally:
|
|
254
|
+
with suppress(asyncio.CancelledError):
|
|
255
|
+
server_task.cancel()
|
|
256
|
+
await server_task
|
|
257
|
+
|
|
258
|
+
|
|
259
|
+
@pytest.mark.asyncio
|
|
260
|
+
async def test_shutdown_handler_runs_once_when_both_paths_fire() -> None:
|
|
261
|
+
"""With SIGTERM flag set, both the lifespan.finally and run_async.finally would
|
|
262
|
+
try to invoke @mcp.shutdown. The per-instance guard must ensure exactly once.
|
|
263
|
+
"""
|
|
264
|
+
port = _free_port()
|
|
265
|
+
mcp = MCPServer(name="ShutdownOnce")
|
|
266
|
+
calls = {"n": 0}
|
|
267
|
+
|
|
268
|
+
@mcp.shutdown
|
|
269
|
+
async def _on_shutdown() -> None:
|
|
270
|
+
calls["n"] += 1
|
|
271
|
+
|
|
272
|
+
server_task = await _start_http_server(mcp, port)
|
|
273
|
+
try:
|
|
274
|
+
# Ensure lifespan started
|
|
275
|
+
cfg = {"srv": {"url": f"http://127.0.0.1:{port}/mcp"}}
|
|
276
|
+
c = MCPClient(mcp_config=cfg, auto_trace=False, verbose=False)
|
|
277
|
+
await c.initialize()
|
|
278
|
+
await c.shutdown()
|
|
279
|
+
|
|
280
|
+
# Arm SIGTERM flag so both code paths believe they should run
|
|
281
|
+
server_mod._sigterm_received = True # type: ignore[attr-defined]
|
|
282
|
+
finally:
|
|
283
|
+
with suppress(asyncio.CancelledError):
|
|
284
|
+
server_task.cancel()
|
|
285
|
+
await server_task
|
|
286
|
+
|
|
287
|
+
# Give the event loop a tick to run both finalizers
|
|
288
|
+
await asyncio.sleep(0.05)
|
|
289
|
+
|
|
290
|
+
try:
|
|
291
|
+
assert calls["n"] == 1, f"shutdown hook must run exactly once, got {calls['n']}"
|
|
292
|
+
finally:
|
|
293
|
+
# Always reset module flag
|
|
294
|
+
server_mod._sigterm_received = False # type: ignore[attr-defined]
|
|
295
|
+
|
|
296
|
+
|
|
297
|
+
@pytest.mark.asyncio
|
|
298
|
+
async def test_initialize_ctx_exposes_client_info() -> None:
|
|
299
|
+
"""Initializer gets a ctx; clientInfo may be absent depending on client implementation."""
|
|
300
|
+
port = _free_port()
|
|
301
|
+
|
|
302
|
+
mcp = MCPServer(name="InitCtx")
|
|
303
|
+
seen = {"has_session": False, "client_name": None}
|
|
304
|
+
|
|
305
|
+
@mcp.initialize
|
|
306
|
+
async def _init(ctx) -> None: # type: ignore[override]
|
|
307
|
+
# Ensure we have a session object
|
|
308
|
+
seen["has_session"] = hasattr(ctx, "session") and ctx.session is not None
|
|
309
|
+
|
|
310
|
+
# Client info is optional; capture it if present
|
|
311
|
+
client_info = getattr(getattr(ctx.session, "client_params", None), "clientInfo", None)
|
|
312
|
+
if client_info is not None:
|
|
313
|
+
seen["client_name"] = getattr(client_info, "name", None)
|
|
314
|
+
|
|
315
|
+
server_task = await _start_http_server(mcp, port)
|
|
316
|
+
try:
|
|
317
|
+
cfg = {"srv": {"url": f"http://127.0.0.1:{port}/mcp"}}
|
|
318
|
+
c = MCPClient(mcp_config=cfg, auto_trace=False, verbose=False)
|
|
319
|
+
await c.initialize()
|
|
320
|
+
await c.shutdown()
|
|
321
|
+
finally:
|
|
322
|
+
with suppress(asyncio.CancelledError):
|
|
323
|
+
server_task.cancel()
|
|
324
|
+
await server_task
|
|
325
|
+
|
|
326
|
+
assert seen["has_session"] is True
|
|
327
|
+
# If present, name should be a string; otherwise None is acceptable.
|
|
328
|
+
assert seen["client_name"] is None or isinstance(seen["client_name"], str)
|
|
329
|
+
|
|
330
|
+
|
|
331
|
+
@pytest.mark.asyncio
|
|
332
|
+
async def test_initialize_redirects_stdout_to_stderr(capsys) -> None:
|
|
333
|
+
"""Initializer prints should be redirected to stderr (never stdout)."""
|
|
334
|
+
port = _free_port()
|
|
335
|
+
|
|
336
|
+
mcp = MCPServer(name="StdoutRedirect")
|
|
337
|
+
|
|
338
|
+
@mcp.initialize
|
|
339
|
+
async def _init(_ctx) -> None:
|
|
340
|
+
# This would normally pollute STDOUT; our server redirects to STDERR
|
|
341
|
+
print("INIT_STDOUT_MARKER") # noqa: T201
|
|
342
|
+
|
|
343
|
+
server_task = await _start_http_server(mcp, port)
|
|
344
|
+
|
|
345
|
+
try:
|
|
346
|
+
cfg = {"srv": {"url": f"http://127.0.0.1:{port}/mcp"}}
|
|
347
|
+
c = MCPClient(mcp_config=cfg, auto_trace=False, verbose=False)
|
|
348
|
+
await c.initialize()
|
|
349
|
+
await c.shutdown()
|
|
350
|
+
finally:
|
|
351
|
+
with suppress(asyncio.CancelledError):
|
|
352
|
+
server_task.cancel()
|
|
353
|
+
await server_task
|
|
354
|
+
|
|
355
|
+
captured = capsys.readouterr()
|
|
356
|
+
assert "INIT_STDOUT_MARKER" in captured.err
|
|
357
|
+
assert "INIT_STDOUT_MARKER" not in captured.out
|
|
358
|
+
|
|
359
|
+
|
|
360
|
+
@pytest.mark.asyncio
|
|
361
|
+
async def test_initialize_callable_form_runs_once() -> None:
|
|
362
|
+
"""Coverage for mcp.initialize(fn) (callable style), not only decorator usage."""
|
|
363
|
+
port = _free_port()
|
|
364
|
+
mcp = MCPServer(name="CallableInit")
|
|
365
|
+
hits = {"n": 0}
|
|
366
|
+
|
|
367
|
+
async def _init(_ctx) -> None:
|
|
368
|
+
hits["n"] += 1
|
|
369
|
+
|
|
370
|
+
# Callable form instead of decorator
|
|
371
|
+
mcp.initialize(_init)
|
|
372
|
+
|
|
373
|
+
server_task = await _start_http_server(mcp, port)
|
|
374
|
+
try:
|
|
375
|
+
cfg = {"srv": {"url": f"http://127.0.0.1:{port}/mcp"}}
|
|
376
|
+
c1 = MCPClient(mcp_config=cfg, auto_trace=False, verbose=False)
|
|
377
|
+
await c1.initialize()
|
|
378
|
+
await c1.shutdown()
|
|
379
|
+
|
|
380
|
+
c2 = MCPClient(mcp_config=cfg, auto_trace=False, verbose=False)
|
|
381
|
+
await c2.initialize()
|
|
382
|
+
await c2.shutdown()
|
|
383
|
+
finally:
|
|
384
|
+
with suppress(asyncio.CancelledError):
|
|
385
|
+
server_task.cancel()
|
|
386
|
+
await server_task
|
|
387
|
+
|
|
388
|
+
assert hits["n"] == 1
|
|
389
|
+
|
|
390
|
+
|
|
391
|
+
@pytest.mark.asyncio
|
|
392
|
+
async def test_notification_handlers_survive_real_replacement() -> None:
|
|
393
|
+
"""End-to-end check that notification handlers survive when initialize is registered."""
|
|
394
|
+
mcp = MCPServer(name="NotifCopy")
|
|
395
|
+
|
|
396
|
+
# Seed a dummy notification handler before replacement
|
|
397
|
+
cast("dict[Any, Any]", mcp._mcp_server.notification_handlers)["hud/notify"] = object()
|
|
398
|
+
assert "hud/notify" in mcp._mcp_server.notification_handlers
|
|
399
|
+
|
|
400
|
+
@mcp.initialize
|
|
401
|
+
async def _init(_ctx) -> None:
|
|
402
|
+
pass
|
|
403
|
+
|
|
404
|
+
# After replacement, the handler should still be there
|
|
405
|
+
assert "hud/notify" in mcp._mcp_server.notification_handlers
|
|
@@ -0,0 +1,247 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import asyncio
|
|
4
|
+
import socket
|
|
5
|
+
from contextlib import asynccontextmanager, suppress
|
|
6
|
+
|
|
7
|
+
import anyio
|
|
8
|
+
import pytest
|
|
9
|
+
|
|
10
|
+
from hud.clients import MCPClient
|
|
11
|
+
from hud.server import MCPServer
|
|
12
|
+
from hud.server import server as server_mod # to toggle _sigterm_received
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
def _free_port() -> int:
|
|
16
|
+
"""Get a free port for testing."""
|
|
17
|
+
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
|
18
|
+
s.bind(("127.0.0.1", 0))
|
|
19
|
+
return s.getsockname()[1]
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
@asynccontextmanager
|
|
23
|
+
async def _fake_stdio_server():
|
|
24
|
+
"""
|
|
25
|
+
Stand-in for mcp.server.stdio.stdio_server that avoids reading real stdin.
|
|
26
|
+
|
|
27
|
+
It yields a pair of in-memory streams (receive, send) so the low-level server
|
|
28
|
+
can start and idle without touching sys.stdin/sys.stdout.
|
|
29
|
+
"""
|
|
30
|
+
# Server reads from recv_in and writes to send_out
|
|
31
|
+
send_in, recv_in = anyio.create_memory_object_stream(100) # stdin → server
|
|
32
|
+
send_out, recv_out = anyio.create_memory_object_stream(100) # server → stdout
|
|
33
|
+
try:
|
|
34
|
+
yield recv_in, send_out
|
|
35
|
+
finally:
|
|
36
|
+
# Best effort close; methods exist across anyio versions
|
|
37
|
+
for s in (send_in, recv_in, send_out, recv_out):
|
|
38
|
+
close = getattr(s, "close", None) or getattr(s, "aclose", None)
|
|
39
|
+
try:
|
|
40
|
+
if close is not None:
|
|
41
|
+
res = close()
|
|
42
|
+
if asyncio.iscoroutine(res):
|
|
43
|
+
await res
|
|
44
|
+
except Exception:
|
|
45
|
+
pass
|
|
46
|
+
|
|
47
|
+
|
|
48
|
+
@pytest.fixture(autouse=True)
|
|
49
|
+
def _patch_stdio(monkeypatch: pytest.MonkeyPatch):
|
|
50
|
+
"""Patch stdio server for all tests to avoid stdin reading issues."""
|
|
51
|
+
monkeypatch.setenv("FASTMCP_DISABLE_BANNER", "1")
|
|
52
|
+
# Patch both the source and the bound symbol FastMCP uses
|
|
53
|
+
monkeypatch.setattr("mcp.server.stdio.stdio_server", _fake_stdio_server)
|
|
54
|
+
monkeypatch.setattr("fastmcp.server.server.stdio_server", _fake_stdio_server)
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
@pytest.mark.asyncio
|
|
58
|
+
async def test_stdio_shutdown_handler_on_sigterm_flag() -> None:
|
|
59
|
+
"""@mcp.shutdown runs on stdio transport when the SIGTERM flag is set."""
|
|
60
|
+
mcp = MCPServer(name="StdIOShutdown")
|
|
61
|
+
calls = {"n": 0}
|
|
62
|
+
|
|
63
|
+
@mcp.shutdown
|
|
64
|
+
async def _on_shutdown() -> None:
|
|
65
|
+
calls["n"] += 1
|
|
66
|
+
|
|
67
|
+
task = asyncio.create_task(mcp.run_async(transport="stdio", show_banner=False))
|
|
68
|
+
try:
|
|
69
|
+
await asyncio.sleep(0.05)
|
|
70
|
+
# Simulate SIGTERM path
|
|
71
|
+
server_mod._sigterm_received = True # type: ignore[attr-defined]
|
|
72
|
+
finally:
|
|
73
|
+
with suppress(asyncio.CancelledError):
|
|
74
|
+
task.cancel()
|
|
75
|
+
await task
|
|
76
|
+
|
|
77
|
+
assert calls["n"] == 1
|
|
78
|
+
assert not getattr(server_mod, "_sigterm_received")
|
|
79
|
+
|
|
80
|
+
|
|
81
|
+
@pytest.mark.asyncio
|
|
82
|
+
async def test_stdio_shutdown_handler_not_called_without_sigterm() -> None:
|
|
83
|
+
"""@mcp.shutdown must NOT run on stdio cancel when no SIGTERM flag."""
|
|
84
|
+
mcp = MCPServer(name="StdIONoSigterm")
|
|
85
|
+
called = {"n": 0}
|
|
86
|
+
|
|
87
|
+
@mcp.shutdown
|
|
88
|
+
async def _on_shutdown() -> None:
|
|
89
|
+
called["n"] += 1
|
|
90
|
+
|
|
91
|
+
task = asyncio.create_task(mcp.run_async(transport="stdio", show_banner=False))
|
|
92
|
+
try:
|
|
93
|
+
await asyncio.sleep(0.05)
|
|
94
|
+
# no SIGTERM flag
|
|
95
|
+
finally:
|
|
96
|
+
with suppress(asyncio.CancelledError):
|
|
97
|
+
task.cancel()
|
|
98
|
+
await task
|
|
99
|
+
|
|
100
|
+
assert called["n"] == 0
|
|
101
|
+
|
|
102
|
+
|
|
103
|
+
@pytest.mark.asyncio
|
|
104
|
+
async def test_last_initialize_handler_wins_and_ctx_shape_exists() -> None:
|
|
105
|
+
"""If multiple @initialize decorators are applied, only the last one should execute.
|
|
106
|
+
Also sanity-check that ctx has the expected core attributes in a version-tolerant way.
|
|
107
|
+
"""
|
|
108
|
+
port = _free_port()
|
|
109
|
+
|
|
110
|
+
mcp = MCPServer(name="InitOverride")
|
|
111
|
+
seen = {"a": False, "b": False, "has_session": False, "has_request": False}
|
|
112
|
+
|
|
113
|
+
@mcp.initialize
|
|
114
|
+
async def _init_a(ctx) -> None: # type: ignore[override]
|
|
115
|
+
# This one should get overridden and never run
|
|
116
|
+
seen["a"] = True
|
|
117
|
+
|
|
118
|
+
@mcp.initialize
|
|
119
|
+
async def _init_b(ctx) -> None: # type: ignore[override]
|
|
120
|
+
# This is the one that should actually run
|
|
121
|
+
seen["b"] = True
|
|
122
|
+
seen["has_session"] = hasattr(ctx, "session") and ctx.session is not None
|
|
123
|
+
seen["has_request"] = hasattr(ctx, "request") and ctx.request is not None
|
|
124
|
+
|
|
125
|
+
# A simple echo tool so we can verify the server works post-init
|
|
126
|
+
@mcp.tool()
|
|
127
|
+
async def echo(text: str = "ok") -> str: # type: ignore[override]
|
|
128
|
+
return f"echo:{text}"
|
|
129
|
+
|
|
130
|
+
# Start HTTP transport (quickest way to use a real client)
|
|
131
|
+
task = asyncio.create_task(
|
|
132
|
+
mcp.run_async(
|
|
133
|
+
transport="http",
|
|
134
|
+
host="127.0.0.1",
|
|
135
|
+
port=port,
|
|
136
|
+
path="/mcp",
|
|
137
|
+
log_level="ERROR",
|
|
138
|
+
show_banner=False,
|
|
139
|
+
)
|
|
140
|
+
)
|
|
141
|
+
await asyncio.sleep(0.05)
|
|
142
|
+
|
|
143
|
+
try:
|
|
144
|
+
cfg = {"srv": {"url": f"http://127.0.0.1:{port}/mcp"}}
|
|
145
|
+
c = MCPClient(mcp_config=cfg, auto_trace=False, verbose=False)
|
|
146
|
+
await c.initialize()
|
|
147
|
+
|
|
148
|
+
# Call a tool to ensure init didn't break anything
|
|
149
|
+
res = await c.call_tool(name="echo", arguments={"text": "ping"})
|
|
150
|
+
text = getattr(res, "content", None)
|
|
151
|
+
if isinstance(text, list) and text and hasattr(text[0], "text"):
|
|
152
|
+
text = text[0].text
|
|
153
|
+
assert text == "echo:ping"
|
|
154
|
+
|
|
155
|
+
await c.shutdown()
|
|
156
|
+
finally:
|
|
157
|
+
with suppress(asyncio.CancelledError):
|
|
158
|
+
task.cancel()
|
|
159
|
+
await task
|
|
160
|
+
|
|
161
|
+
# Only the last initializer should have run
|
|
162
|
+
assert seen["a"] is False
|
|
163
|
+
assert seen["b"] is True
|
|
164
|
+
# And the ctx had the key attributes (shape may vary by lib version; just presence)
|
|
165
|
+
assert seen["has_session"] is True
|
|
166
|
+
assert seen["has_request"] is True
|
|
167
|
+
|
|
168
|
+
|
|
169
|
+
@pytest.mark.asyncio
|
|
170
|
+
async def test_stdio_shutdown_handler_runs_once_when_both_paths_fire() -> None:
|
|
171
|
+
"""Even on stdio, when SIGTERM is set, ensure shutdown runs exactly once."""
|
|
172
|
+
mcp = MCPServer(name="StdIOOnce")
|
|
173
|
+
calls = {"n": 0}
|
|
174
|
+
|
|
175
|
+
@mcp.shutdown
|
|
176
|
+
async def _on_shutdown() -> None:
|
|
177
|
+
calls["n"] += 1
|
|
178
|
+
|
|
179
|
+
task = asyncio.create_task(mcp.run_async(transport="stdio", show_banner=False))
|
|
180
|
+
try:
|
|
181
|
+
await asyncio.sleep(0.05)
|
|
182
|
+
# Make both the lifespan.finally and run_async.finally want to execute
|
|
183
|
+
server_mod._sigterm_received = True # type: ignore[attr-defined]
|
|
184
|
+
finally:
|
|
185
|
+
with suppress(asyncio.CancelledError):
|
|
186
|
+
task.cancel()
|
|
187
|
+
await task
|
|
188
|
+
|
|
189
|
+
assert calls["n"] == 1
|
|
190
|
+
# Reset global flag always
|
|
191
|
+
server_mod._sigterm_received = False # type: ignore[attr-defined]
|
|
192
|
+
|
|
193
|
+
|
|
194
|
+
@pytest.mark.asyncio
|
|
195
|
+
async def test_run_async_defaults_to_stdio_and_uses_patched_stdio(monkeypatch: pytest.MonkeyPatch):
|
|
196
|
+
"""transport=None should default to stdio and use our patched stdio server."""
|
|
197
|
+
entered = {"v": False}
|
|
198
|
+
|
|
199
|
+
@asynccontextmanager
|
|
200
|
+
async def tracking_stdio():
|
|
201
|
+
entered["v"] = True
|
|
202
|
+
async with _fake_stdio_server() as streams:
|
|
203
|
+
yield streams
|
|
204
|
+
|
|
205
|
+
# Override the autouse fixture for this test to track entry
|
|
206
|
+
monkeypatch.setattr("fastmcp.server.server.stdio_server", tracking_stdio)
|
|
207
|
+
|
|
208
|
+
mcp = MCPServer(name="DefaultStdio")
|
|
209
|
+
task = asyncio.create_task(mcp.run_async(transport=None, show_banner=False))
|
|
210
|
+
try:
|
|
211
|
+
await asyncio.sleep(0.05)
|
|
212
|
+
assert entered["v"] is True, "Expected stdio transport to be used by default"
|
|
213
|
+
finally:
|
|
214
|
+
with suppress(asyncio.CancelledError):
|
|
215
|
+
task.cancel()
|
|
216
|
+
await task
|
|
217
|
+
|
|
218
|
+
|
|
219
|
+
@pytest.mark.asyncio
|
|
220
|
+
async def test_custom_lifespan_relies_on_run_async_fallback_for_sigterm() -> None:
|
|
221
|
+
"""When a custom lifespan is supplied, run_async's finally path must still call
|
|
222
|
+
@shutdown on SIGTERM."""
|
|
223
|
+
|
|
224
|
+
@asynccontextmanager
|
|
225
|
+
async def custom_lifespan(_):
|
|
226
|
+
# No shutdown call here on purpose
|
|
227
|
+
yield {}
|
|
228
|
+
|
|
229
|
+
mcp = MCPServer(name="CustomLS", lifespan=custom_lifespan)
|
|
230
|
+
calls = {"n": 0}
|
|
231
|
+
|
|
232
|
+
@mcp.shutdown
|
|
233
|
+
async def _s() -> None:
|
|
234
|
+
calls["n"] += 1
|
|
235
|
+
|
|
236
|
+
task = asyncio.create_task(mcp.run_async(transport="stdio", show_banner=False))
|
|
237
|
+
try:
|
|
238
|
+
await asyncio.sleep(0.05)
|
|
239
|
+
# Ensure finalizer believes SIGTERM happened
|
|
240
|
+
server_mod._sigterm_received = True # type: ignore[attr-defined]
|
|
241
|
+
finally:
|
|
242
|
+
with suppress(asyncio.CancelledError):
|
|
243
|
+
task.cancel()
|
|
244
|
+
await task
|
|
245
|
+
|
|
246
|
+
assert calls["n"] == 1
|
|
247
|
+
assert not getattr(server_mod, "_sigterm_received")
|
|
@@ -0,0 +1,53 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import asyncio
|
|
4
|
+
from typing import TYPE_CHECKING
|
|
5
|
+
|
|
6
|
+
from hud.server import MCPServer
|
|
7
|
+
|
|
8
|
+
if TYPE_CHECKING:
|
|
9
|
+
import pytest
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
def test_run_uses_sigterm_wrapper(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
13
|
+
"""MCPServer.run should delegate through _run_with_sigterm (don't actually start a server)."""
|
|
14
|
+
called = {"hit": False, "args": None, "kwargs": None}
|
|
15
|
+
|
|
16
|
+
def fake_wrapper(coro_fn, *args, **kwargs):
|
|
17
|
+
called["hit"] = True
|
|
18
|
+
called["args"] = args
|
|
19
|
+
called["kwargs"] = kwargs
|
|
20
|
+
# Do not execute the bootstrap coroutine; this is unit wiring only.
|
|
21
|
+
|
|
22
|
+
monkeypatch.setattr("hud.server.server._run_with_sigterm", fake_wrapper)
|
|
23
|
+
|
|
24
|
+
mcp = MCPServer(name="RunWrapper")
|
|
25
|
+
# Should immediately return after calling our fake wrapper
|
|
26
|
+
mcp.run(transport="http", host="127.0.0.1", port=9999, path="/mcp", show_banner=False)
|
|
27
|
+
|
|
28
|
+
assert called["hit"] is True
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
def test_run_defaults_to_stdio(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
32
|
+
"""transport=None in .run should resolve to 'stdio' and forward to run_async."""
|
|
33
|
+
seen = {}
|
|
34
|
+
|
|
35
|
+
async def fake_run_async(self, *, transport, show_banner, **kwargs):
|
|
36
|
+
seen["transport"] = transport
|
|
37
|
+
seen["show_banner"] = show_banner
|
|
38
|
+
seen["kwargs"] = kwargs
|
|
39
|
+
|
|
40
|
+
# Replace bound method on the instance class
|
|
41
|
+
monkeypatch.setattr(MCPServer, "run_async", fake_run_async, raising=False)
|
|
42
|
+
|
|
43
|
+
# Execute the bootstrap coroutine immediately (no real server)
|
|
44
|
+
def fake_wrapper(coro_fn, *args, **kwargs):
|
|
45
|
+
asyncio.run(coro_fn())
|
|
46
|
+
|
|
47
|
+
monkeypatch.setattr("hud.server.server._run_with_sigterm", fake_wrapper)
|
|
48
|
+
|
|
49
|
+
mcp = MCPServer(name="RunDefaultTransport")
|
|
50
|
+
mcp.run(transport=None, show_banner=False)
|
|
51
|
+
|
|
52
|
+
assert seen["transport"] == "stdio"
|
|
53
|
+
assert seen["show_banner"] is False
|