proxymountregistry

This commit is contained in:
lda
2026-05-15 23:29:42 +07:00 Verified
parent 3f83afc100
commit 6935329d44
4 changed files with 220 additions and 17 deletions
+13
View File
@@ -4,3 +4,16 @@ no: assert dict == dict
yes: assert dict['field'] == dict['field'] unless we know better yes: assert dict['field'] == dict['field'] unless we know better
more later more later
# Test suite
```bash
uv run /* --env-file .env */ pytest -q
(uv run / uvx) ruff check / format
uv run basedpyright --level error # error to cut spam
# maybe uvx ty
```
or so i think.
<!-- if this file comes with every request, tell me, and you have perms to cut the files down. goo goo ga ga. use caveman skill (only) for repeated artifacts. -->
+105
View File
@@ -0,0 +1,105 @@
from __future__ import annotations
import json
from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Generic, TypeVar
from fastmcp import FastMCP
from fastmcp.client import Client
from fastmcp.client.transports.config import MCPConfigTransport
from fastmcp.server import create_proxy
from fastmcp.server.transforms import Namespace
from ..models import BrokerConfig, ConnectionConfig
from ..proxy_config import broker_config_to_fastmcp_config
ProxyT = TypeVar("ProxyT")
ProxyMountFactory = Callable[[ConnectionConfig, Path], "ProxyMount[ProxyT]"]
@dataclass(frozen=True, slots=True)
class ProxyMount(Generic[ProxyT]):
"""Reusable proxy mount for one enabled upstream connection."""
connection_id: str
fingerprint: str
proxy: ProxyT
class ProxyMountRegistry(Generic[ProxyT]):
"""Reuse unchanged proxy mounts while keeping retirement explicit."""
def __init__(self, factory: ProxyMountFactory[ProxyT]) -> None:
self._factory = factory
self._mounts: dict[str, ProxyMount[ProxyT]] = {}
def active_mounts_for(self, config: BrokerConfig) -> list[ProxyMount[ProxyT]]:
"""Return enabled mounts, reusing unchanged connection fingerprints."""
active: list[ProxyMount[ProxyT]] = []
for connection in config.connections:
if not connection.enabled:
continue
active.append(self.get_or_create(connection, store_root=config.store_root))
return active
def get_or_create(
self,
connection: ConnectionConfig,
*,
store_root: Path,
) -> ProxyMount[ProxyT]:
"""Return a cached mount when connection transport identity is unchanged."""
fingerprint = connection_fingerprint(connection)
current = self._mounts.get(connection.id)
if current is not None and current.fingerprint == fingerprint:
return current
created = self._factory(connection, store_root)
mount = ProxyMount(
connection_id=created.connection_id,
fingerprint=fingerprint,
proxy=created.proxy,
)
self._mounts[connection.id] = mount
return mount
def retired_connection_ids(self, active_connection_ids: set[str]) -> set[str]:
"""Return cached connection ids that are not part of the active reload set."""
return set(self._mounts) - active_connection_ids
def connection_fingerprint(connection: ConnectionConfig) -> str:
"""Return a deterministic internal reuse key for one connection config."""
return json.dumps(
{
"id": connection.id,
"server": connection.server,
"account": connection.account,
"enabled": connection.enabled,
"metadata": connection.metadata,
},
sort_keys=True,
separators=(",", ":"),
default=str,
)
def create_proxy_mount(
connection: ConnectionConfig,
store_root: Path,
) -> ProxyMount[FastMCP[Any]]:
"""Create one FastMCP proxy mount for an enabled upstream connection."""
server_config = broker_config_to_fastmcp_config(
BrokerConfig(store_root=store_root, connections=[connection])
)
transport = MCPConfigTransport(server_config, name_as_prefix=False)
client = Client(transport=transport, name=f"wf-mcp:{connection.id}")
proxy: FastMCP[Any] = create_proxy(client, name=f"Proxy-{connection.id}")
proxy.add_transform(Namespace(connection.id))
return ProxyMount(
connection_id=connection.id,
fingerprint=connection_fingerprint(connection),
proxy=proxy,
)
+9 -17
View File
@@ -5,19 +5,17 @@ from typing import Any
from fastmcp import FastMCP from fastmcp import FastMCP
from fastmcp.client import Client from fastmcp.client import Client
from fastmcp.client.transports.config import MCPConfigTransport
from fastmcp.client.transports.memory import FastMCPTransport from fastmcp.client.transports.memory import FastMCPTransport
from fastmcp.server import create_proxy from fastmcp.server.transforms import PromptsAsTools, ResourcesAsTools
from fastmcp.server.transforms import Namespace, PromptsAsTools, ResourcesAsTools
from fastmcp.server.transforms.search import BM25SearchTransform from fastmcp.server.transforms.search import BM25SearchTransform
from ..control import BrokerConfigManager, ConfigMutationError from ..control import BrokerConfigManager, ConfigMutationError
from ..events import EventBus from ..events import EventBus
from ..models import BrokerConfig from ..models import BrokerConfig
from ..shared.names import ADMIN_NAMESPACE, LdaNamespace from ..shared.names import ADMIN_NAMESPACE, LdaNamespace
from ..proxy_config import broker_config_to_fastmcp_config
from ..proxy_validation import validate_transparent_proxy_config from ..proxy_validation import validate_transparent_proxy_config
from .admin import create_proxy_admin_server from .admin import create_proxy_admin_server
from .mounts import ProxyMountRegistry, create_proxy_mount
from .tools import ( from .tools import (
ProxyToolPayload, ProxyToolPayload,
collect_proxy_tools, collect_proxy_tools,
@@ -71,6 +69,9 @@ class ProxyRuntime:
) )
self.admin_tools = admin_tools self.admin_tools = admin_tools
self.event_bus = event_bus self.event_bus = event_bus
self.mounts: ProxyMountRegistry[FastMCP[Any]] = ProxyMountRegistry(
create_proxy_mount
)
self.reload() self.reload()
if resources_as_tools: if resources_as_tools:
self.server.add_transform(ResourcesAsTools(self.server)) self.server.add_transform(ResourcesAsTools(self.server))
@@ -104,19 +105,10 @@ class ProxyRuntime:
admin.add_transform(LdaNamespace(ADMIN_NAMESPACE)) admin.add_transform(LdaNamespace(ADMIN_NAMESPACE))
self.server.mount(admin) self.server.mount(admin)
mounted_connections: list[str] = [] mounts = self.mounts.active_mounts_for(config)
for connection in config.connections: for mount in mounts:
if not connection.enabled: self.server.mount(mount.proxy)
continue mounted_connections = [mount.connection_id for mount in mounts]
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)
result = ProxyReloadResult( result = ProxyReloadResult(
mounted_connections=mounted_connections, mounted_connections=mounted_connections,
+93
View File
@@ -0,0 +1,93 @@
from __future__ import annotations
from pathlib import Path
from wf_mcp.models import BrokerConfig, ConnectionConfig
from wf_mcp.transparent_proxy.mounts import ProxyMount, ProxyMountRegistry
def test_registry_reuses_unchanged_enabled_mount() -> None:
created: list[str] = []
registry = ProxyMountRegistry[object](
lambda connection, store_root: _fake_mount(connection, store_root, created)
)
config = BrokerConfig(
store_root=Path(".wf_mcp_store"),
connections=[_connection()],
)
first = registry.active_mounts_for(config)
second = registry.active_mounts_for(config)
assert first[0] is second[0]
assert created == ["fixture.personal"]
def test_registry_replaces_mount_when_connection_changes() -> None:
created: list[str] = []
registry = ProxyMountRegistry[object](
lambda connection, store_root: _fake_mount(connection, store_root, created)
)
initial = BrokerConfig(
store_root=Path(".wf_mcp_store"),
connections=[_connection()],
)
changed = BrokerConfig(
store_root=Path(".wf_mcp_store"),
connections=[_connection(metadata={"transport": "stdio", "args": ["new.py"]})],
)
first = registry.active_mounts_for(initial)
second = registry.active_mounts_for(changed)
assert first[0] is not second[0]
assert created == ["fixture.personal", "fixture.personal"]
def test_registry_skips_disabled_mounts_and_reports_retired_connections() -> None:
registry = ProxyMountRegistry[object](
lambda connection, store_root: _fake_mount(connection, store_root)
)
initial = BrokerConfig(
store_root=Path(".wf_mcp_store"),
connections=[_connection()],
)
disabled = BrokerConfig(
store_root=Path(".wf_mcp_store"),
connections=[_connection(enabled=False)],
)
first = registry.active_mounts_for(initial)
second = registry.active_mounts_for(disabled)
assert [mount.connection_id for mount in first] == ["fixture.personal"]
assert second == []
assert registry.retired_connection_ids(set()) == {"fixture.personal"}
def _connection(
*,
enabled: bool = True,
metadata: dict[str, object] | None = None,
) -> ConnectionConfig:
return ConnectionConfig(
id="fixture.personal",
server="fixture",
account="personal",
enabled=enabled,
metadata=metadata or {"transport": "stdio"},
)
def _fake_mount(
connection: ConnectionConfig,
store_root: Path,
created: list[str] | None = None,
) -> ProxyMount[object]:
if created is not None:
created.append(connection.id)
return ProxyMount(
connection_id=connection.id,
fingerprint=f"{connection.id}:{store_root}:{connection.metadata}",
proxy=object(),
)