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.

Files changed (63) hide show
  1. hud/agents/__init__.py +2 -0
  2. hud/agents/lite_llm.py +72 -0
  3. hud/agents/openai_chat_generic.py +21 -7
  4. hud/agents/tests/test_claude.py +32 -7
  5. hud/agents/tests/test_openai.py +29 -6
  6. hud/cli/__init__.py +228 -79
  7. hud/cli/build.py +26 -6
  8. hud/cli/dev.py +21 -40
  9. hud/cli/eval.py +96 -15
  10. hud/cli/flows/tasks.py +198 -65
  11. hud/cli/init.py +222 -629
  12. hud/cli/pull.py +6 -0
  13. hud/cli/push.py +11 -1
  14. hud/cli/rl/__init__.py +14 -4
  15. hud/cli/rl/celebrate.py +187 -0
  16. hud/cli/rl/config.py +15 -8
  17. hud/cli/rl/local_runner.py +44 -20
  18. hud/cli/rl/remote_runner.py +166 -87
  19. hud/cli/rl/viewer.py +141 -0
  20. hud/cli/rl/wait_utils.py +89 -0
  21. hud/cli/tests/test_build.py +3 -27
  22. hud/cli/tests/test_mcp_server.py +1 -12
  23. hud/cli/utils/config.py +85 -0
  24. hud/cli/utils/docker.py +21 -39
  25. hud/cli/utils/env_check.py +196 -0
  26. hud/cli/utils/environment.py +4 -3
  27. hud/cli/utils/interactive.py +2 -1
  28. hud/cli/utils/local_runner.py +204 -0
  29. hud/cli/utils/metadata.py +3 -1
  30. hud/cli/utils/package_runner.py +292 -0
  31. hud/cli/utils/remote_runner.py +4 -1
  32. hud/cli/utils/source_hash.py +108 -0
  33. hud/clients/base.py +1 -1
  34. hud/clients/fastmcp.py +1 -1
  35. hud/clients/mcp_use.py +30 -7
  36. hud/datasets/parallel.py +3 -1
  37. hud/datasets/runner.py +4 -1
  38. hud/otel/config.py +1 -1
  39. hud/otel/context.py +40 -6
  40. hud/rl/buffer.py +3 -0
  41. hud/rl/tests/test_learner.py +1 -1
  42. hud/rl/vllm_adapter.py +1 -1
  43. hud/server/server.py +234 -7
  44. hud/server/tests/test_add_tool.py +60 -0
  45. hud/server/tests/test_context.py +128 -0
  46. hud/server/tests/test_mcp_server_handlers.py +44 -0
  47. hud/server/tests/test_mcp_server_integration.py +405 -0
  48. hud/server/tests/test_mcp_server_more.py +247 -0
  49. hud/server/tests/test_run_wrapper.py +53 -0
  50. hud/server/tests/test_server_extra.py +166 -0
  51. hud/server/tests/test_sigterm_runner.py +78 -0
  52. hud/settings.py +38 -0
  53. hud/shared/hints.py +2 -2
  54. hud/telemetry/job.py +2 -2
  55. hud/types.py +9 -2
  56. hud/utils/tasks.py +32 -24
  57. hud/utils/tests/test_version.py +1 -1
  58. hud/version.py +1 -1
  59. {hud_python-0.4.35.dist-info → hud_python-0.4.37.dist-info}/METADATA +43 -23
  60. {hud_python-0.4.35.dist-info → hud_python-0.4.37.dist-info}/RECORD +63 -46
  61. {hud_python-0.4.35.dist-info → hud_python-0.4.37.dist-info}/WHEEL +0 -0
  62. {hud_python-0.4.35.dist-info → hud_python-0.4.37.dist-info}/entry_points.txt +0 -0
  63. {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