refactor: split generic source registry mechanics

src/wf_api/source_registry.py

@@ -0,0 +1,107 @@
+from __future__ import annotations
+
+import json
+import re
+from collections.abc import Callable, Sequence
+from pathlib import Path
+from typing import Generic, Protocol, TypeVar
+from uuid import uuid4
+
+from pydantic import BaseModel, ConfigDict
+
+
+class SourceRegistryBaseModel(BaseModel):
+    """Base model for persisted source registry state; reject misspelled fields."""
+
+    model_config = ConfigDict(extra="forbid")
+
+
+SOURCE_REGISTRY_ID_PATTERN = r"^[A-Za-z0-9_][A-Za-z0-9_.-]*$"
+
+
+def validate_source_registry_id(value: str) -> str:
+    """Validate ids that are safe as registry keys and filesystem path segments.
+
+    This helper intentionally does not parse provider/account meaning. MCP can
+    layer stricter `parse_connection_id` validation on top while other future
+    source families can reuse the safe-id rule.
+    """
+
+    if not re.fullmatch(SOURCE_REGISTRY_ID_PATTERN, value):
+        raise ValueError(
+            "source id must start with alphanumeric or underscore and contain "
+            "only [A-Za-z0-9_.-]"
+        )
+    return value
+
+
+def validate_unique_source_ids(entries: Sequence[object]) -> None:
+    """Reject duplicate `id` fields without owning the entry model shape."""
+
+    seen: set[str] = set()
+    for entry in entries:
+        source_id = getattr(entry, "id", None)
+        if not isinstance(source_id, str):
+            raise ValueError("source registry entries must expose string id")
+        if source_id in seen:
+            raise ValueError(f"duplicate source id {source_id!r}")
+        seen.add(source_id)
+
+
+RegistryT = TypeVar("RegistryT", bound=BaseModel)
+
+
+class SourceRegistryStore(Protocol[RegistryT]):
+    def load_registry(self) -> RegistryT: ...
+
+    def save_registry(self, registry: RegistryT) -> None: ...
+
+
+class AtomicJsonRegistryStore(Generic[RegistryT]):
+    """Filesystem implementation for small desired-registry documents."""
+
+    def __init__(
+        self,
+        root: Path,
+        *,
+        filename: str,
+        registry_type: type[RegistryT],
+        empty_factory: Callable[[], RegistryT],
+        corrupt_label: str,
+    ) -> None:
+        self.root = root
+        self.filename = filename
+        self.registry_type = registry_type
+        self.empty_factory = empty_factory
+        self.corrupt_label = corrupt_label
+        self.root.mkdir(parents=True, exist_ok=True)
+
+    @property
+    def path(self) -> Path:
+        return self.root / self.filename
+
+    def load_registry(self) -> RegistryT:
+        if not self.path.exists():
+            return self.empty_factory()
+        try:
+            data = json.loads(self.path.read_text(encoding="utf-8"))
+        except json.JSONDecodeError as exc:
+            raise ValueError(f"{self.corrupt_label} is corrupted: {self.path}") from exc
+        return self.registry_type.model_validate(data)
+
+    def save_registry(self, registry: RegistryT) -> None:
+        validated = self.registry_type.model_validate(registry.model_dump(mode="json"))
+        payload = json.dumps(validated.model_dump(mode="json"), indent=2)
+        tmp_path = self.path.with_name(f"{self.path.name}.{uuid4().hex}.tmp")
+        tmp_path.write_text(payload, encoding="utf-8")
+        tmp_path.replace(self.path)
+
+
+__all__ = [
+    "AtomicJsonRegistryStore",
+    "SourceRegistryBaseModel",
+    "SourceRegistryStore",
+    "SOURCE_REGISTRY_ID_PATTERN",
+    "validate_source_registry_id",
+    "validate_unique_source_ids",
+]

src/wf_mcp/source_registry.py

@@ -1,37 +1,38 @@
 from __future__ import annotations
 
-import json
+from typing import TYPE_CHECKING, Annotated, Literal, Protocol
+
 from pathlib import Path
