hud-python 0.4.10__py3-none-any.whl → 0.4.12__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/__main__.py +8 -0
- hud/agents/base.py +7 -8
- hud/agents/langchain.py +2 -2
- hud/agents/tests/test_openai.py +3 -1
- hud/cli/__init__.py +106 -51
- hud/cli/build.py +121 -71
- hud/cli/debug.py +2 -2
- hud/cli/{mcp_server.py → dev.py} +60 -25
- hud/cli/eval.py +148 -68
- hud/cli/init.py +0 -1
- hud/cli/list_func.py +72 -71
- hud/cli/pull.py +1 -2
- hud/cli/push.py +35 -23
- hud/cli/remove.py +35 -41
- hud/cli/tests/test_analyze.py +2 -1
- hud/cli/tests/test_analyze_metadata.py +42 -49
- hud/cli/tests/test_build.py +28 -52
- hud/cli/tests/test_cursor.py +1 -1
- hud/cli/tests/test_debug.py +1 -1
- hud/cli/tests/test_list_func.py +75 -64
- hud/cli/tests/test_main_module.py +30 -0
- hud/cli/tests/test_mcp_server.py +3 -3
- hud/cli/tests/test_pull.py +30 -61
- hud/cli/tests/test_push.py +70 -89
- hud/cli/tests/test_registry.py +36 -38
- hud/cli/tests/test_utils.py +1 -1
- hud/cli/utils/__init__.py +1 -0
- hud/cli/{docker_utils.py → utils/docker.py} +36 -0
- hud/cli/{env_utils.py → utils/environment.py} +7 -7
- hud/cli/{interactive.py → utils/interactive.py} +91 -19
- hud/cli/{analyze_metadata.py → utils/metadata.py} +12 -8
- hud/cli/{registry.py → utils/registry.py} +28 -30
- hud/cli/{remote_runner.py → utils/remote_runner.py} +1 -1
- hud/cli/utils/runner.py +134 -0
- hud/cli/utils/server.py +250 -0
- hud/clients/base.py +1 -1
- hud/clients/fastmcp.py +7 -5
- hud/clients/mcp_use.py +8 -6
- hud/server/server.py +34 -4
- hud/shared/exceptions.py +11 -0
- hud/shared/tests/test_exceptions.py +22 -0
- hud/telemetry/tests/__init__.py +0 -0
- hud/telemetry/tests/test_replay.py +40 -0
- hud/telemetry/tests/test_trace.py +63 -0
- hud/tools/base.py +20 -3
- hud/tools/computer/hud.py +15 -6
- hud/tools/executors/tests/test_base_executor.py +27 -0
- hud/tools/response.py +15 -4
- hud/tools/tests/test_response.py +60 -0
- hud/tools/tests/test_tools_init.py +49 -0
- hud/utils/design.py +19 -8
- hud/utils/mcp.py +17 -5
- hud/utils/tests/test_mcp.py +112 -0
- hud/utils/tests/test_version.py +1 -1
- hud/version.py +1 -1
- {hud_python-0.4.10.dist-info → hud_python-0.4.12.dist-info}/METADATA +14 -10
- {hud_python-0.4.10.dist-info → hud_python-0.4.12.dist-info}/RECORD +62 -52
- hud/cli/runner.py +0 -160
- /hud/cli/{cursor.py → utils/cursor.py} +0 -0
- /hud/cli/{utils.py → utils/logging.py} +0 -0
- {hud_python-0.4.10.dist-info → hud_python-0.4.12.dist-info}/WHEEL +0 -0
- {hud_python-0.4.10.dist-info → hud_python-0.4.12.dist-info}/entry_points.txt +0 -0
- {hud_python-0.4.10.dist-info → hud_python-0.4.12.dist-info}/licenses/LICENSE +0 -0
hud/cli/utils/runner.py
ADDED
|
@@ -0,0 +1,134 @@
|
|
|
1
|
+
"""Run Docker images as MCP servers."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import asyncio
|
|
6
|
+
import subprocess
|
|
7
|
+
import sys
|
|
8
|
+
|
|
9
|
+
from hud.utils.design import HUDDesign
|
|
10
|
+
|
|
11
|
+
from .logging import find_free_port
|
|
12
|
+
from .server import MCPServerManager, run_server_with_interactive
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
def run_stdio_server(image: str, docker_args: list[str], verbose: bool) -> None:
|
|
16
|
+
"""Run Docker image as stdio MCP server (direct passthrough)."""
|
|
17
|
+
design = HUDDesign() # Use stderr for stdio mode
|
|
18
|
+
|
|
19
|
+
# Build docker command
|
|
20
|
+
docker_cmd = ["docker", "run", "--rm", "-i", *docker_args, image]
|
|
21
|
+
|
|
22
|
+
if verbose:
|
|
23
|
+
design.info(f"🐳 Running: {' '.join(docker_cmd)}")
|
|
24
|
+
|
|
25
|
+
# Run docker directly with stdio passthrough
|
|
26
|
+
try:
|
|
27
|
+
result = subprocess.run(docker_cmd, stdin=sys.stdin) # noqa: S603
|
|
28
|
+
sys.exit(result.returncode)
|
|
29
|
+
except KeyboardInterrupt:
|
|
30
|
+
design.info("\n👋 Shutting down...")
|
|
31
|
+
sys.exit(0)
|
|
32
|
+
except Exception as e:
|
|
33
|
+
design.error(f"Error: {e}")
|
|
34
|
+
sys.exit(1)
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
async def run_http_server(image: str, docker_args: list[str], port: int, verbose: bool) -> None:
|
|
38
|
+
"""Run Docker image as HTTP MCP server (proxy mode)."""
|
|
39
|
+
design = HUDDesign()
|
|
40
|
+
|
|
41
|
+
# Create server manager
|
|
42
|
+
server_manager = MCPServerManager(image, docker_args)
|
|
43
|
+
|
|
44
|
+
# Find available port
|
|
45
|
+
actual_port = find_free_port(port)
|
|
46
|
+
if actual_port is None:
|
|
47
|
+
design.error(f"No available ports found starting from {port}")
|
|
48
|
+
return
|
|
49
|
+
|
|
50
|
+
if actual_port != port:
|
|
51
|
+
design.warning(f"Port {port} in use, using port {actual_port} instead")
|
|
52
|
+
|
|
53
|
+
# Clean up any existing container
|
|
54
|
+
server_manager.cleanup_container()
|
|
55
|
+
|
|
56
|
+
# Build docker command
|
|
57
|
+
docker_cmd = server_manager.build_docker_command()
|
|
58
|
+
|
|
59
|
+
# Create MCP config
|
|
60
|
+
config = server_manager.create_mcp_config(docker_cmd)
|
|
61
|
+
|
|
62
|
+
# Create proxy
|
|
63
|
+
proxy = server_manager.create_proxy(config)
|
|
64
|
+
|
|
65
|
+
# Show header
|
|
66
|
+
design.info("") # Empty line
|
|
67
|
+
design.header("HUD MCP Server", icon="🌐")
|
|
68
|
+
|
|
69
|
+
# Show configuration
|
|
70
|
+
design.section_title("Server Information")
|
|
71
|
+
design.info(f"Port: {actual_port}")
|
|
72
|
+
design.info(f"URL: http://localhost:{actual_port}/mcp")
|
|
73
|
+
design.info(f"Container: {server_manager.container_name}")
|
|
74
|
+
design.info("")
|
|
75
|
+
design.progress_message("Press Ctrl+C to stop")
|
|
76
|
+
|
|
77
|
+
try:
|
|
78
|
+
await server_manager.run_http_server(proxy, actual_port, verbose)
|
|
79
|
+
except KeyboardInterrupt:
|
|
80
|
+
design.info("\n👋 Shutting down...")
|
|
81
|
+
finally:
|
|
82
|
+
# Clean up container
|
|
83
|
+
server_manager.cleanup_container()
|
|
84
|
+
|
|
85
|
+
|
|
86
|
+
async def run_http_server_interactive(
|
|
87
|
+
image: str, docker_args: list[str], port: int, verbose: bool
|
|
88
|
+
) -> None:
|
|
89
|
+
"""Run Docker image as HTTP MCP server with interactive testing."""
|
|
90
|
+
# Create server manager
|
|
91
|
+
server_manager = MCPServerManager(image, docker_args)
|
|
92
|
+
|
|
93
|
+
# Use the shared utility function
|
|
94
|
+
await run_server_with_interactive(server_manager, port, verbose)
|
|
95
|
+
|
|
96
|
+
|
|
97
|
+
def run_mcp_server(
|
|
98
|
+
image: str,
|
|
99
|
+
docker_args: list[str],
|
|
100
|
+
transport: str,
|
|
101
|
+
port: int,
|
|
102
|
+
verbose: bool,
|
|
103
|
+
interactive: bool = False,
|
|
104
|
+
) -> None:
|
|
105
|
+
"""Run Docker image as MCP server with specified transport."""
|
|
106
|
+
if transport == "stdio":
|
|
107
|
+
if interactive:
|
|
108
|
+
design = HUDDesign()
|
|
109
|
+
design.error("Interactive mode requires HTTP transport")
|
|
110
|
+
sys.exit(1)
|
|
111
|
+
run_stdio_server(image, docker_args, verbose)
|
|
112
|
+
elif transport == "http":
|
|
113
|
+
if interactive:
|
|
114
|
+
# Run in interactive mode
|
|
115
|
+
asyncio.run(run_http_server_interactive(image, docker_args, port, verbose))
|
|
116
|
+
else:
|
|
117
|
+
try:
|
|
118
|
+
asyncio.run(run_http_server(image, docker_args, port, verbose))
|
|
119
|
+
except Exception as e:
|
|
120
|
+
# Suppress the graceful shutdown errors
|
|
121
|
+
if not any(
|
|
122
|
+
x in str(e)
|
|
123
|
+
for x in [
|
|
124
|
+
"timeout graceful shutdown exceeded",
|
|
125
|
+
"Cancel 0 running task(s)",
|
|
126
|
+
"Application shutdown complete",
|
|
127
|
+
]
|
|
128
|
+
):
|
|
129
|
+
design = HUDDesign()
|
|
130
|
+
design.error(f"Unexpected error: {e}")
|
|
131
|
+
else:
|
|
132
|
+
design = HUDDesign()
|
|
133
|
+
design.error(f"Unknown transport: {transport}")
|
|
134
|
+
sys.exit(1)
|
hud/cli/utils/server.py
ADDED
|
@@ -0,0 +1,250 @@
|
|
|
1
|
+
"""Common server utilities for HUD CLI."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import asyncio
|
|
6
|
+
from typing import Any
|
|
7
|
+
|
|
8
|
+
from fastmcp import FastMCP
|
|
9
|
+
|
|
10
|
+
from hud.utils.design import HUDDesign
|
|
11
|
+
|
|
12
|
+
from .docker import generate_container_name, remove_container
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
class MCPServerManager:
|
|
16
|
+
"""Manages MCP server lifecycle and configuration."""
|
|
17
|
+
|
|
18
|
+
def __init__(self, image: str, docker_args: list[str] | None = None) -> None:
|
|
19
|
+
"""Initialize server manager.
|
|
20
|
+
|
|
21
|
+
Args:
|
|
22
|
+
image: Docker image name
|
|
23
|
+
docker_args: Additional Docker arguments
|
|
24
|
+
"""
|
|
25
|
+
self.image = image
|
|
26
|
+
self.docker_args = docker_args or []
|
|
27
|
+
self.design = HUDDesign()
|
|
28
|
+
self.container_name = self._generate_container_name()
|
|
29
|
+
|
|
30
|
+
def _generate_container_name(self) -> str:
|
|
31
|
+
"""Generate a unique container name from image."""
|
|
32
|
+
return generate_container_name(self.image)
|
|
33
|
+
|
|
34
|
+
def cleanup_container(self) -> None:
|
|
35
|
+
"""Remove any existing container with the same name."""
|
|
36
|
+
remove_container(self.container_name)
|
|
37
|
+
|
|
38
|
+
def build_docker_command(
|
|
39
|
+
self,
|
|
40
|
+
extra_args: list[str] | None = None,
|
|
41
|
+
entrypoint: list[str] | None = None,
|
|
42
|
+
) -> list[str]:
|
|
43
|
+
"""Build Docker run command.
|
|
44
|
+
|
|
45
|
+
Args:
|
|
46
|
+
extra_args: Additional arguments to add before image
|
|
47
|
+
entrypoint: Custom entrypoint override
|
|
48
|
+
|
|
49
|
+
Returns:
|
|
50
|
+
Complete docker command as list
|
|
51
|
+
"""
|
|
52
|
+
cmd = [
|
|
53
|
+
"docker",
|
|
54
|
+
"run",
|
|
55
|
+
"--rm",
|
|
56
|
+
"-i",
|
|
57
|
+
"--name",
|
|
58
|
+
self.container_name,
|
|
59
|
+
]
|
|
60
|
+
|
|
61
|
+
# Add extra args (like volume mounts, env vars)
|
|
62
|
+
if extra_args:
|
|
63
|
+
cmd.extend(extra_args)
|
|
64
|
+
|
|
65
|
+
# Add user-provided docker args
|
|
66
|
+
cmd.extend(self.docker_args)
|
|
67
|
+
|
|
68
|
+
# Add entrypoint if specified
|
|
69
|
+
if entrypoint:
|
|
70
|
+
cmd.extend(["--entrypoint", entrypoint[0]])
|
|
71
|
+
|
|
72
|
+
# Add image
|
|
73
|
+
cmd.append(self.image)
|
|
74
|
+
|
|
75
|
+
# Add entrypoint args if specified
|
|
76
|
+
if entrypoint and len(entrypoint) > 1:
|
|
77
|
+
cmd.extend(entrypoint[1:])
|
|
78
|
+
|
|
79
|
+
return cmd
|
|
80
|
+
|
|
81
|
+
def create_mcp_config(self, docker_cmd: list[str]) -> dict[str, Any]:
|
|
82
|
+
"""Create MCP configuration for stdio transport.
|
|
83
|
+
|
|
84
|
+
Args:
|
|
85
|
+
docker_cmd: Docker command to run
|
|
86
|
+
|
|
87
|
+
Returns:
|
|
88
|
+
MCP configuration dict
|
|
89
|
+
"""
|
|
90
|
+
return {
|
|
91
|
+
"mcpServers": {
|
|
92
|
+
"default": {
|
|
93
|
+
"command": docker_cmd[0],
|
|
94
|
+
"args": docker_cmd[1:] if len(docker_cmd) > 1 else [],
|
|
95
|
+
# transport defaults to stdio
|
|
96
|
+
}
|
|
97
|
+
}
|
|
98
|
+
}
|
|
99
|
+
|
|
100
|
+
def create_proxy(self, config: dict[str, Any], name: str | None = None) -> FastMCP:
|
|
101
|
+
"""Create FastMCP proxy server.
|
|
102
|
+
|
|
103
|
+
Args:
|
|
104
|
+
config: MCP configuration
|
|
105
|
+
name: Optional server name
|
|
106
|
+
|
|
107
|
+
Returns:
|
|
108
|
+
FastMCP proxy instance
|
|
109
|
+
"""
|
|
110
|
+
proxy_name = name or f"HUD Server - {self.image}"
|
|
111
|
+
return FastMCP.as_proxy(config, name=proxy_name)
|
|
112
|
+
|
|
113
|
+
async def run_http_server(
|
|
114
|
+
self,
|
|
115
|
+
proxy: FastMCP,
|
|
116
|
+
port: int,
|
|
117
|
+
verbose: bool = False,
|
|
118
|
+
path: str = "/mcp",
|
|
119
|
+
) -> None:
|
|
120
|
+
"""Run HTTP server with proper shutdown handling.
|
|
121
|
+
|
|
122
|
+
Args:
|
|
123
|
+
proxy: FastMCP proxy instance
|
|
124
|
+
port: Port to listen on
|
|
125
|
+
verbose: Enable verbose logging
|
|
126
|
+
path: URL path for MCP endpoint
|
|
127
|
+
"""
|
|
128
|
+
# Set up logging
|
|
129
|
+
import logging
|
|
130
|
+
import os
|
|
131
|
+
|
|
132
|
+
os.environ["FASTMCP_DISABLE_BANNER"] = "1"
|
|
133
|
+
|
|
134
|
+
if not verbose:
|
|
135
|
+
logging.getLogger("fastmcp").setLevel(logging.ERROR)
|
|
136
|
+
logging.getLogger("mcp").setLevel(logging.ERROR)
|
|
137
|
+
logging.getLogger("uvicorn").setLevel(logging.ERROR)
|
|
138
|
+
logging.getLogger("uvicorn.access").setLevel(logging.ERROR)
|
|
139
|
+
logging.getLogger("uvicorn.error").setLevel(logging.ERROR)
|
|
140
|
+
|
|
141
|
+
import warnings
|
|
142
|
+
|
|
143
|
+
warnings.filterwarnings("ignore", category=DeprecationWarning)
|
|
144
|
+
|
|
145
|
+
try:
|
|
146
|
+
await proxy.run_async(
|
|
147
|
+
transport="http",
|
|
148
|
+
host="0.0.0.0", # noqa: S104
|
|
149
|
+
port=port,
|
|
150
|
+
path=path,
|
|
151
|
+
log_level="error" if not verbose else "info",
|
|
152
|
+
show_banner=False,
|
|
153
|
+
)
|
|
154
|
+
except asyncio.CancelledError:
|
|
155
|
+
pass # Normal cancellation
|
|
156
|
+
except Exception as e:
|
|
157
|
+
if verbose:
|
|
158
|
+
self.design.error(f"Server error: {e}")
|
|
159
|
+
raise
|
|
160
|
+
|
|
161
|
+
|
|
162
|
+
async def run_server_with_interactive(
|
|
163
|
+
server_manager: MCPServerManager,
|
|
164
|
+
port: int,
|
|
165
|
+
verbose: bool = False,
|
|
166
|
+
) -> None:
|
|
167
|
+
"""Run server with interactive testing mode.
|
|
168
|
+
|
|
169
|
+
Args:
|
|
170
|
+
server_manager: Server manager instance
|
|
171
|
+
port: Port to listen on
|
|
172
|
+
verbose: Enable verbose logging
|
|
173
|
+
"""
|
|
174
|
+
from .interactive import run_interactive_mode
|
|
175
|
+
from .logging import find_free_port
|
|
176
|
+
|
|
177
|
+
design = HUDDesign()
|
|
178
|
+
|
|
179
|
+
# Find available port
|
|
180
|
+
actual_port = find_free_port(port)
|
|
181
|
+
if actual_port is None:
|
|
182
|
+
design.error(f"No available ports found starting from {port}")
|
|
183
|
+
return
|
|
184
|
+
|
|
185
|
+
if actual_port != port:
|
|
186
|
+
design.warning(f"Port {port} in use, using port {actual_port} instead")
|
|
187
|
+
|
|
188
|
+
# Clean up any existing container
|
|
189
|
+
server_manager.cleanup_container()
|
|
190
|
+
|
|
191
|
+
# Build docker command
|
|
192
|
+
docker_cmd = server_manager.build_docker_command()
|
|
193
|
+
|
|
194
|
+
# Create MCP config
|
|
195
|
+
config = server_manager.create_mcp_config(docker_cmd)
|
|
196
|
+
|
|
197
|
+
# Create proxy
|
|
198
|
+
proxy = server_manager.create_proxy(config, f"HUD Interactive - {server_manager.image}")
|
|
199
|
+
|
|
200
|
+
# Show header
|
|
201
|
+
design.info("") # Empty line
|
|
202
|
+
design.header("HUD MCP Server - Interactive Mode", icon="🎮")
|
|
203
|
+
|
|
204
|
+
# Show configuration
|
|
205
|
+
design.section_title("Server Information")
|
|
206
|
+
design.info(f"Image: {server_manager.image}")
|
|
207
|
+
design.info(f"Port: {actual_port}")
|
|
208
|
+
design.info(f"URL: http://localhost:{actual_port}/mcp")
|
|
209
|
+
design.info(f"Container: {server_manager.container_name}")
|
|
210
|
+
design.info("")
|
|
211
|
+
|
|
212
|
+
# Create event to signal server is ready
|
|
213
|
+
server_ready = asyncio.Event()
|
|
214
|
+
server_task = None
|
|
215
|
+
|
|
216
|
+
async def start_server() -> None:
|
|
217
|
+
"""Start the proxy server."""
|
|
218
|
+
nonlocal server_task
|
|
219
|
+
try:
|
|
220
|
+
# Signal that we're ready before starting
|
|
221
|
+
server_ready.set()
|
|
222
|
+
await server_manager.run_http_server(proxy, actual_port, verbose)
|
|
223
|
+
except asyncio.CancelledError:
|
|
224
|
+
pass
|
|
225
|
+
|
|
226
|
+
try:
|
|
227
|
+
# Start server in background
|
|
228
|
+
server_task = asyncio.create_task(start_server())
|
|
229
|
+
|
|
230
|
+
# Wait for server to be ready
|
|
231
|
+
await server_ready.wait()
|
|
232
|
+
await asyncio.sleep(0.5) # Give it a moment to fully start
|
|
233
|
+
|
|
234
|
+
# Run interactive mode
|
|
235
|
+
server_url = f"http://localhost:{actual_port}/mcp"
|
|
236
|
+
await run_interactive_mode(server_url, verbose=verbose)
|
|
237
|
+
|
|
238
|
+
except KeyboardInterrupt:
|
|
239
|
+
design.info("\n👋 Shutting down...")
|
|
240
|
+
finally:
|
|
241
|
+
# Cancel server task
|
|
242
|
+
if server_task and not server_task.done():
|
|
243
|
+
server_task.cancel()
|
|
244
|
+
try:
|
|
245
|
+
await server_task
|
|
246
|
+
except asyncio.CancelledError:
|
|
247
|
+
design.error("Server task cancelled")
|
|
248
|
+
|
|
249
|
+
# Clean up container
|
|
250
|
+
server_manager.cleanup_container()
|
hud/clients/base.py
CHANGED
hud/clients/fastmcp.py
CHANGED
|
@@ -94,25 +94,27 @@ class FastMCPHUDClient(BaseHUDClient):
|
|
|
94
94
|
if "mcp.hud.so" in url:
|
|
95
95
|
raise RuntimeError(
|
|
96
96
|
"Authentication failed for HUD API. "
|
|
97
|
-
"Please ensure your HUD_API_KEY environment variable is set correctly.
|
|
97
|
+
"Please ensure your HUD_API_KEY environment variable is set correctly." # noqa: E501
|
|
98
98
|
"You can get an API key at https://app.hud.so"
|
|
99
99
|
) from e
|
|
100
100
|
# Generic 401 error
|
|
101
101
|
raise RuntimeError(
|
|
102
|
-
|
|
103
|
-
|
|
102
|
+
"Authentication failed (401 Unauthorized). "
|
|
103
|
+
"Please check your credentials or API key."
|
|
104
104
|
) from e
|
|
105
105
|
raise
|
|
106
106
|
|
|
107
107
|
# Configure validation for output schemas based on client setting
|
|
108
108
|
try:
|
|
109
|
-
from hud_mcp.client.session import
|
|
109
|
+
from hud_mcp.client.session import ( # type: ignore[import-not-found]
|
|
110
|
+
ValidationOptions, # type: ignore[import-not-found]
|
|
111
|
+
)
|
|
110
112
|
|
|
111
113
|
if (
|
|
112
114
|
hasattr(self._client, "_session_state")
|
|
113
115
|
and self._client._session_state.session is not None
|
|
114
116
|
):
|
|
115
|
-
self._client._session_state.session._validation_options = ValidationOptions(
|
|
117
|
+
self._client._session_state.session._validation_options = ValidationOptions( # type: ignore[attr-defined]
|
|
116
118
|
strict_output_validation=self._strict_validation
|
|
117
119
|
)
|
|
118
120
|
except ImportError:
|
hud/clients/mcp_use.py
CHANGED
|
@@ -16,12 +16,12 @@ from .base import BaseHUDClient
|
|
|
16
16
|
|
|
17
17
|
if TYPE_CHECKING:
|
|
18
18
|
from mcp import types
|
|
19
|
-
from mcp_use.client import MCPClient as MCPUseClient
|
|
20
|
-
from mcp_use.session import MCPSession as MCPUseSession
|
|
19
|
+
from mcp_use.client import MCPClient as MCPUseClient # type: ignore[attr-defined]
|
|
20
|
+
from mcp_use.session import MCPSession as MCPUseSession # type: ignore[attr-defined]
|
|
21
21
|
|
|
22
22
|
try:
|
|
23
|
-
from mcp_use.client import MCPClient as MCPUseClient
|
|
24
|
-
from mcp_use.session import MCPSession as MCPUseSession
|
|
23
|
+
from mcp_use.client import MCPClient as MCPUseClient # type: ignore[attr-defined]
|
|
24
|
+
from mcp_use.session import MCPSession as MCPUseSession # type: ignore[attr-defined]
|
|
25
25
|
except ImportError:
|
|
26
26
|
MCPUseClient = None # type: ignore[misc, assignment]
|
|
27
27
|
MCPUseSession = None # type: ignore[misc, assignment]
|
|
@@ -67,13 +67,15 @@ class MCPUseHUDClient(BaseHUDClient):
|
|
|
67
67
|
raise ImportError("MCPUseClient is not available")
|
|
68
68
|
self._client = MCPUseClient.from_dict(config)
|
|
69
69
|
try:
|
|
70
|
-
assert self._client is not None #
|
|
70
|
+
assert self._client is not None # noqa: S101
|
|
71
71
|
self._sessions = await self._client.create_all_sessions()
|
|
72
72
|
logger.info("Created %d MCP sessions", len(self._sessions))
|
|
73
73
|
|
|
74
74
|
# Configure validation for all sessions based on client setting
|
|
75
75
|
try:
|
|
76
|
-
from hud_mcp.client.session import
|
|
76
|
+
from hud_mcp.client.session import ( # type: ignore[import-not-found]
|
|
77
|
+
ValidationOptions, # type: ignore[import-not-found]
|
|
78
|
+
)
|
|
77
79
|
|
|
78
80
|
for session in self._sessions.values():
|
|
79
81
|
if (
|
hud/server/server.py
CHANGED
|
@@ -3,11 +3,11 @@
|
|
|
3
3
|
from __future__ import annotations
|
|
4
4
|
|
|
5
5
|
import asyncio
|
|
6
|
+
import contextlib
|
|
6
7
|
import logging
|
|
7
8
|
import os
|
|
8
9
|
import signal
|
|
9
10
|
import sys
|
|
10
|
-
import contextlib
|
|
11
11
|
from contextlib import asynccontextmanager
|
|
12
12
|
from typing import TYPE_CHECKING, Any
|
|
13
13
|
|
|
@@ -25,9 +25,13 @@ __all__ = ["MCPServer"]
|
|
|
25
25
|
|
|
26
26
|
logger = logging.getLogger(__name__)
|
|
27
27
|
|
|
28
|
+
# Global flag to track if shutdown was triggered by SIGTERM
|
|
29
|
+
_sigterm_received = False
|
|
30
|
+
|
|
28
31
|
|
|
29
32
|
def _run_with_sigterm(coro_fn: Callable[..., Any], *args: Any, **kwargs: Any) -> None:
|
|
30
33
|
"""Run *coro_fn* via anyio.run() and cancel on SIGTERM or SIGINT (POSIX)."""
|
|
34
|
+
global _sigterm_received
|
|
31
35
|
|
|
32
36
|
async def _runner() -> None:
|
|
33
37
|
stop_evt: asyncio.Event | None = None
|
|
@@ -35,9 +39,16 @@ def _run_with_sigterm(coro_fn: Callable[..., Any], *args: Any, **kwargs: Any) ->
|
|
|
35
39
|
loop = asyncio.get_running_loop()
|
|
36
40
|
stop_evt = asyncio.Event()
|
|
37
41
|
|
|
42
|
+
# Handle SIGTERM for production shutdown
|
|
43
|
+
def handle_sigterm() -> None:
|
|
44
|
+
global _sigterm_received
|
|
45
|
+
_sigterm_received = True
|
|
46
|
+
logger.info("Received SIGTERM signal")
|
|
47
|
+
stop_evt.set()
|
|
48
|
+
|
|
38
49
|
# Handle both SIGTERM and SIGINT for graceful shutdown
|
|
39
50
|
if signal.getsignal(signal.SIGTERM) is signal.SIG_DFL:
|
|
40
|
-
loop.add_signal_handler(signal.SIGTERM,
|
|
51
|
+
loop.add_signal_handler(signal.SIGTERM, handle_sigterm)
|
|
41
52
|
if signal.getsignal(signal.SIGINT) is signal.SIG_DFL:
|
|
42
53
|
loop.add_signal_handler(signal.SIGINT, stop_evt.set)
|
|
43
54
|
|
|
@@ -62,6 +73,7 @@ class MCPServer(FastMCP):
|
|
|
62
73
|
"""FastMCP wrapper that adds helpful functionality for dockerized environments.
|
|
63
74
|
This works with any MCP client, and adds just a few extra server-side features:
|
|
64
75
|
1. SIGTERM handling for graceful shutdown in container runtimes.
|
|
76
|
+
Note: SIGINT (Ctrl+C) is not handled, allowing normal hot reload behavior.
|
|
65
77
|
2. ``@MCPServer.initialize`` decorator that registers an async initializer
|
|
66
78
|
executed during the MCP *initialize* request. The initializer function receives
|
|
67
79
|
a single ``ctx`` parameter (RequestContext) from which you can access:
|
|
@@ -69,7 +81,7 @@ class MCPServer(FastMCP):
|
|
|
69
81
|
- ``ctx.meta.progressToken``: Token for progress notifications (if provided)
|
|
70
82
|
- ``ctx.session.client_params.clientInfo``: Client information
|
|
71
83
|
3. ``@MCPServer.shutdown`` decorator that registers a coroutine to run during
|
|
72
|
-
server teardown
|
|
84
|
+
server teardown ONLY when SIGTERM is received (not on hot reload/SIGINT).
|
|
73
85
|
4. Enhanced ``add_tool`` that accepts instances of
|
|
74
86
|
:class:`hud.tools.base.BaseTool` which are classes that implement the
|
|
75
87
|
FastMCP ``FunctionTool`` interface.
|
|
@@ -84,11 +96,17 @@ class MCPServer(FastMCP):
|
|
|
84
96
|
|
|
85
97
|
@asynccontextmanager
|
|
86
98
|
async def _lifespan(_: Any) -> AsyncGenerator[dict[str, Any], None]:
|
|
99
|
+
global _sigterm_received
|
|
87
100
|
try:
|
|
88
101
|
yield {}
|
|
89
102
|
finally:
|
|
90
|
-
if
|
|
103
|
+
# Only call shutdown handler if SIGTERM was received
|
|
104
|
+
if self._shutdown_fn is not None and _sigterm_received:
|
|
105
|
+
logger.info("SIGTERM received, calling shutdown handler")
|
|
91
106
|
await self._shutdown_fn()
|
|
107
|
+
_sigterm_received = False
|
|
108
|
+
elif self._shutdown_fn is not None:
|
|
109
|
+
logger.debug("Normal shutdown (hot reload), skipping shutdown handler")
|
|
92
110
|
|
|
93
111
|
fastmcp_kwargs["lifespan"] = _lifespan
|
|
94
112
|
|
|
@@ -138,6 +156,18 @@ class MCPServer(FastMCP):
|
|
|
138
156
|
# Shutdown decorator: runs after server stops
|
|
139
157
|
# Supports dockerized SIGTERM handling
|
|
140
158
|
def shutdown(self, fn: Callable | None = None) -> Callable | None:
|
|
159
|
+
"""Register a shutdown handler that runs ONLY on SIGTERM.
|
|
160
|
+
|
|
161
|
+
This handler will be called when the server receives a SIGTERM signal
|
|
162
|
+
(e.g., during container shutdown). It will NOT be called on:
|
|
163
|
+
- SIGINT (Ctrl+C or hot reload)
|
|
164
|
+
- Normal client disconnects
|
|
165
|
+
- Other graceful shutdowns
|
|
166
|
+
|
|
167
|
+
This ensures that persistent resources (like browser sessions) are only
|
|
168
|
+
cleaned up during actual termination, not during development hot reloads.
|
|
169
|
+
"""
|
|
170
|
+
|
|
141
171
|
def decorator(func: Callable) -> Callable:
|
|
142
172
|
self._shutdown_fn = func
|
|
143
173
|
return func
|
hud/shared/exceptions.py
CHANGED
|
@@ -18,6 +18,17 @@ class HudException(Exception):
|
|
|
18
18
|
Consumers should be able to catch this exception to handle any HUD-related error.
|
|
19
19
|
"""
|
|
20
20
|
|
|
21
|
+
def __init__(self, message: str, response_json: dict[str, Any] | None = None) -> None:
|
|
22
|
+
super().__init__(message)
|
|
23
|
+
self.message = message
|
|
24
|
+
self.response_json = response_json
|
|
25
|
+
|
|
26
|
+
def __str__(self) -> str:
|
|
27
|
+
parts = [self.message]
|
|
28
|
+
if self.response_json:
|
|
29
|
+
parts.append(f"Response: {self.response_json}")
|
|
30
|
+
return " | ".join(parts)
|
|
31
|
+
|
|
21
32
|
|
|
22
33
|
class HudRequestError(Exception):
|
|
23
34
|
"""Any request to the HUD API can raise this exception."""
|
|
@@ -155,3 +155,25 @@ class TestGymMakeException:
|
|
|
155
155
|
assert "env_id" in error_str
|
|
156
156
|
assert "test-env" in error_str
|
|
157
157
|
assert "invalid config" in error_str
|
|
158
|
+
|
|
159
|
+
|
|
160
|
+
class TestHudException:
|
|
161
|
+
"""Test base HudException class."""
|
|
162
|
+
|
|
163
|
+
def test_str_with_response_json(self):
|
|
164
|
+
"""Test HudException string representation with response_json."""
|
|
165
|
+
response_data = {"error": "test error", "code": 42}
|
|
166
|
+
error = HudException("Base error message", response_data)
|
|
167
|
+
|
|
168
|
+
error_str = str(error)
|
|
169
|
+
assert "Base error message" in error_str
|
|
170
|
+
assert "error" in error_str
|
|
171
|
+
assert "test error" in error_str
|
|
172
|
+
|
|
173
|
+
def test_str_without_response_json(self):
|
|
174
|
+
"""Test HudException string representation without response_json."""
|
|
175
|
+
error = HudException("Just a message")
|
|
176
|
+
|
|
177
|
+
error_str = str(error)
|
|
178
|
+
assert error_str == "Just a message"
|
|
179
|
+
assert "Response:" not in error_str
|
|
File without changes
|
|
@@ -0,0 +1,40 @@
|
|
|
1
|
+
"""Tests for telemetry replay functionality."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
from unittest.mock import patch
|
|
6
|
+
|
|
7
|
+
from hud.telemetry.replay import clear_trace, get_trace
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
class TestReplayAPI:
|
|
11
|
+
"""Tests for replay API functions."""
|
|
12
|
+
|
|
13
|
+
def test_get_trace_calls_internal(self):
|
|
14
|
+
"""Test that get_trace calls the internal _get_trace function."""
|
|
15
|
+
with patch("hud.telemetry.replay._get_trace") as mock_get:
|
|
16
|
+
mock_get.return_value = None
|
|
17
|
+
|
|
18
|
+
result = get_trace("test-task-id")
|
|
19
|
+
|
|
20
|
+
mock_get.assert_called_once_with("test-task-id")
|
|
21
|
+
assert result is None
|
|
22
|
+
|
|
23
|
+
def test_clear_trace_calls_internal(self):
|
|
24
|
+
"""Test that clear_trace calls the internal _clear_trace function."""
|
|
25
|
+
with patch("hud.telemetry.replay._clear_trace") as mock_clear:
|
|
26
|
+
clear_trace("test-task-id")
|
|
27
|
+
|
|
28
|
+
mock_clear.assert_called_once_with("test-task-id")
|
|
29
|
+
|
|
30
|
+
def test_get_trace_with_data(self):
|
|
31
|
+
"""Test get_trace with mock data."""
|
|
32
|
+
mock_trace = {"trace": [{"step": 1}], "task_run_id": "test-123"}
|
|
33
|
+
|
|
34
|
+
with patch("hud.telemetry.replay._get_trace") as mock_get:
|
|
35
|
+
mock_get.return_value = mock_trace
|
|
36
|
+
|
|
37
|
+
result = get_trace("test-123")
|
|
38
|
+
|
|
39
|
+
assert result == mock_trace
|
|
40
|
+
mock_get.assert_called_once_with("test-123")
|