refactor: compartmentalize mcp models

src/wf_mcp/broker/__init__.py

@@ -12,13 +12,17 @@ from .server import (
     workflow_server_from_service,
 )
 from .config import build_service_from_config, load_broker_config
+from .models import BrokerConfig, ConnectionConfig, SourceConfigOwnership
 from .transport import normalize_transport
 from .service import WfMcpService
 
 __all__ = [
+    "BrokerConfig",
     "CombinedCatalog",
+    "ConnectionConfig",
     "DiscoveredConnectionCapabilities",
     "McpEvent",
+    "SourceConfigOwnership",
     "WfMcpService",
     "build_service_from_config",
     "build_workflow_server_from_config",

src/wf_mcp/broker/models.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Literal
+
+
+SourceConfigOwnership = Literal["locked", "seed"]
+
+
+@dataclass(slots=True)
+class ConnectionConfig:
+    id: str
+    server: str
+    account: str
+    enabled: bool = True
+    metadata: dict[str, Any] = field(default_factory=dict)
+    source_config_ownership: SourceConfigOwnership = "locked"
+
+
+@dataclass(slots=True)
+class BrokerConfig:
+    store_root: Path
+    connections: list[ConnectionConfig] = field(default_factory=list)
+
+
+__all__ = [
+    "BrokerConfig",
+    "ConnectionConfig",
+    "SourceConfigOwnership",
+]

src/wf_mcp/broker/service/core.py

@@ -16,6 +16,7 @@ from wf_core import (
     Workflow,
 )
 from wf_api.models import RawWorkflowPlan
+from wf_mcp.capabilities import CatalogNodeEntry, CatalogPromptEntry, CatalogResourceEntry
 from wf_platform import (
     CapabilitySource,
 )
@@ -23,9 +24,6 @@ from ...connections import ConnectionRegistry
 from ...events import EventBus, McpEvent
 from ...models import (
     AuthRecord,
-    CatalogNodeEntry,
-    CatalogPromptEntry,
-    CatalogResourceEntry,
     CatalogSnapshot,
     BrokerConfig,
     ConnectionConfig,

src/wf_mcp/broker/service/source_catalog.py

@@ -7,6 +7,7 @@ from typing import Any
 
 from pydantic import BaseModel
 from wf_authoring import NodeReturn, NodeSpec
+from wf_mcp.capabilities import CatalogNodeEntry, CatalogPromptEntry, CatalogResourceEntry
 from wf_platform import (
     CapabilityBuckets,
     CapabilitySource,
@@ -21,9 +22,6 @@ from ...connections import ConnectionConfig, qualify_node_name
 from ...events import McpEvent, make_event
 from ...models import (
     AuthRecord,
-    CatalogNodeEntry,
-    CatalogPromptEntry,
-    CatalogResourceEntry,
     CatalogSnapshot,
 )
 from ...runtime import ToolExecutor

src/wf_mcp/catalog/__init__.py

@@ -0,0 +1,6 @@
+from .models import CatalogSnapshot, dump_catalog_snapshot
+
+__all__ = [
+    "CatalogSnapshot",
+    "dump_catalog_snapshot",
+]

src/wf_mcp/catalog/models.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass, field
+from typing import Any
+
+from wf_mcp.capabilities import (
+    CatalogNodeEntry,
+    CatalogPromptEntry,
+    CatalogResourceEntry,
+)
+
+
+@dataclass(slots=True)
+class CatalogSnapshot:
+    connection_id: str
+    fetched_at_epoch_ms: int
+    max_age_seconds: int
+    nodes: list[CatalogNodeEntry] = field(default_factory=list)
+    resources: list[CatalogResourceEntry] = field(default_factory=list)
+    prompts: list[CatalogPromptEntry] = field(default_factory=list)
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    def is_stale(self, now_epoch_ms: int) -> bool:
+        age_ms = now_epoch_ms - self.fetched_at_epoch_ms
+        return age_ms > self.max_age_seconds * 1000
+
+
+def dump_catalog_snapshot(snapshot: CatalogSnapshot) -> dict[str, Any]:
+    return {
+        "connection_id": snapshot.connection_id,
+        "fetched_at_epoch_ms": snapshot.fetched_at_epoch_ms,
+        "max_age_seconds": snapshot.max_age_seconds,
+        "nodes": [asdict(node) for node in snapshot.nodes],
+        "resources": [asdict(resource) for resource in snapshot.resources],
+        "prompts": [asdict(prompt) for prompt in snapshot.prompts],
+        "metadata": snapshot.metadata,
+    }
+
+
+__all__ = [
+    "CatalogSnapshot",
+    "dump_catalog_snapshot",
+]

src/wf_mcp/models.py

@@ -1,26 +1,12 @@
 from __future__ import annotations
 
-from dataclasses import asdict, dataclass, field
+from dataclasses import dataclass, field
-from pathlib import Path
+from typing import Any
-from typing import Any, Literal
-
-from .capabilities import CatalogNodeEntry, CatalogPromptEntry, CatalogResourceEntry
 
 # RawWorkflowPlan moved to wf_api.models; re-exported here for backward compat.
 from wf_api.models import RawWorkflowPlan  # noqa: F401
-
+from wf_mcp.broker.models import BrokerConfig, ConnectionConfig, SourceConfigOwnership
-
+from wf_mcp.catalog.models import CatalogSnapshot, dump_catalog_snapshot
-SourceConfigOwnership = Literal["locked", "seed"]
-
-
-@dataclass(slots=True)
-class ConnectionConfig:
-    id: str
-    server: str
-    account: str
-    enabled: bool = True
-    metadata: dict[str, Any] = field(default_factory=dict)
-    source_config_ownership: SourceConfigOwnership = "locked"
 
 
 @dataclass(slots=True)
@@ -30,34 +16,12 @@ class AuthRecord:
     payload: dict[str, Any] = field(default_factory=dict)
 
 
-@dataclass(slots=True)
+__all__ = [
-class CatalogSnapshot:
+    "AuthRecord",
-    connection_id: str
+    "BrokerConfig",
-    fetched_at_epoch_ms: int
+    "CatalogSnapshot",
-    max_age_seconds: int
+    "ConnectionConfig",
-    nodes: list[CatalogNodeEntry] = field(default_factory=list)
+    "RawWorkflowPlan",
-    resources: list[CatalogResourceEntry] = field(default_factory=list)
+    "SourceConfigOwnership",
-    prompts: list[CatalogPromptEntry] = field(default_factory=list)
+    "dump_catalog_snapshot",
-    metadata: dict[str, Any] = field(default_factory=dict)
+]
-
-    def is_stale(self, now_epoch_ms: int) -> bool:
-        age_ms = now_epoch_ms - self.fetched_at_epoch_ms
-        return age_ms > self.max_age_seconds * 1000
-
-
-@dataclass(slots=True)
-class BrokerConfig:
-    store_root: Path
-    connections: list[ConnectionConfig] = field(default_factory=list)
-
-
-def dump_catalog_snapshot(snapshot: CatalogSnapshot) -> dict[str, Any]:
-    return {
-        "connection_id": snapshot.connection_id,
-        "fetched_at_epoch_ms": snapshot.fetched_at_epoch_ms,
-        "max_age_seconds": snapshot.max_age_seconds,
-        "nodes": [asdict(node) for node in snapshot.nodes],
-        "resources": [asdict(resource) for resource in snapshot.resources],
-        "prompts": [asdict(prompt) for prompt in snapshot.prompts],
-        "metadata": snapshot.metadata,
-    }

src/wf_mcp/storage/store.py

@@ -4,14 +4,16 @@ import json
 from pathlib import Path
 
 from wf_api.auth import AuthRecord as NeutralAuthRecord
+from wf_mcp.capabilities import (
+    CatalogNodeEntry,
+    CatalogPromptEntry,
+    CatalogResourceEntry,
+)
 
 from ..auth import mcp_auth_from_neutral, neutral_auth_from_mcp
 from ..connections import parse_connection_id
 from ..models import (
     AuthRecord,
-    CatalogNodeEntry,
-    CatalogPromptEntry,
-    CatalogResourceEntry,
     CatalogSnapshot,
     dump_catalog_snapshot,
 )

tests/wf_mcp/service/test_content_access.py

@@ -40,7 +40,7 @@ def _make_content_access(
         connection_list_enabled=connection_service.list_enabled,
         connection_list_all=connection_service.list_all,
         tool_executor_for=upstream.tool_executor_for,
-        load_auth=upstream.load_auth,
+        load_auth=upstream.load_connection_auth,
         emit_event=events.record_event,
     )
     connection_service.bind_source_catalog(source_catalog)