-from typing import Annotated, Literal, Protocol
-from uuid import uuid4
 
 from pydantic import (
     AnyHttpUrl,
-    BaseModel,
-    ConfigDict,
     Field,
     field_validator,
     model_validator,
 )
 
+from wf_api.source_registry import (
+    AtomicJsonRegistryStore,
+    SourceRegistryBaseModel,
+    SourceRegistryStore as GenericSourceRegistryStore,
+    validate_unique_source_ids,
+)
+
 from .connections import parse_connection_id
 from .shared.names import RESERVED_CONNECTION_IDS
 
-
+if TYPE_CHECKING:
-class SourceRegistryModel(BaseModel):
+    from .models import ConnectionConfig
-    """Base model for persisted source registry state; reject misspelled fields."""
-
-    model_config = ConfigDict(extra="forbid")
 
 
-class StdioSourceTransport(SourceRegistryModel):
+class StdioSourceTransport(SourceRegistryBaseModel):
     kind: Literal["stdio"] = "stdio"
     command: str = Field(min_length=1)
     args: tuple[str, ...] = ()
     env: dict[str, str] = Field(default_factory=dict)
 
 
-class HttpSourceTransport(SourceRegistryModel):
+class HttpSourceTransport(SourceRegistryBaseModel):
     kind: Literal["http"] = "http"
     url: AnyHttpUrl
     headers: dict[str, str] = Field(default_factory=dict)
@@ -43,7 +44,7 @@ SourceTransport = Annotated[
 ]
 
 
-class McpSourceRegistryEntry(SourceRegistryModel):
+class McpSourceRegistryEntry(SourceRegistryBaseModel):
     """Desired MCP source configuration persisted by server-owned mutation."""
 
     id: str
@@ -65,63 +66,65 @@ class McpSourceRegistryEntry(SourceRegistryModel):
         return value
 
 
-class SourceRegistryFile(SourceRegistryModel):
+class SourceRegistryFile(SourceRegistryBaseModel):
     version: Literal[1] = 1
     sources: list[McpSourceRegistryEntry] = Field(default_factory=list)
 
     @model_validator(mode="after")
     def validate_unique_source_ids(self) -> SourceRegistryFile:
-        seen: set[str] = set()
+        validate_unique_source_ids(self.sources)
-        for source in self.sources:
-            if source.id in seen:
-                raise ValueError(f"duplicate source id {source.id!r}")
-            seen.add(source.id)
         return self
 
     def source_map(self) -> dict[str, McpSourceRegistryEntry]:
         return {source.id: source for source in self.sources}
 
 
-class SourceRegistryStore(Protocol):
+class SourceRegistryStore(GenericSourceRegistryStore[SourceRegistryFile], Protocol):
-    """Persistence boundary for desired server-owned source configuration."""
+    """MCP-specialized persistence boundary for desired source configuration."""
-
-    def load_registry(self) -> SourceRegistryFile:
-        """Return the stored registry, or an empty registry when absent."""
-        ...
-
-    def save_registry(self, registry: SourceRegistryFile) -> None:
-        """Persist one validated registry atomically."""
-        ...
 
 
 class FileSourceRegistryStore:
     """Filesystem implementation for desired source registry state."""
 
     def __init__(self, root: Path) -> None:
-        self.root = root
+        self._delegate = AtomicJsonRegistryStore(
-        self.root.mkdir(parents=True, exist_ok=True)
+            root,
+            filename="source_registry.json",
+            registry_type=SourceRegistryFile,
+            empty_factory=SourceRegistryFile,
+            corrupt_label="source registry file",
+        )
 
     @property
     def path(self) -> Path:
-        return self.root / "source_registry.json"
+        return self._delegate.path
 
     def load_registry(self) -> SourceRegistryFile:
-        if not self.path.exists():
+        return self._delegate.load_registry()
-            return SourceRegistryFile()
-        try:
-            data = json.loads(self.path.read_text(encoding="utf-8"))
-        except json.JSONDecodeError as exc:
-            raise ValueError(f"source registry file is corrupted: {self.path}") from exc
-        return SourceRegistryFile.model_validate(data)
 
     def save_registry(self, registry: SourceRegistryFile) -> None:
