hot reloading the BRUTAL way

This commit is contained in:
lda
2026-04-30 06:53:32 +07:00 Verified
parent ef2f179110
commit 0d6f04b50e
12 changed files with 715 additions and 49 deletions
+5
View File
@@ -19,6 +19,7 @@ from .capabilities import (
)
from .catalog import CombinedCatalog
from .connections import ConnectionRegistry, parse_connection_id, qualify_node_name
from .config_manager import BrokerConfigManager, ConfigMutationError
from .config_models import (
BrokerConfigFile,
ConnectionConfigFile,
@@ -43,6 +44,7 @@ from .proxy_validation import ProxyConfigError, validate_transparent_proxy_confi
from .service import WfMcpService
from .store import FileStore, Store
from .transparent_proxy import (
TransparentProxyRuntime,
broker_config_to_fastmcp_config,
connection_to_fastmcp_server_config,
create_proxy_admin_server,
@@ -55,6 +57,7 @@ __all__ = [
"AuthRecord",
"BackendAdapter",
"BrokerConfig",
"BrokerConfigManager",
"CatalogNodeEntry",
"CatalogPromptEntry",
"CatalogResourceEntry",
@@ -63,6 +66,7 @@ __all__ = [
"ConnectionConfig",
"ConnectionConfigFile",
"ConnectionRegistry",
"ConfigMutationError",
"DiscoveredConnectionCapabilities",
"DiscoveredPrompt",
"DiscoveredResource",
@@ -77,6 +81,7 @@ __all__ = [
"Store",
"StdioConnectionMetadata",
"ToolCallResult",
"TransparentProxyRuntime",
"WfMcpService",
"build_service_from_config",
"broker_config_to_fastmcp_config",
+1
View File
@@ -228,6 +228,7 @@ def run_transparent_proxy_server(
config = load_broker_config(config_path)
server = create_transparent_proxy_server(
config,
config_path=config_path,
resources_as_tools=resources_as_tools,
prompts_as_tools=prompts_as_tools,
search_tools=search_tools,
+125
View File
@@ -0,0 +1,125 @@
from __future__ import annotations
import json
from pathlib import Path
from typing import Any
from .config_models import BrokerConfigFile, ConnectionConfigFile
from .models import BrokerConfig
class ConfigMutationError(ValueError):
"""Raised when a requested config mutation cannot be applied."""
class BrokerConfigManager:
def __init__(self, config_path: str | Path) -> None:
self.config_path = Path(config_path)
def load_file(self) -> BrokerConfigFile:
data = json.loads(self.config_path.read_text(encoding="utf-8"))
return BrokerConfigFile.model_validate(data)
def load_runtime(self) -> BrokerConfig:
return self.load_file().to_runtime(config_path=self.config_path)
def write_file(self, config: BrokerConfigFile) -> None:
payload = config.model_dump(mode="json", exclude_none=True)
text = json.dumps(payload, indent=2) + "\n"
self.config_path.write_text(text, encoding="utf-8")
def get_payload(self) -> dict[str, Any]:
return self.load_file().model_dump(mode="json", exclude_none=True)
def add_connection(
self,
*,
connection_id: str,
server: str,
account: str,
metadata: dict[str, Any] | None = None,
enabled: bool = True,
) -> dict[str, Any]:
config = self.load_file()
if _find_connection(config, connection_id) is not None:
raise ConfigMutationError(f"connection {connection_id!r} already exists")
connection = ConnectionConfigFile(
id=connection_id,
server=server,
account=account,
enabled=enabled,
metadata={} if metadata is None else metadata,
)
config.connections.append(connection)
self.write_file(config)
return _mutation_payload("add_connection", connection_id)
def update_connection(
self,
*,
connection_id: str,
server: str | None = None,
account: str | None = None,
metadata: dict[str, Any] | None = None,
enabled: bool | None = None,
) -> dict[str, Any]:
config = self.load_file()
index = _find_connection_index(config, connection_id)
if index is None:
raise ConfigMutationError(f"connection {connection_id!r} does not exist")
existing = config.connections[index]
config.connections[index] = ConnectionConfigFile(
id=existing.id,
server=existing.server if server is None else server,
account=existing.account if account is None else account,
enabled=existing.enabled if enabled is None else enabled,
metadata=existing.metadata if metadata is None else metadata,
)
self.write_file(config)
return _mutation_payload("update_connection", connection_id)
def set_connection_enabled(
self,
connection_id: str,
*,
enabled: bool,
) -> dict[str, Any]:
return self.update_connection(connection_id=connection_id, enabled=enabled)
def remove_connection(self, connection_id: str) -> dict[str, Any]:
config = self.load_file()
index = _find_connection_index(config, connection_id)
if index is None:
raise ConfigMutationError(f"connection {connection_id!r} does not exist")
del config.connections[index]
self.write_file(config)
return _mutation_payload("remove_connection", connection_id)
def _find_connection(
config: BrokerConfigFile,
connection_id: str,
) -> ConnectionConfigFile | None:
index = _find_connection_index(config, connection_id)
if index is None:
return None
return config.connections[index]
def _find_connection_index(
config: BrokerConfigFile,
connection_id: str,
) -> int | None:
for index, connection in enumerate(config.connections):
if connection.id == connection_id:
return index
return None
def _mutation_payload(action: str, connection_id: str) -> dict[str, Any]:
return {
"action": action,
"connection_id": connection_id,
"ok": True,
"requires_reload": True,
}
+161 -36
View File
@@ -1,6 +1,7 @@
from __future__ import annotations
from dataclasses import asdict
from pathlib import Path
from typing import Any
from fastmcp import FastMCP
@@ -12,6 +13,7 @@ from fastmcp.server import create_proxy
from fastmcp.server.transforms import Namespace, PromptsAsTools, ResourcesAsTools
from fastmcp.server.transforms.search import BM25SearchTransform
from .config_manager import BrokerConfigManager, ConfigMutationError
from .models import BrokerConfig, ConnectionConfig
from .proxy_validation import validate_transparent_proxy_config
@@ -19,10 +21,92 @@ _ADMIN_NAMESPACE = "wf.mcp"
_ADMIN_TOOL_NAMES = [
f"{_ADMIN_NAMESPACE}_list_connections",
f"{_ADMIN_NAMESPACE}_get_connection_statuses",
f"{_ADMIN_NAMESPACE}_get_config",
f"{_ADMIN_NAMESPACE}_reload_config",
f"{_ADMIN_NAMESPACE}_add_connection",
f"{_ADMIN_NAMESPACE}_update_connection",
f"{_ADMIN_NAMESPACE}_enable_connection",
f"{_ADMIN_NAMESPACE}_disable_connection",
f"{_ADMIN_NAMESPACE}_remove_connection",
]
def create_proxy_admin_server(config: BrokerConfig) -> FastMCP[Any]:
class TransparentProxyRuntime:
def __init__(
self,
config: BrokerConfig,
*,
config_path: str | Path | None = None,
resources_as_tools: bool = False,
prompts_as_tools: bool = False,
search_tools: bool = False,
) -> None:
self.config = config
self.manager = None if config_path is None else BrokerConfigManager(config_path)
self.server: FastMCP[Any] = FastMCP(
"wf-mcp-transparent-proxy",
instructions=(
"Transparent MCP proxy over configured upstream MCP connections. "
"Upstream tools, resources, and prompts are exposed as first-class "
"broker capabilities with connection-qualified names."
),
)
self.reload()
if resources_as_tools:
self.server.add_transform(ResourcesAsTools(self.server))
if prompts_as_tools:
self.server.add_transform(PromptsAsTools(self.server))
if search_tools:
self.server.add_transform(BM25SearchTransform(always_visible=_ADMIN_TOOL_NAMES))
def current_config(self) -> BrokerConfig:
if self.manager is None:
return self.config
self.config = self.manager.load_runtime()
return self.config
def require_manager(self) -> BrokerConfigManager:
if self.manager is None:
raise ConfigMutationError(
"config mutation tools require a config path-backed proxy"
)
return self.manager
def reload(self) -> dict[str, Any]:
config = self.current_config()
validate_transparent_proxy_config(config)
self.server.providers[:] = [self.server.local_provider]
admin = create_proxy_admin_server(self)
admin.add_transform(Namespace(_ADMIN_NAMESPACE))
self.server.mount(admin)
mounted_connections: list[str] = []
for connection in config.connections:
if not connection.enabled:
continue
server_config = broker_config_to_fastmcp_config(
BrokerConfig(store_root=config.store_root, connections=[connection])
)
transport = MCPConfigTransport(server_config, name_as_prefix=False)
client = Client(transport=transport, name=f"wf-mcp:{connection.id}")
proxy = create_proxy(client, name=f"Proxy-{connection.id}")
proxy.add_transform(Namespace(connection.id))
self.server.mount(proxy)
mounted_connections.append(connection.id)
return {
"ok": True,
"reloaded": True,
"mounted_connections": mounted_connections,
"connection_count": len(config.connections),
"enabled_connection_count": len(mounted_connections),
}
def create_proxy_admin_server(
runtime: TransparentProxyRuntime,
) -> FastMCP[Any]:
admin = FastMCP(
"wf-mcp-admin",
instructions="Administrative tools for this wf-mcp proxy instance.",
@@ -33,7 +117,7 @@ def create_proxy_admin_server(config: BrokerConfig) -> FastMCP[Any]:
return [
asdict(connection)
for connection in sorted(
config.connections,
runtime.current_config().connections,
key=lambda connection: connection.id,
)
]
@@ -49,11 +133,75 @@ def create_proxy_admin_server(config: BrokerConfig) -> FastMCP[Any]:
"transport": connection.metadata.get("transport"),
}
for connection in sorted(
config.connections,
runtime.current_config().connections,
key=lambda connection: connection.id,
)
]
@admin.tool()
async def get_config() -> dict[str, Any]:
if runtime.manager is not None:
return runtime.manager.get_payload()
config = runtime.current_config()
return {
"store_root": str(config.store_root),
"connections": [asdict(connection) for connection in config.connections],
}
@admin.tool()
async def reload_config() -> dict[str, Any]:
return runtime.reload()
@admin.tool()
async def add_connection(
connection_id: str,
server: str,
account: str,
metadata: dict[str, Any] | None = None,
enabled: bool = True,
) -> dict[str, Any]:
return runtime.require_manager().add_connection(
connection_id=connection_id,
server=server,
account=account,
metadata=metadata,
enabled=enabled,
)
@admin.tool()
async def update_connection(
connection_id: str,
server: str | None = None,
account: str | None = None,
metadata: dict[str, Any] | None = None,
enabled: bool | None = None,
) -> dict[str, Any]:
return runtime.require_manager().update_connection(
connection_id=connection_id,
server=server,
account=account,
metadata=metadata,
enabled=enabled,
)
@admin.tool()
async def enable_connection(connection_id: str) -> dict[str, Any]:
return runtime.require_manager().set_connection_enabled(
connection_id,
enabled=True,
)
@admin.tool()
async def disable_connection(connection_id: str) -> dict[str, Any]:
return runtime.require_manager().set_connection_enabled(
connection_id,
enabled=False,
)
@admin.tool()
async def remove_connection(connection_id: str) -> dict[str, Any]:
return runtime.require_manager().remove_connection(connection_id)
return admin
@@ -99,6 +247,7 @@ def broker_config_to_fastmcp_config(config: BrokerConfig) -> MCPConfig:
def create_transparent_proxy_server(
config: BrokerConfig,
*,
config_path: str | Path | None = None,
resources_as_tools: bool = False,
prompts_as_tools: bool = False,
search_tools: bool = False,
@@ -107,45 +256,20 @@ def create_transparent_proxy_server(
config,
resources_as_tools=resources_as_tools,
prompts_as_tools=prompts_as_tools,
# not yet idk codex help
)
root = FastMCP(
"wf-mcp-transparent-proxy",
instructions=(
"Transparent MCP proxy over configured upstream MCP connections. "
"Upstream tools, resources, and prompts are exposed as first-class "
"broker capabilities with connection-qualified names."
),
)
admin = create_proxy_admin_server(config)
admin.add_transform(Namespace(_ADMIN_NAMESPACE))
root.mount(admin)
for connection in config.connections:
if not connection.enabled:
continue
server_config = broker_config_to_fastmcp_config(
BrokerConfig(store_root=config.store_root, connections=[connection])
)
transport = MCPConfigTransport(server_config, name_as_prefix=False)
client = Client(transport=transport, name=f"wf-mcp:{connection.id}")
proxy = create_proxy(client, name=f"Proxy-{connection.id}")
proxy.add_transform(Namespace(connection.id))
root.mount(proxy)
if resources_as_tools:
root.add_transform(ResourcesAsTools(root))
if prompts_as_tools:
root.add_transform(PromptsAsTools(root))
if search_tools:
root.add_transform(BM25SearchTransform(always_visible=_ADMIN_TOOL_NAMES))
return root
return TransparentProxyRuntime(
config,
config_path=config_path,
resources_as_tools=resources_as_tools,
prompts_as_tools=prompts_as_tools,
search_tools=search_tools,
).server
def create_transparent_proxy_client(
config: BrokerConfig,
*,
config_path: str | Path | None = None,
resources_as_tools: bool = False,
prompts_as_tools: bool = False,
search_tools: bool = False,
@@ -154,6 +278,7 @@ def create_transparent_proxy_client(
FastMCPTransport(
create_transparent_proxy_server(
config,
config_path=config_path,
resources_as_tools=resources_as_tools,
prompts_as_tools=prompts_as_tools,
search_tools=search_tools,