feat(hub): add hub integration and rename project to FuzzForge AI

This commit is contained in:
AFredefon
2026-02-25 23:12:42 +01:00
parent 04c8383739
commit f3899279d5
33 changed files with 1904 additions and 73 deletions
@@ -0,0 +1,42 @@
"""FuzzForge Hub - Generic MCP server bridge.
This module provides a generic bridge to connect FuzzForge with any MCP server.
It allows AI agents to discover and execute tools from external MCP servers
(like mcp-security-hub) through the same interface as native FuzzForge modules.
The hub is server-agnostic: it doesn't hardcode any specific tools or servers.
Instead, it dynamically discovers tools by connecting to configured MCP servers
and calling their `list_tools()` method.
Supported transport types:
- docker: Run MCP server as a Docker container with stdio transport
- command: Run MCP server as a local process with stdio transport
- sse: Connect to a remote MCP server via Server-Sent Events
"""
from fuzzforge_common.hub.client import HubClient, HubClientError
from fuzzforge_common.hub.executor import HubExecutionResult, HubExecutor
from fuzzforge_common.hub.models import (
HubConfig,
HubServer,
HubServerConfig,
HubServerType,
HubTool,
HubToolParameter,
)
from fuzzforge_common.hub.registry import HubRegistry
__all__ = [
"HubClient",
"HubClientError",
"HubConfig",
"HubExecutionResult",
"HubExecutor",
"HubRegistry",
"HubServer",
"HubServerConfig",
"HubServerType",
"HubTool",
"HubToolParameter",
]
@@ -0,0 +1,443 @@
"""MCP client for communicating with hub servers.
This module provides a generic MCP client that can connect to any MCP server
via stdio (docker/command) or SSE transport. It handles:
- Starting containers/processes for stdio transport
- Connecting to SSE endpoints
- Discovering tools via list_tools()
- Executing tools via call_tool()
"""
from __future__ import annotations
import asyncio
import json
import subprocess
from contextlib import asynccontextmanager
from typing import TYPE_CHECKING, Any, cast
from fuzzforge_common.hub.models import (
HubServer,
HubServerConfig,
HubServerType,
HubTool,
)
if TYPE_CHECKING:
from asyncio.subprocess import Process
from collections.abc import AsyncGenerator
from structlog.stdlib import BoundLogger
def get_logger() -> BoundLogger:
"""Get structlog logger instance.
:returns: Configured structlog logger.
"""
from structlog import get_logger # noqa: PLC0415
return cast("BoundLogger", get_logger())
class HubClientError(Exception):
"""Error in hub client operations."""
class HubClient:
"""Client for communicating with MCP hub servers.
Supports stdio (via docker/command) and SSE transports.
Uses the MCP protocol for tool discovery and execution.
"""
#: Default timeout for operations.
DEFAULT_TIMEOUT: int = 30
def __init__(self, timeout: int = DEFAULT_TIMEOUT) -> None:
"""Initialize the hub client.
:param timeout: Default timeout for operations in seconds.
"""
self._timeout = timeout
async def discover_tools(self, server: HubServer) -> list[HubTool]:
"""Discover tools from a hub server.
Connects to the server, calls list_tools(), and returns
parsed HubTool instances.
:param server: Hub server to discover tools from.
:returns: List of discovered tools.
:raises HubClientError: If discovery fails.
"""
logger = get_logger()
config = server.config
logger.info("Discovering tools", server=config.name, type=config.type.value)
try:
async with self._connect(config) as (reader, writer):
# Initialize MCP session
await self._initialize_session(reader, writer, config.name)
# List tools
tools_data = await self._call_method(
reader,
writer,
"tools/list",
{},
)
# Parse tools
tools = []
for tool_data in tools_data.get("tools", []):
tool = HubTool.from_mcp_tool(
server_name=config.name,
name=tool_data["name"],
description=tool_data.get("description"),
input_schema=tool_data.get("inputSchema", {}),
)
tools.append(tool)
logger.info(
"Discovered tools",
server=config.name,
count=len(tools),
)
return tools
except Exception as e:
logger.error(
"Tool discovery failed",
server=config.name,
error=str(e),
)
raise HubClientError(f"Discovery failed for {config.name}: {e}") from e
async def execute_tool(
self,
server: HubServer,
tool_name: str,
arguments: dict[str, Any],
*,
timeout: int | None = None,
) -> dict[str, Any]:
"""Execute a tool on a hub server.
:param server: Hub server to execute on.
:param tool_name: Name of the tool to execute.
:param arguments: Tool arguments.
:param timeout: Execution timeout (uses default if None).
:returns: Tool execution result.
:raises HubClientError: If execution fails.
"""
logger = get_logger()
config = server.config
exec_timeout = timeout or self._timeout
logger.info(
"Executing hub tool",
server=config.name,
tool=tool_name,
timeout=exec_timeout,
)
try:
async with self._connect(config) as (reader, writer):
# Initialize MCP session
await self._initialize_session(reader, writer, config.name)
# Call tool
result = await asyncio.wait_for(
self._call_method(
reader,
writer,
"tools/call",
{"name": tool_name, "arguments": arguments},
),
timeout=exec_timeout,
)
logger.info(
"Tool execution completed",
server=config.name,
tool=tool_name,
)
return result
except asyncio.TimeoutError as e:
logger.error(
"Tool execution timed out",
server=config.name,
tool=tool_name,
timeout=exec_timeout,
)
raise HubClientError(
f"Execution timed out for {config.name}:{tool_name}"
) from e
except Exception as e:
logger.error(
"Tool execution failed",
server=config.name,
tool=tool_name,
error=str(e),
)
raise HubClientError(
f"Execution failed for {config.name}:{tool_name}: {e}"
) from e
@asynccontextmanager
async def _connect(
self,
config: HubServerConfig,
) -> AsyncGenerator[tuple[asyncio.StreamReader, asyncio.StreamWriter], None]:
"""Connect to an MCP server.
:param config: Server configuration.
:yields: Tuple of (reader, writer) for communication.
"""
if config.type == HubServerType.DOCKER:
async with self._connect_docker(config) as streams:
yield streams
elif config.type == HubServerType.COMMAND:
async with self._connect_command(config) as streams:
yield streams
elif config.type == HubServerType.SSE:
async with self._connect_sse(config) as streams:
yield streams
else:
msg = f"Unsupported server type: {config.type}"
raise HubClientError(msg)
@asynccontextmanager
async def _connect_docker(
self,
config: HubServerConfig,
) -> AsyncGenerator[tuple[asyncio.StreamReader, asyncio.StreamWriter], None]:
"""Connect to a Docker-based MCP server.
:param config: Server configuration with image name.
:yields: Tuple of (reader, writer) for stdio communication.
"""
if not config.image:
msg = f"Docker image not specified for server '{config.name}'"
raise HubClientError(msg)
# Build docker command
cmd = ["docker", "run", "-i", "--rm"]
# Add capabilities
for cap in config.capabilities:
cmd.extend(["--cap-add", cap])
# Add volumes
for volume in config.volumes:
cmd.extend(["-v", volume])
# Add environment variables
for key, value in config.environment.items():
cmd.extend(["-e", f"{key}={value}"])
cmd.append(config.image)
process: Process = await asyncio.create_subprocess_exec(
*cmd,
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
)
try:
if process.stdin is None or process.stdout is None:
msg = "Failed to get process streams"
raise HubClientError(msg)
# Create asyncio streams from process pipes
reader = process.stdout
writer = process.stdin
yield reader, writer # type: ignore[misc]
finally:
process.terminate()
try:
await asyncio.wait_for(process.wait(), timeout=5)
except asyncio.TimeoutError:
process.kill()
@asynccontextmanager
async def _connect_command(
self,
config: HubServerConfig,
) -> AsyncGenerator[tuple[asyncio.StreamReader, asyncio.StreamWriter], None]:
"""Connect to a command-based MCP server.
:param config: Server configuration with command.
:yields: Tuple of (reader, writer) for stdio communication.
"""
if not config.command:
msg = f"Command not specified for server '{config.name}'"
raise HubClientError(msg)
# Set up environment
env = dict(config.environment) if config.environment else None
process: Process = await asyncio.create_subprocess_exec(
*config.command,
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
env=env,
)
try:
if process.stdin is None or process.stdout is None:
msg = "Failed to get process streams"
raise HubClientError(msg)
reader = process.stdout
writer = process.stdin
yield reader, writer # type: ignore[misc]
finally:
process.terminate()
try:
await asyncio.wait_for(process.wait(), timeout=5)
except asyncio.TimeoutError:
process.kill()
@asynccontextmanager
async def _connect_sse(
self,
config: HubServerConfig,
) -> AsyncGenerator[tuple[asyncio.StreamReader, asyncio.StreamWriter], None]:
"""Connect to an SSE-based MCP server.
:param config: Server configuration with URL.
:yields: Tuple of (reader, writer) for SSE communication.
"""
# SSE support requires additional dependencies
# For now, raise not implemented
msg = "SSE transport not yet implemented"
raise NotImplementedError(msg)
async def _initialize_session(
self,
reader: asyncio.StreamReader,
writer: asyncio.StreamWriter,
server_name: str,
) -> dict[str, Any]:
"""Initialize MCP session with the server.
:param reader: Stream reader.
:param writer: Stream writer.
:param server_name: Server name for logging.
:returns: Server capabilities.
"""
# Send initialize request
result = await self._call_method(
reader,
writer,
"initialize",
{
"protocolVersion": "2024-11-05",
"capabilities": {},
"clientInfo": {
"name": "fuzzforge-hub",
"version": "0.1.0",
},
},
)
# Send initialized notification
await self._send_notification(reader, writer, "notifications/initialized", {})
return result
async def _call_method(
self,
reader: asyncio.StreamReader,
writer: asyncio.StreamWriter,
method: str,
params: dict[str, Any],
) -> dict[str, Any]:
"""Call an MCP method.
:param reader: Stream reader.
:param writer: Stream writer.
:param method: Method name.
:param params: Method parameters.
:returns: Method result.
"""
# Create JSON-RPC request
request = {
"jsonrpc": "2.0",
"id": 1,
"method": method,
"params": params,
}
# Send request
request_line = json.dumps(request) + "\n"
writer.write(request_line.encode())
await writer.drain()
# Read response
response_line = await asyncio.wait_for(
reader.readline(),
timeout=self._timeout,
)
if not response_line:
msg = "Empty response from server"
raise HubClientError(msg)
response = json.loads(response_line.decode())
if "error" in response:
error = response["error"]
msg = f"MCP error: {error.get('message', 'Unknown error')}"
raise HubClientError(msg)
return response.get("result", {})
async def _send_notification(
self,
reader: asyncio.StreamReader,
writer: asyncio.StreamWriter,
method: str,
params: dict[str, Any],
) -> None:
"""Send an MCP notification (no response expected).
:param reader: Stream reader (unused but kept for consistency).
:param writer: Stream writer.
:param method: Notification method name.
:param params: Notification parameters.
"""
# Create JSON-RPC notification (no id)
notification = {
"jsonrpc": "2.0",
"method": method,
"params": params,
}
notification_line = json.dumps(notification) + "\n"
writer.write(notification_line.encode())
await writer.drain()
@@ -0,0 +1,334 @@
"""Hub executor for managing MCP server lifecycle and tool execution.
This module provides a high-level interface for:
- Discovering tools from all registered hub servers
- Executing tools with proper error handling
- Managing the lifecycle of hub operations
"""
from __future__ import annotations
from pathlib import Path
from typing import TYPE_CHECKING, Any, cast
from fuzzforge_common.hub.client import HubClient, HubClientError
from fuzzforge_common.hub.models import HubServer, HubServerConfig, HubTool
from fuzzforge_common.hub.registry import HubRegistry
if TYPE_CHECKING:
from structlog.stdlib import BoundLogger
def get_logger() -> BoundLogger:
"""Get structlog logger instance.
:returns: Configured structlog logger.
"""
from structlog import get_logger # noqa: PLC0415
return cast("BoundLogger", get_logger())
class HubExecutionResult:
"""Result of a hub tool execution."""
def __init__(
self,
*,
success: bool,
server_name: str,
tool_name: str,
result: dict[str, Any] | None = None,
error: str | None = None,
) -> None:
"""Initialize execution result.
:param success: Whether execution succeeded.
:param server_name: Name of the hub server.
:param tool_name: Name of the executed tool.
:param result: Tool execution result data.
:param error: Error message if execution failed.
"""
self.success = success
self.server_name = server_name
self.tool_name = tool_name
self.result = result or {}
self.error = error
@property
def identifier(self) -> str:
"""Get full tool identifier."""
return f"hub:{self.server_name}:{self.tool_name}"
def to_dict(self) -> dict[str, Any]:
"""Convert to dictionary.
:returns: Dictionary representation.
"""
return {
"success": self.success,
"identifier": self.identifier,
"server": self.server_name,
"tool": self.tool_name,
"result": self.result,
"error": self.error,
}
class HubExecutor:
"""Executor for hub server operations.
Provides high-level methods for discovering and executing
tools from hub servers.
"""
#: Hub registry instance.
_registry: HubRegistry
#: MCP client instance.
_client: HubClient
def __init__(
self,
config_path: Path | None = None,
timeout: int = 300,
) -> None:
"""Initialize the hub executor.
:param config_path: Path to hub-servers.json config file.
:param timeout: Default timeout for tool execution.
"""
self._registry = HubRegistry(config_path)
self._client = HubClient(timeout=timeout)
@property
def registry(self) -> HubRegistry:
"""Get the hub registry.
:returns: Hub registry instance.
"""
return self._registry
def add_server(self, config: HubServerConfig) -> HubServer:
"""Add a server to the registry.
:param config: Server configuration.
:returns: Created HubServer instance.
"""
return self._registry.add_server(config)
async def discover_all_tools(self) -> dict[str, list[HubTool]]:
"""Discover tools from all enabled servers.
:returns: Dict mapping server names to lists of discovered tools.
"""
logger = get_logger()
results: dict[str, list[HubTool]] = {}
for server in self._registry.enabled_servers:
try:
tools = await self._client.discover_tools(server)
self._registry.update_server_tools(server.name, tools)
results[server.name] = tools
except HubClientError as e:
logger.warning(
"Failed to discover tools",
server=server.name,
error=str(e),
)
self._registry.update_server_tools(server.name, [], error=str(e))
results[server.name] = []
return results
async def discover_server_tools(self, server_name: str) -> list[HubTool]:
"""Discover tools from a specific server.
:param server_name: Name of the server.
:returns: List of discovered tools.
:raises ValueError: If server not found.
"""
server = self._registry.get_server(server_name)
if not server:
msg = f"Server '{server_name}' not found"
raise ValueError(msg)
try:
tools = await self._client.discover_tools(server)
self._registry.update_server_tools(server_name, tools)
return tools
except HubClientError as e:
self._registry.update_server_tools(server_name, [], error=str(e))
raise
async def execute_tool(
self,
identifier: str,
arguments: dict[str, Any] | None = None,
*,
timeout: int | None = None,
) -> HubExecutionResult:
"""Execute a hub tool.
:param identifier: Tool identifier (hub:server:tool or server:tool).
:param arguments: Tool arguments.
:param timeout: Execution timeout.
:returns: Execution result.
"""
logger = get_logger()
arguments = arguments or {}
# Parse identifier and find tool
server, tool = self._registry.find_tool(identifier)
if not server or not tool:
# Try to parse as server:tool and discover
parts = identifier.replace("hub:", "").split(":")
if len(parts) == 2: # noqa: PLR2004
server_name, tool_name = parts
server = self._registry.get_server(server_name)
if server and not server.discovered:
# Try to discover tools first
try:
await self.discover_server_tools(server_name)
tool = server.get_tool(tool_name)
except HubClientError:
pass
if server and not tool:
# Tool not found, but server exists - try to execute anyway
# The server might have the tool even if discovery failed
tool_name_to_use = tool_name
else:
tool_name_to_use = tool.name if tool else ""
if not server:
return HubExecutionResult(
success=False,
server_name=server_name,
tool_name=tool_name,
error=f"Server '{server_name}' not found",
)
# Execute even if tool wasn't discovered (server might still have it)
try:
result = await self._client.execute_tool(
server,
tool_name_to_use or tool_name,
arguments,
timeout=timeout,
)
return HubExecutionResult(
success=True,
server_name=server.name,
tool_name=tool_name_to_use or tool_name,
result=result,
)
except HubClientError as e:
return HubExecutionResult(
success=False,
server_name=server.name,
tool_name=tool_name_to_use or tool_name,
error=str(e),
)
else:
return HubExecutionResult(
success=False,
server_name="unknown",
tool_name=identifier,
error=f"Invalid tool identifier: {identifier}",
)
# Execute the tool
logger.info(
"Executing hub tool",
server=server.name,
tool=tool.name,
arguments=arguments,
)
try:
result = await self._client.execute_tool(
server,
tool.name,
arguments,
timeout=timeout,
)
return HubExecutionResult(
success=True,
server_name=server.name,
tool_name=tool.name,
result=result,
)
except HubClientError as e:
return HubExecutionResult(
success=False,
server_name=server.name,
tool_name=tool.name,
error=str(e),
)
def list_servers(self) -> list[dict[str, Any]]:
"""List all registered servers with their status.
:returns: List of server info dicts.
"""
servers = []
for server in self._registry.servers:
servers.append({
"name": server.name,
"identifier": server.identifier,
"type": server.config.type.value,
"enabled": server.config.enabled,
"category": server.config.category,
"description": server.config.description,
"discovered": server.discovered,
"tool_count": len(server.tools),
"error": server.discovery_error,
})
return servers
def list_tools(self) -> list[dict[str, Any]]:
"""List all discovered tools.
:returns: List of tool info dicts.
"""
tools = []
for tool in self._registry.get_all_tools():
tools.append({
"identifier": tool.identifier,
"name": tool.name,
"server": tool.server_name,
"description": tool.description,
"parameters": [p.model_dump() for p in tool.parameters],
})
return tools
def get_tool_schema(self, identifier: str) -> dict[str, Any] | None:
"""Get the JSON Schema for a tool's input.
:param identifier: Tool identifier.
:returns: JSON Schema dict or None if not found.
"""
_, tool = self._registry.find_tool(identifier)
if tool:
return tool.input_schema
return None
@@ -0,0 +1,284 @@
"""Data models for FuzzForge Hub.
This module defines the Pydantic models used to represent MCP servers
and their tools in the hub registry.
"""
from __future__ import annotations
from enum import Enum
from typing import Any
from pydantic import BaseModel, Field
class HubServerType(str, Enum):
"""Type of MCP server connection."""
#: Run as Docker container with stdio transport.
DOCKER = "docker"
#: Run as local command/process with stdio transport.
COMMAND = "command"
#: Connect via Server-Sent Events (HTTP).
SSE = "sse"
class HubServerConfig(BaseModel):
"""Configuration for an MCP server in the hub.
This defines how to connect to an MCP server, not what tools it provides.
Tools are discovered dynamically at runtime.
"""
#: Unique identifier for this server (e.g., "nmap", "nuclei").
name: str = Field(description="Unique server identifier")
#: Human-readable description of the server.
description: str | None = Field(
default=None,
description="Human-readable description",
)
#: Type of connection to use.
type: HubServerType = Field(description="Connection type")
#: Docker image name (for type=docker).
image: str | None = Field(
default=None,
description="Docker image name (for docker type)",
)
#: Command to run (for type=command).
command: list[str] | None = Field(
default=None,
description="Command and args (for command type)",
)
#: URL endpoint (for type=sse).
url: str | None = Field(
default=None,
description="SSE endpoint URL (for sse type)",
)
#: Environment variables to pass to the server.
environment: dict[str, str] = Field(
default_factory=dict,
description="Environment variables",
)
#: Docker capabilities to add (e.g., ["NET_RAW"] for nmap).
capabilities: list[str] = Field(
default_factory=list,
description="Docker capabilities to add",
)
#: Volume mounts for Docker (e.g., ["/host/path:/container/path:ro"]).
volumes: list[str] = Field(
default_factory=list,
description="Docker volume mounts",
)
#: Whether this server is enabled.
enabled: bool = Field(
default=True,
description="Whether server is enabled",
)
#: Category for grouping (e.g., "reconnaissance", "web-security").
category: str | None = Field(
default=None,
description="Category for grouping servers",
)
class HubToolParameter(BaseModel):
"""A parameter for an MCP tool.
Parsed from the tool's JSON Schema inputSchema.
"""
#: Parameter name.
name: str
#: Parameter type (string, integer, boolean, array, object).
type: str
#: Human-readable description.
description: str | None = None
#: Whether this parameter is required.
required: bool = False
#: Default value if any.
default: Any = None
#: Enum values if constrained.
enum: list[Any] | None = None
class HubTool(BaseModel):
"""An MCP tool discovered from a hub server.
This is populated by calling `list_tools()` on the MCP server.
"""
#: Tool name as defined by the MCP server.
name: str = Field(description="Tool name from MCP server")
#: Human-readable description.
description: str | None = Field(
default=None,
description="Tool description",
)
#: Name of the hub server this tool belongs to.
server_name: str = Field(description="Parent server name")
#: Parsed parameters from inputSchema.
parameters: list[HubToolParameter] = Field(
default_factory=list,
description="Tool parameters",
)
#: Raw JSON Schema for the tool input.
input_schema: dict[str, Any] = Field(
default_factory=dict,
description="Raw JSON Schema from MCP",
)
@property
def identifier(self) -> str:
"""Get the full tool identifier (hub:server:tool)."""
return f"hub:{self.server_name}:{self.name}"
@classmethod
def from_mcp_tool(
cls,
server_name: str,
name: str,
description: str | None,
input_schema: dict[str, Any],
) -> HubTool:
"""Create a HubTool from MCP tool metadata.
:param server_name: Name of the parent hub server.
:param name: Tool name.
:param description: Tool description.
:param input_schema: JSON Schema for tool input.
:returns: HubTool instance.
"""
parameters = cls._parse_parameters(input_schema)
return cls(
name=name,
description=description,
server_name=server_name,
parameters=parameters,
input_schema=input_schema,
)
@staticmethod
def _parse_parameters(schema: dict[str, Any]) -> list[HubToolParameter]:
"""Parse parameters from JSON Schema.
:param schema: JSON Schema dict.
:returns: List of parsed parameters.
"""
parameters: list[HubToolParameter] = []
properties = schema.get("properties", {})
required_params = set(schema.get("required", []))
for name, prop in properties.items():
param = HubToolParameter(
name=name,
type=prop.get("type", "string"),
description=prop.get("description"),
required=name in required_params,
default=prop.get("default"),
enum=prop.get("enum"),
)
parameters.append(param)
return parameters
class HubServer(BaseModel):
"""A hub server with its discovered tools.
Combines configuration with dynamically discovered tools.
"""
#: Server configuration.
config: HubServerConfig
#: Tools discovered from the server (populated at runtime).
tools: list[HubTool] = Field(
default_factory=list,
description="Discovered tools",
)
#: Whether tools have been discovered.
discovered: bool = Field(
default=False,
description="Whether tools have been discovered",
)
#: Error message if discovery failed.
discovery_error: str | None = Field(
default=None,
description="Error message if discovery failed",
)
@property
def name(self) -> str:
"""Get server name."""
return self.config.name
@property
def identifier(self) -> str:
"""Get server identifier for module listing."""
return f"hub:{self.config.name}"
def get_tool(self, tool_name: str) -> HubTool | None:
"""Get a tool by name.
:param tool_name: Name of the tool.
:returns: HubTool if found, None otherwise.
"""
for tool in self.tools:
if tool.name == tool_name:
return tool
return None
class HubConfig(BaseModel):
"""Configuration for the entire hub.
Loaded from hub-servers.json or similar config file.
"""
#: List of configured servers.
servers: list[HubServerConfig] = Field(
default_factory=list,
description="Configured MCP servers",
)
#: Default timeout for tool execution (seconds).
default_timeout: int = Field(
default=300,
description="Default execution timeout",
)
#: Whether to cache discovered tools.
cache_tools: bool = Field(
default=True,
description="Cache discovered tools",
)
@@ -0,0 +1,258 @@
"""Hub registry for managing MCP server configurations.
The registry loads server configurations from a JSON file and provides
methods to access and manage them. It does not hardcode any specific
servers or tools - everything is configured by the user.
"""
from __future__ import annotations
import json
from pathlib import Path
from typing import TYPE_CHECKING, cast
from fuzzforge_common.hub.models import (
HubConfig,
HubServer,
HubServerConfig,
)
if TYPE_CHECKING:
from structlog.stdlib import BoundLogger
def get_logger() -> BoundLogger:
"""Get structlog logger instance.
:returns: Configured structlog logger.
"""
from structlog import get_logger # noqa: PLC0415
return cast("BoundLogger", get_logger())
class HubRegistry:
"""Registry for MCP hub servers.
Manages the configuration and state of hub servers.
Configurations are loaded from a JSON file.
"""
#: Loaded hub configuration.
_config: HubConfig
#: Server instances with discovered tools.
_servers: dict[str, HubServer]
#: Path to the configuration file.
_config_path: Path | None
def __init__(self, config_path: Path | str | None = None) -> None:
"""Initialize the hub registry.
:param config_path: Path to hub-servers.json config file.
If None, starts with empty configuration.
"""
if config_path is not None:
self._config_path = Path(config_path)
else:
self._config_path = None
self._servers = {}
self._config = HubConfig()
if self._config_path and self._config_path.exists():
self._load_config(self._config_path)
def _load_config(self, config_path: Path) -> None:
"""Load configuration from JSON file.
:param config_path: Path to config file.
"""
logger = get_logger()
try:
with config_path.open() as f:
data = json.load(f)
self._config = HubConfig.model_validate(data)
# Create server instances from config
for server_config in self._config.servers:
if server_config.enabled:
self._servers[server_config.name] = HubServer(
config=server_config,
)
logger.info(
"Loaded hub configuration",
path=str(config_path),
servers=len(self._servers),
)
except Exception as e:
logger.error(
"Failed to load hub configuration",
path=str(config_path),
error=str(e),
)
raise
def reload(self) -> None:
"""Reload configuration from file."""
if self._config_path and self._config_path.exists():
self._servers.clear()
self._load_config(self._config_path)
@property
def servers(self) -> list[HubServer]:
"""Get all registered servers.
:returns: List of hub servers.
"""
return list(self._servers.values())
@property
def enabled_servers(self) -> list[HubServer]:
"""Get all enabled servers.
:returns: List of enabled hub servers.
"""
return [s for s in self._servers.values() if s.config.enabled]
def get_server(self, name: str) -> HubServer | None:
"""Get a server by name.
:param name: Server name.
:returns: HubServer if found, None otherwise.
"""
return self._servers.get(name)
def add_server(self, config: HubServerConfig) -> HubServer:
"""Add a server to the registry.
:param config: Server configuration.
:returns: Created HubServer instance.
:raises ValueError: If server with same name exists.
"""
if config.name in self._servers:
msg = f"Server '{config.name}' already exists"
raise ValueError(msg)
server = HubServer(config=config)
self._servers[config.name] = server
self._config.servers.append(config)
get_logger().info("Added hub server", name=config.name, type=config.type)
return server
def remove_server(self, name: str) -> bool:
"""Remove a server from the registry.
:param name: Server name.
:returns: True if removed, False if not found.
"""
if name not in self._servers:
return False
del self._servers[name]
self._config.servers = [s for s in self._config.servers if s.name != name]
get_logger().info("Removed hub server", name=name)
return True
def save_config(self, path: Path | None = None) -> None:
"""Save current configuration to file.
:param path: Path to save to. Uses original path if None.
"""
save_path = path or self._config_path
if not save_path:
msg = "No config path specified"
raise ValueError(msg)
with save_path.open("w") as f:
json.dump(
self._config.model_dump(mode="json"),
f,
indent=2,
)
get_logger().info("Saved hub configuration", path=str(save_path))
def update_server_tools(
self,
server_name: str,
tools: list,
*,
error: str | None = None,
) -> None:
"""Update discovered tools for a server.
Called by the hub client after tool discovery.
:param server_name: Server name.
:param tools: List of HubTool instances.
:param error: Error message if discovery failed.
"""
server = self._servers.get(server_name)
if not server:
return
if error:
server.discovered = False
server.discovery_error = error
server.tools = []
else:
server.discovered = True
server.discovery_error = None
server.tools = tools
def get_all_tools(self) -> list:
"""Get all discovered tools from all servers.
:returns: Flat list of all HubTool instances.
"""
tools = []
for server in self._servers.values():
if server.discovered:
tools.extend(server.tools)
return tools
def find_tool(self, identifier: str):
"""Find a tool by its full identifier.
:param identifier: Full identifier (hub:server:tool or server:tool).
:returns: Tuple of (HubServer, HubTool) if found, (None, None) otherwise.
"""
# Parse identifier
parts = identifier.split(":")
if len(parts) == 3 and parts[0] == "hub": # noqa: PLR2004
# hub:server:tool format
server_name = parts[1]
tool_name = parts[2]
elif len(parts) == 2: # noqa: PLR2004
# server:tool format
server_name = parts[0]
tool_name = parts[1]
else:
return None, None
server = self._servers.get(server_name)
if not server:
return None, None
tool = server.get_tool(tool_name)
return server, tool
@@ -272,6 +272,23 @@ class AbstractFuzzForgeSandboxEngine(ABC):
message: str = f"method 'read_file_from_container' is not implemented for class '{self.__class__.__name__}'"
raise NotImplementedError(message)
@abstractmethod
def tail_file_from_container(self, identifier: str, path: str, start_line: int = 1) -> str:
"""Read a file from a running container starting at a given line number.
Uses ``tail -n +{start_line}`` to avoid re-reading the entire file on
every poll. This is the preferred method for incremental reads of
append-only files such as ``stream.jsonl``.
:param identifier: Container identifier.
:param path: Path to file inside container.
:param start_line: 1-based line number to start reading from.
:returns: File contents from *start_line* onwards (may be empty).
"""
message: str = f"method 'tail_file_from_container' is not implemented for class '{self.__class__.__name__}'"
raise NotImplementedError(message)
@abstractmethod
def list_containers(self, all_containers: bool = True) -> list[dict]:
"""List containers.
@@ -389,6 +389,24 @@ class DockerCLI(AbstractFuzzForgeSandboxEngine):
return ""
return result.stdout
def tail_file_from_container(self, identifier: str, path: str, start_line: int = 1) -> str:
"""Read a file from a container starting at a given line number.
:param identifier: Container identifier.
:param path: Path to file in container.
:param start_line: 1-based line number to start reading from.
:returns: File contents from *start_line* onwards.
"""
result = self._run(
["exec", identifier, "tail", "-n", f"+{start_line}", path],
check=False,
)
if result.returncode != 0:
get_logger().debug("failed to tail file from container", path=path, start_line=start_line)
return ""
return result.stdout
def list_containers(self, all_containers: bool = True) -> list[dict]:
"""List containers.
@@ -168,6 +168,11 @@ class Docker(AbstractFuzzForgeSandboxEngine):
message: str = "Docker engine read_file_from_container is not yet implemented"
raise NotImplementedError(message)
def tail_file_from_container(self, identifier: str, path: str, start_line: int = 1) -> str:
"""Read a file from a container starting at a given line number."""
message: str = "Docker engine tail_file_from_container is not yet implemented"
raise NotImplementedError(message)
def list_containers(self, all_containers: bool = True) -> list[dict]:
"""List containers."""
message: str = "Docker engine list_containers is not yet implemented"
@@ -449,6 +449,24 @@ class PodmanCLI(AbstractFuzzForgeSandboxEngine):
return ""
return result.stdout
def tail_file_from_container(self, identifier: str, path: str, start_line: int = 1) -> str:
"""Read a file from a container starting at a given line number.
:param identifier: Container identifier.
:param path: Path to file in container.
:param start_line: 1-based line number to start reading from.
:returns: File contents from *start_line* onwards.
"""
result = self._run(
["exec", identifier, "tail", "-n", f"+{start_line}", path],
check=False,
)
if result.returncode != 0:
get_logger().debug("failed to tail file from container", path=path, start_line=start_line)
return ""
return result.stdout
def list_containers(self, all_containers: bool = True) -> list[dict]:
"""List containers.
@@ -475,6 +475,30 @@ class Podman(AbstractFuzzForgeSandboxEngine):
return ""
return stdout.decode("utf-8", errors="replace") if stdout else ""
def tail_file_from_container(self, identifier: str, path: str, start_line: int = 1) -> str:
"""Read a file from a container starting at a given line number.
:param identifier: Container identifier.
:param path: Path to file inside container.
:param start_line: 1-based line number to start reading from.
:returns: File contents from *start_line* onwards.
"""
client: PodmanClient = self.get_client()
with client:
container: Container = client.containers.get(key=identifier)
(status, (stdout, stderr)) = container.exec_run(
cmd=["tail", "-n", f"+{start_line}", path],
demux=True,
)
if status != 0:
error_msg = stderr.decode("utf-8", errors="replace") if stderr else "File not found"
get_logger().debug(
"failed to tail file from container", path=path, start_line=start_line, error=error_msg,
)
return ""
return stdout.decode("utf-8", errors="replace") if stdout else ""
def list_containers(self, all_containers: bool = True) -> list[dict]:
"""List containers.