feat: merge source registry at startup

This commit is contained in:
lda
2026-06-04 00:29:44 +07:00 Verified
parent 4e8726f36d
commit 2433d58988
10 changed files with 273 additions and 24 deletions
+48
View File
@@ -9,6 +9,12 @@ from mcp import types as mcp_types
from wf_mcp.broker.config import load_broker_config
from wf_mcp.models import BrokerConfig
from wf_mcp.server import create_server_client
from wf_mcp.source_registry import (
FileSourceRegistryStore,
McpSourceRegistryEntry,
SourceRegistryFile,
StdioSourceTransport,
)
from ..test_support import fixture_server_path, local_temp_root
from .conftest import structured
@@ -178,3 +184,45 @@ def test_server_reload_syncs_service_connection_source_enabled_state() -> None:
assert "fixture.personal.echo_tool" in names
asyncio.run(run_proxy())
def test_server_reload_preserves_source_registry_connections() -> None:
tmp_path = local_temp_root() / "unified_reload_registry_source_store"
tmp_path.mkdir(parents=True, exist_ok=True)
config_path = tmp_path / "wf_mcp.config.json"
config_path.write_text(
json.dumps({"store_root": ".wf_mcp_store", "connections": []}),
encoding="utf-8",
)
config = load_broker_config(config_path)
FileSourceRegistryStore(config.store_root).save_registry(
SourceRegistryFile(
sources=[
McpSourceRegistryEntry(
id="fixture.registry",
kind="mcp",
enabled=True,
provider="fixture",
account="registry",
transport=StdioSourceTransport(command=sys.executable),
)
]
)
)
async def run_proxy() -> None:
client = create_server_client(config, config_path=config_path)
async with client:
before = await client.call_tool("wf.admin.list_sources", {"limit": 100})
before_ids = {
source["id"] for source in structured(before)["sources"]
}
assert "fixture.registry" in before_ids
await client.call_tool("wf.admin.reload_config")
after = await client.call_tool("wf.admin.list_sources", {"limit": 100})
after_ids = {source["id"] for source in structured(after)["sources"]}
assert "fixture.registry" in after_ids
asyncio.run(run_proxy())
@@ -7,6 +7,12 @@ from wf_mcp.broker.service.source_catalog import SourceCatalogService
from wf_mcp.events import EventBus
from wf_mcp.models import BrokerConfig, ConnectionConfig
from wf_mcp.runtime import ToolExecutor
from wf_mcp.source_registry import (
FileSourceRegistryStore,
McpSourceRegistryEntry,
SourceRegistryFile,
StdioSourceTransport,
)
from wf_mcp.storage import FileStore
from ..test_support import local_temp_root
@@ -174,3 +180,93 @@ def test_wfmcpservice_sync_connections_delegates_to_connection_service() -> None
]
assert "demo.personal" not in service.capability_sources
assert "demo.work" in service.capability_sources
# ---------------------------------------------------------------------------
# Source registry merge helpers and tests
# ---------------------------------------------------------------------------
def _registry_entry(
source_id: str = "demo.registry",
*,
enabled: bool = True,
) -> McpSourceRegistryEntry:
return McpSourceRegistryEntry(
id=source_id,
kind="mcp",
enabled=enabled,
provider="demo",
account=source_id.rsplit(".", 1)[-1],
transport=StdioSourceTransport(command="demo-server"),
)
def test_connection_service_sync_merges_registry_entries() -> None:
service = ConnectionService(events=BrokerEventRecorder(EventBus()))
catalog = _source_catalog(service)
store = FileSourceRegistryStore(local_temp_root() / "registry_merge")
store.save_registry(SourceRegistryFile(sources=[_registry_entry()]))
service.sync_connections_from_config(
BrokerConfig(store_root=local_temp_root(), connections=[]),
source_registry_store=store,
)
assert [connection.id for connection in service.list_all()] == ["demo.registry"]
assert "demo.registry" in catalog.capability_sources
def test_connection_service_sync_config_shadows_registry_entry() -> None:
service = ConnectionService(events=BrokerEventRecorder(EventBus()))
_source_catalog(service)
store = FileSourceRegistryStore(local_temp_root() / "registry_shadow")
store.save_registry(SourceRegistryFile(sources=[_registry_entry("demo.same")]))
service.sync_connections_from_config(
BrokerConfig(
store_root=local_temp_root(),
connections=[
ConnectionConfig(id="demo.same", server="demo", account="config"),
],
),
source_registry_store=store,
)
assert service.get("demo.same").account == "config"
assert any(
event.kind == "source_registry_ignored_config_shadow"
and event.connection_id == "demo.same"
for event in service.events.list_events()
)
def test_connection_service_sync_registry_disabled_entry_hydrates_disabled_source() -> None:
service = ConnectionService(events=BrokerEventRecorder(EventBus()))
catalog = _source_catalog(service)
store = FileSourceRegistryStore(local_temp_root() / "registry_disabled")
store.save_registry(SourceRegistryFile(sources=[_registry_entry(enabled=False)]))
service.sync_connections_from_config(
BrokerConfig(store_root=local_temp_root(), connections=[]),
source_registry_store=store,
)
assert service.get("demo.registry").enabled is False
assert catalog.capability_sources["demo.registry"].enabled is False
def test_wfmcpservice_sync_connections_delegates_registry_store() -> None:
service = WfMcpService(store=FileStore(local_temp_root() / "facade_registry"))
store = FileSourceRegistryStore(local_temp_root() / "facade_registry_store")
store.save_registry(SourceRegistryFile(sources=[_registry_entry()]))
service.sync_connections_from_config(
BrokerConfig(store_root=local_temp_root(), connections=[]),
source_registry_store=store,
)
assert [connection.id for connection in service.connections.list_all()] == [
"demo.registry"
]
assert "demo.registry" in service.capability_sources
+64
View File
@@ -19,6 +19,12 @@ from wf_mcp.broker import (
load_broker_config,
)
from wf_mcp.models import BrokerConfig, ConnectionConfig
from wf_mcp.source_registry import (
FileSourceRegistryStore,
McpSourceRegistryEntry,
SourceRegistryFile,
StdioSourceTransport,
)
from wf_mcp.storage import FileStore
from .test_support import (
@@ -485,6 +491,64 @@ def test_build_service_from_config_uses_store_root_for_workflow_stores() -> None
assert service.run_store.root == store_root
# ---------------------------------------------------------------------------
# Source registry integration tests
# ---------------------------------------------------------------------------
def _registry_entry(
source_id: str = "demo.registry",
*,
enabled: bool = True,
) -> McpSourceRegistryEntry:
return McpSourceRegistryEntry(
id=source_id,
kind="mcp",
enabled=enabled,
provider="demo",
account=source_id.rsplit(".", 1)[-1],
transport=StdioSourceTransport(command="demo-server"),
)
def test_build_service_from_config_loads_source_registry_entries() -> None:
tmp_path = local_temp_root() / "broker_config_registry_load"
tmp_path.mkdir(parents=True, exist_ok=True)
config = BrokerConfig(store_root=tmp_path, connections=[])
FileSourceRegistryStore(tmp_path).save_registry(
SourceRegistryFile(sources=[_registry_entry("fixture.registry")])
)
service = build_service_from_config(config)
assert service.connections.get("fixture.registry").server == "demo"
assert "demo" in service.adapters
assert "fixture.registry" in service.capability_sources
def test_build_service_from_config_config_shadows_registry() -> None:
tmp_path = local_temp_root() / "broker_config_registry_shadow"
tmp_path.mkdir(parents=True, exist_ok=True)
config = BrokerConfig(
store_root=tmp_path,
connections=[
ConnectionConfig(id="fixture.same", server="fixture", account="config"),
],
)
FileSourceRegistryStore(tmp_path).save_registry(
SourceRegistryFile(sources=[_registry_entry("fixture.same")])
)
service = build_service_from_config(config)
assert service.connections.get("fixture.same").account == "config"
assert any(
event.kind == "source_registry_ignored_config_shadow"
and event.connection_id == "fixture.same"
for event in service.list_events()
)
def _artifact() -> WorkflowArtifact:
return WorkflowArtifact(
id="summarize_docs",