refactor: split generic source registry mechanics
This commit is contained in:
@@ -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",
|
||||
]
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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"})
|
||||
@@ -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:
|
||||
store = FileSourceRegistryStore(tmp_path)
|
||||
store.path.write_text("not json{{{", encoding="utf-8")
|
||||
def test_registry_entry_to_connection_config_preserves_identity() -> None:
|
||||
entry = _entry()
|
||||
config = registry_entry_to_connection_config(entry)
|
||||
|
||||
with pytest.raises(ValueError, match="corrupted"):
|
||||
store.load_registry()
|
||||
assert config.id == "github.work"
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user