connection service

This commit is contained in:
lda
2026-06-02 23:34:50 +07:00 Verified
parent 83d97ecd62
commit e7f7ff4851
6 changed files with 845 additions and 43 deletions
@@ -0,0 +1,77 @@
from __future__ import annotations
from dataclasses import dataclass, field
from ...connections import ConnectionRegistry, parse_connection_id
from ...models import BrokerConfig, ConnectionConfig
from ...shared.names import RESERVED_CONNECTION_IDS
from .events import BrokerEventRecorder
from .source_catalog import SourceCatalogService
@dataclass(slots=True)
class ConnectionService:
"""Own broker connection registration and config reconciliation.
SourceCatalogService needs connection lookup callbacks during construction,
while registering a connection needs source-catalog hydration. The catalog is
therefore bound after both services exist; `_source_catalog()` makes that
construction cycle explicit and fail-fast.
"""
events: BrokerEventRecorder
connections: ConnectionRegistry = field(default_factory=ConnectionRegistry)
source_catalog: SourceCatalogService | None = None
def bind_source_catalog(self, source_catalog: SourceCatalogService) -> None:
self.source_catalog = source_catalog
def get(self, connection_id: str) -> ConnectionConfig:
return self.connections.get(connection_id)
def list_all(self) -> list[ConnectionConfig]:
return self.connections.list_all()
def list_enabled(self) -> list[ConnectionConfig]:
return self.connections.list_enabled()
def register_connection(self, connection: ConnectionConfig) -> None:
self._validate_connection_id(connection.id)
self.connections.register(connection)
self._source_catalog().hydrate_connection_source_from_snapshot(connection)
self.events.record_kind(
"connection_registered",
connection_id=connection.id,
payload={"server": connection.server, "account": connection.account},
)
def sync_connections_from_config(self, config: BrokerConfig) -> None:
"""Reconcile registry/source state after the public server reloads config."""
source_catalog = self._source_catalog()
next_ids = {connection.id for connection in config.connections}
previous_ids = set(self.connections.connections)
for connection_id in previous_ids - next_ids:
del self.connections.connections[connection_id]
source_catalog.capability_sources.pop(connection_id, None)
for connection in config.connections:
self._validate_connection_id(connection.id)
self.connections.register(connection)
source = source_catalog.capability_sources.get(connection.id)
if source is None:
source_catalog.hydrate_connection_source_from_snapshot(connection)
else:
source.enabled = connection.enabled
def _source_catalog(self) -> SourceCatalogService:
if self.source_catalog is None:
raise RuntimeError(
"ConnectionService requires a bound SourceCatalogService"
)
return self.source_catalog
@staticmethod
def _validate_connection_id(connection_id: str) -> None:
parse_connection_id(connection_id)
if connection_id in RESERVED_CONNECTION_IDS:
raise ValueError(f"connection id {connection_id!r} is reserved by wf-mcp")
+20 -43
View File
@@ -21,7 +21,7 @@ from wf_api.models import RawWorkflowPlan
from wf_platform import (
CapabilitySource,
)
from ...connections import ConnectionRegistry, parse_connection_id
from ...connections import ConnectionRegistry
from ...events import EventBus, McpEvent, make_event
from ...models import (
AuthRecord,
@@ -34,7 +34,7 @@ from ...models import (
)
from ...sdk import BackendAdapter
from ...runtime import ToolExecutor
from ...shared.names import RESERVED_CONNECTION_IDS
from .connection_service import ConnectionService
from ...storage import Store
from wf_api.saved_subgraphs import SavedSubgraphTree
from ..admin_capabilities import admin_source
@@ -50,7 +50,6 @@ from .workflow_runtime import WorkflowRuntimeService
class WfMcpService:
store: Store
default_catalog_max_age_seconds: int = 300
connections: ConnectionRegistry = field(default_factory=ConnectionRegistry)
event_bus: EventBus = field(default_factory=EventBus)
include_builtin_specs: bool = True
artifact_store: WorkflowArtifactStore | None = None
@@ -58,6 +57,7 @@ class WfMcpService:
run_store: RunStore | None = None
tool_executor: ToolExecutor | None = None
events: BrokerEventRecorder = field(init=False)
connection_service: ConnectionService = field(init=False)
upstream: UpstreamTransportService = field(init=False)
source_catalog: SourceCatalogService = field(init=False)
workflow_runtime: WorkflowRuntimeService = field(init=False)
@@ -70,6 +70,7 @@ class WfMcpService:
CLI, MCP, and future HTTP frontends may share or swap those stores.
"""
self.events = BrokerEventRecorder(self.event_bus)
self.connection_service = ConnectionService(events=self.events)
self.upstream = UpstreamTransportService(
store=self.store,
event_sink=self.events.record_event,
@@ -77,14 +78,15 @@ class WfMcpService:
)
self.source_catalog = SourceCatalogService(
store=self.store,
connection_lookup=self.connections.get,
connection_list_enabled=self.connections.list_enabled,
connection_list_all=self.connections.list_all,
connection_lookup=self.connection_service.get,
connection_list_enabled=self.connection_service.list_enabled,
connection_list_all=self.connection_service.list_all,
tool_executor_for=self.upstream.tool_executor_for,
load_auth=self.upstream.load_auth,
emit_event=self.events.record_event,
default_catalog_max_age_seconds=self.default_catalog_max_age_seconds,
)
self.connection_service.bind_source_catalog(self.source_catalog)
if self.include_builtin_specs:
for source in builtin_sources().values():
self.register_capability_source(source)
@@ -104,51 +106,26 @@ class WfMcpService:
"""
return self.source_catalog.capability_sources
@property
def connections(self) -> ConnectionRegistry:
"""Compatibility view of the broker connection registry.
Connection lifecycle ownership has moved to ConnectionService. Keep this
property because admin handlers, CLI helpers, and tests still inspect the
registry through the service facade.
"""
return self.connection_service.connections
@property
def adapters(self) -> dict[str, BackendAdapter]:
"""Compatibility view of upstream adapter registry."""
return self.upstream.adapters
def register_connection(self, connection: ConnectionConfig) -> None:
parse_connection_id(connection.id)
if connection.id in RESERVED_CONNECTION_IDS:
raise ValueError(f"connection id {connection.id!r} is reserved by wf-mcp")
self.connections.register(connection)
self.source_catalog.hydrate_connection_source_from_snapshot(connection)
self._record_event(
make_event(
"connection_registered",
connection_id=connection.id,
payload={"server": connection.server, "account": connection.account},
)
)
self.connection_service.register_connection(connection)
def sync_connections_from_config(self, config: BrokerConfig) -> None:
"""Reconcile connection sources after the public server reloads config.
The public server has two cooperating views: proxy mounts read the live
file-backed config, while workflow discovery reads this service's source
registry. Reload must keep those views aligned, or raw proxy tools can be
enabled while planner-visible workflow capabilities remain disabled.
"""
next_ids = {connection.id for connection in config.connections}
previous_ids = set(self.connections.connections)
for connection_id in previous_ids - next_ids:
del self.connections.connections[connection_id]
self.capability_sources.pop(connection_id, None)
for connection in config.connections:
parse_connection_id(connection.id)
if connection.id in RESERVED_CONNECTION_IDS:
raise ValueError(
f"connection id {connection.id!r} is reserved by wf-mcp"
)
self.connections.register(connection)
source = self.capability_sources.get(connection.id)
if source is None:
self.source_catalog.hydrate_connection_source_from_snapshot(connection)
else:
source.enabled = connection.enabled
self.connection_service.sync_connections_from_config(config)
def register_adapter(self, server: str, adapter: BackendAdapter) -> None:
self.upstream.register_adapter(server, adapter)