-        validated = SourceRegistryFile.model_validate(registry.model_dump(mode="json"))
+        self._delegate.save_registry(registry)
-        payload = json.dumps(validated.model_dump(mode="json"), indent=2)
+
-        # Use a unique temp file so multiple store objects pointing at the same
+
-        # root do not trample each other's pending writes before replacement.
+def registry_entry_to_connection_config(
-        tmp_path = self.path.with_name(f"{self.path.name}.{uuid4().hex}.tmp")
+    entry: McpSourceRegistryEntry,
-        tmp_path.write_text(payload, encoding="utf-8")
+) -> ConnectionConfig:
-        tmp_path.replace(self.path)
+    """Convert a registry entry to a broker connection config."""
+    from .models import ConnectionConfig
+
+    return ConnectionConfig(
+        id=entry.id,
+        server=entry.provider,
+        account=entry.account,
+        enabled=entry.enabled,
+        metadata={
+            **entry.metadata,
+            "auth_ref": entry.auth_ref,
+            "profile": entry.profile,
+            "transport": entry.transport.model_dump(mode="json"),
+            "source_registry": True,
+        },
+    )
 
 
 __all__ = [
@@ -132,4 +135,5 @@ __all__ = [
     "SourceRegistryStore",
     "SourceTransport",
     "StdioSourceTransport",
+    "registry_entry_to_connection_config",
 ]

tests/wf_api/test_source_registry.py

@@ -0,0 +1,108 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+from pydantic import Field
+
+from wf_api.source_registry import (
+    AtomicJsonRegistryStore,
+    SourceRegistryBaseModel,
+    validate_source_registry_id,
+    validate_unique_source_ids,
+)
+
+
+class FakeEntry(SourceRegistryBaseModel):
+    id: str
+
+
+class FakeRegistry(SourceRegistryBaseModel):
+    version: int = 1
+    entries: list[FakeEntry] = Field(default_factory=list)
+
+
+def test_validate_source_registry_id_accepts_normal_ids() -> None:
+    assert validate_source_registry_id("github.work") == "github.work"
+    assert validate_source_registry_id("my_source") == "my_source"
+    assert validate_source_registry_id("test-v1") == "test-v1"
+
+
+def test_validate_source_registry_id_rejects_unsafe_ids() -> None:
+    with pytest.raises(ValueError, match="source id must start"):
+        validate_source_registry_id("../bad")
+    with pytest.raises(ValueError, match="source id must start"):
+        validate_source_registry_id("has space")
+    with pytest.raises(ValueError, match="source id must start"):
+        validate_source_registry_id("")
+
+
+def test_validate_unique_source_ids_accepts_unique_ids() -> None:
+    entries = [FakeEntry(id="a"), FakeEntry(id="b")]
+    validate_unique_source_ids(entries)
+
+
+def test_validate_unique_source_ids_rejects_duplicate_ids() -> None:
+    entries = [FakeEntry(id="a"), FakeEntry(id="a")]
+    with pytest.raises(ValueError, match="duplicate source id 'a'"):
+        validate_unique_source_ids(entries)
+
+
+def test_validate_unique_source_ids_rejects_non_string_id() -> None:
+    entries: list[object] = [FakeEntry(id="a"), object()]
+    with pytest.raises(
+        ValueError, match="source registry entries must expose string id"
+    ):
+        validate_unique_source_ids(entries)
+
+
+def test_atomic_json_registry_store_loads_empty_when_missing(tmp_path: Path) -> None:
+    store = AtomicJsonRegistryStore(
+        tmp_path,
+        filename="registry.json",
+        registry_type=FakeRegistry,
+        empty_factory=FakeRegistry,
+        corrupt_label="test registry",
+    )
+
+    registry = store.load_registry()
+
+    assert registry.version == 1
+    assert registry.entries == []
+    assert store.path == tmp_path / "registry.json"
+
+
+def test_atomic_json_registry_store_round_trips(tmp_path: Path) -> None:
+    store = AtomicJsonRegistryStore(
+        tmp_path,
+        filename="registry.json",
+        registry_type=FakeRegistry,
+        empty_factory=FakeRegistry,
+        corrupt_label="test registry",
+    )
+    registry = FakeRegistry(entries=[FakeEntry(id="test.entry")])
+
+    store.save_registry(registry)
+    loaded = store.load_registry()
+
+    assert len(loaded.entries) == 1
+    assert loaded.entries[0].id == "test.entry"
+
+
+def test_atomic_json_registry_store_rejects_corrupted_json(tmp_path: Path) -> None:
+    store = AtomicJsonRegistryStore(
+        tmp_path,
+        filename="registry.json",
+        registry_type=FakeRegistry,
+        empty_factory=FakeRegistry,
+        corrupt_label="test registry",
+    )
+    store.path.write_text("not json{{{", encoding="utf-8")
+
+    with pytest.raises(ValueError, match="corrupted"):
+        store.load_registry()
+
+
+def test_source_registry_base_model_rejects_extra_fields() -> None:
+    with pytest.raises(ValueError, match="Extra inputs are not permitted"):
+        SourceRegistryBaseModel.model_validate({"unknown_field": "value"})

tests/wf_mcp/test_source_registry.py

@@ -10,6 +10,7 @@ from wf_mcp.source_registry import (
     McpSourceRegistryEntry,
     SourceRegistryFile,
     StdioSourceTransport,
+    registry_entry_to_connection_config,
 )
 
 
@@ -52,33 +53,11 @@ def test_source_registry_accepts_http_transport() -> None:
     assert str(entry.transport.url) == "https://example.test/mcp"
 
 
-def test_source_registry_rejects_duplicate_ids() -> None:
-    with pytest.raises(ValueError, match="duplicate source id 'github.work'"):
-        SourceRegistryFile(sources=[_entry("github.work"), _entry("github.work")])
-
-
 def test_source_registry_rejects_reserved_ids() -> None:
     with pytest.raises(ValueError, match="reserved"):
         _entry("wf.admin")
 
 
-def test_source_registry_rejects_unsafe_ids() -> None:
-    with pytest.raises(ValueError, match="connection id"):
-        _entry("../bad")
-
-
-def test_file_source_registry_store_loads_empty_registry_when_missing(
-    tmp_path: Path,
-) -> None:
-    store = FileSourceRegistryStore(tmp_path)
-
-    registry = store.load_registry()
-
-    assert registry.version == 1
-    assert registry.sources == []
-    assert store.path == tmp_path / "source_registry.json"
-
-
 def test_file_source_registry_store_round_trips_registry(tmp_path: Path) -> None:
     store = FileSourceRegistryStore(tmp_path)
     registry = SourceRegistryFile(sources=[_entry("github.work")])
@@ -102,9 +81,38 @@ def test_file_source_registry_store_validates_loaded_registry(tmp_path: Path) ->
         store.load_registry()
 
 
-def test_file_source_registry_store_rejects_corrupted_json(tmp_path: Path) -> None:
+def test_registry_entry_to_connection_config_preserves_identity() -> None:
-    store = FileSourceRegistryStore(tmp_path)
+    entry = _entry()
-    store.path.write_text("not json{{{", encoding="utf-8")
+    config = registry_entry_to_connection_config(entry)
 
-    with pytest.raises(ValueError, match="corrupted"):
+    assert config.id == "github.work"
-        store.load_registry()
+    assert config.server == "github"
+    assert config.account == "work"
+    assert config.enabled is True
+
+
+def test_registry_entry_to_connection_config_preserves_transport_metadata() -> None:
+    entry = _entry()
+    entry.auth_ref = "github.work.auth"
+    config = registry_entry_to_connection_config(entry)
+
+    assert config.metadata["auth_ref"] == "github.work.auth"
+    assert config.metadata["profile"] is None
+    assert config.metadata["transport"]["kind"] == "stdio"
+    assert config.metadata["transport"]["command"] == "npx"
+    assert config.metadata["source_registry"] is True
+
+
+def test_registry_entry_to_connection_config_preserves_user_metadata() -> None:
+    entry = _entry()
+    config = registry_entry_to_connection_config(entry)
+
+    assert config.metadata["purpose"] == "tests"
+
+
+def test_registry_entry_to_connection_config_disabled_entry() -> None:
+    entry = _entry()
+    entry.enabled = False
+    config = registry_entry_to_connection_config(entry)
+
+    assert config.enabled is False