refactor: move mcp catalog aggregation to wf_sources_mcp

2026-06-08 07:24:59 +07:00
parent 8ac44fdf2d
commit 0e20a42c38
12 changed files with 1064 additions and 177 deletions
@@ -0,0 +1,144 @@
+from __future__ import annotations
+
+import dataclasses
+
+from pydantic import BaseModel
+
+from wf_authoring import NodeSpec
+from wf_sources_mcp.catalog import (
+    CombinedCatalog,
+    DiscoveredPrompt,
+    DiscoveredResource,
+    snapshot_from_specs,
+)
+
+
+class EchoInput(BaseModel):
+    message: str
+
+
+class EchoOutput(BaseModel):
+    text: str
+
+
+async def _echo(payload: EchoInput) -> EchoOutput:
+    return EchoOutput(text=payload.message)
+
+
+def _echo_spec() -> NodeSpec[EchoInput, EchoOutput]:
+    return NodeSpec(
+        name="echo",
+        input_model=EchoInput,
+        output_model=EchoOutput,
+        outcomes=("ok",),
+        fn=_echo,
+        description="Echo message",
+        is_async=True,
+        input_schema_contract={
+            "type": "object",
+            "properties": {"message": {"type": "string"}},
+            "required": ["message"],
+        },
+        output_schema_contract={
+            "type": "object",
+            "properties": {"text": {"type": "string"}},
+        },
+    )
+
+
+def test_snapshot_from_specs_qualifies_nodes_resources_and_prompts() -> None:
+    snapshot = snapshot_from_specs(
+        "demo.default",
+        specs={"echo": _echo_spec()},
+        tool_display_names={"echo": "Echo Tool"},
+        resources=[
+            DiscoveredResource(
+                uri="demo://docs/guide",
+                name="guide",
+                title="Guide",
+                description="Read me",
+                mime_type="text/markdown",
+                metadata={"kind": "doc"},
+            )
+        ],
+        prompts=[
+            DiscoveredPrompt(
+                name="summarize",
+                title="Summarize",
+                description="Summarize text",
+                arguments=[{"name": "topic"}],
+                metadata={"kind": "prompt"},
+            )
+        ],
+        metadata={"source": "test"},
+        fetched_at_epoch_ms=123,
+        max_age_seconds=60,
+    )
+
+    assert snapshot.connection_id == "demo.default"
+    assert snapshot.nodes[0].qualified_name == "demo.default.echo"
+    assert snapshot.nodes[0].local_name == "echo"
+    assert snapshot.nodes[0].title == "Echo Tool"
+    assert snapshot.resources[0].qualified_name == "demo.default.guide"
+    assert snapshot.resources[0].uri == "demo://docs/guide"
+    assert snapshot.prompts[0].qualified_name == "demo.default.summarize"
+    assert snapshot.prompts[0].arguments == [{"name": "topic"}]
+    assert snapshot.metadata == {"source": "test"}
+
+
+def test_snapshot_from_specs_preserves_already_qualified_node_name() -> None:
+    spec = dataclasses.replace(_echo_spec(), name="demo.default.echo")
+
+    snapshot = snapshot_from_specs(
+        "demo.default",
+        specs={"demo.default.echo": spec},
+        fetched_at_epoch_ms=123,
+        max_age_seconds=60,
+    )
+
+    assert snapshot.nodes[0].qualified_name == "demo.default.echo"
+    assert snapshot.nodes[0].local_name == "echo"
+
+
+def test_combined_catalog_sorts_entries_and_serializes_payload() -> None:
+    first = snapshot_from_specs(
+        "zeta.default",
+        specs={"echo": _echo_spec()},
+        resources=[DiscoveredResource(uri="zeta://guide", name="guide", title=None, description=None)],
+        prompts=[DiscoveredPrompt(name="prompt", title=None, description=None)],
+        metadata={"order": "second"},
+        fetched_at_epoch_ms=2,
+        max_age_seconds=60,
+    )
+    second = snapshot_from_specs(
+        "alpha.default",
+        specs={"echo": _echo_spec()},
+        resources=[DiscoveredResource(uri="alpha://guide", name="guide", title=None, description=None)],
+        prompts=[DiscoveredPrompt(name="prompt", title=None, description=None)],
+        metadata={"order": "first"},
+        fetched_at_epoch_ms=1,
+        max_age_seconds=60,
+    )
+
+    catalog = CombinedCatalog(
+        snapshots={
+            first.connection_id: first,
+            second.connection_id: second,
+        }
+    )
+    payload = catalog.as_payload()
+
+    assert [entry.qualified_name for entry in catalog.entries()] == [
+        "alpha.default.echo",
+        "zeta.default.echo",
+    ]
+    assert catalog.find_resource("alpha.default.guide") is not None
+    assert catalog.find_prompt("zeta.default.prompt") is not None
+    assert [node["qualified_name"] for node in payload["nodes"]] == [
+        "alpha.default.echo",
+        "zeta.default.echo",
+    ]
+    assert [item["connection_id"] for item in payload["connections"]] == [
+        "alpha.default",
+        "zeta.default",
+    ]