refactor: split generic source registry mechanics

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
 
-
-class SourceRegistryModel(BaseModel):
-    """Base model for persisted source registry state; reject misspelled fields."""
-
-    model_config = ConfigDict(extra="forbid")
+if TYPE_CHECKING:
+    from .models import ConnectionConfig
 
 
-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()
-        for source in self.sources:
-            if source.id in seen:
-                raise ValueError(f"duplicate source id {source.id!r}")
-            seen.add(source.id)
+        validate_unique_source_ids(self.sources)
         return self
 
     def source_map(self) -> dict[str, McpSourceRegistryEntry]:
         return {source.id: source for source in self.sources}
 
 
-class SourceRegistryStore(Protocol):
-    """Persistence boundary for desired server-owned 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 SourceRegistryStore(GenericSourceRegistryStore[SourceRegistryFile], Protocol):
+    """MCP-specialized persistence boundary for desired source configuration."""
 
 
 class FileSourceRegistryStore:
     """Filesystem implementation for desired source registry state."""
 
     def __init__(self, root: Path) -> None:
-        self.root = root
-        self.root.mkdir(parents=True, exist_ok=True)
+        self._delegate = AtomicJsonRegistryStore(
+            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 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)
+        return self._delegate.load_registry()
 
     def save_registry(self, registry: SourceRegistryFile) -> None:
-        validated = SourceRegistryFile.model_validate(registry.model_dump(mode="json"))
-        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.
-        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)
+        self._delegate.save_registry(registry)
+
+
+def registry_entry_to_connection_config(
+    entry: McpSourceRegistryEntry,
+) -> ConnectionConfig:
+    """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",
 ]