hot reloading the BRUTAL way
This commit is contained in:
@@ -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",
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user