lda.chat/lda-wf
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: add typed mcp source connection seam

Browse Source
This commit is contained in:
lda
2026-06-07 13:50:57 +07:00 Verified
parent f79741eb2a
commit 195a967527
27 changed files with 690 additions and 217 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/wf_mcp/service/conftest.py
+4 -4
View File
@@ -28,8 +28,8 @@ class ContentOnlyOutputAdapter(FakeAdapter):
async def list_tools(
self,
connection: ConnectionConfig,
auth: AuthRecord | None,
connection,
auth,
) -> list[DiscoveredTool]:
return [
DiscoveredTool(
@@ -51,8 +51,8 @@ class ContentOnlyOutputAdapter(FakeAdapter):
async def call_tool(
self,
connection: ConnectionConfig,
auth: AuthRecord | None,
connection,
auth,
tool_name: str,
payload: dict[str, Any],
) -> ToolCallResult:
tests/wf_mcp/service/test_upstream_transport.py
+19 -5
View File
@@ -16,6 +16,10 @@ from ..test_support import FakeAdapter, local_temp_root
from ..workflow_surface.conftest import echo_artifact
def _fake_transport_metadata() -> dict[str, object]:
return {"transport": "stdio", "command": "fake-mcp-server"}
def _transport(root: Path) -> UpstreamTransportService:
events: list[McpEvent] = []
return UpstreamTransportService(
@@ -74,7 +78,12 @@ async def test_upstream_transport_invokes_raw_method_and_records_events() -> Non
events: list[McpEvent] = []
connections = ConnectionRegistry()
connections.register(
ConnectionConfig(id="demo.personal", server="demo", account="personal")
ConnectionConfig(
id="demo.personal",
server="demo",
account="personal",
metadata=_fake_transport_metadata(),
)
)
transport = UpstreamTransportService(
auth_store=FileStore(local_temp_root() / "upstream_raw_method"),
@@ -100,7 +109,12 @@ async def test_upstream_transport_refreshes_catalog_directly() -> None:
events: list[McpEvent] = []
store = FileStore(local_temp_root() / "upstream_refresh")
connections = ConnectionRegistry()
connection = ConnectionConfig(id="demo.personal", server="demo", account="personal")
connection = ConnectionConfig(
id="demo.personal",
server="demo",
account="personal",
metadata=_fake_transport_metadata(),
)
connections.register(connection)
transport = UpstreamTransportService(
auth_store=store,
@@ -197,7 +211,7 @@ def test_upstream_load_connection_auth_prefers_auth_ref(tmp_path: Path) -> None:
id="github.work",
server="github",
account="work",
metadata={"auth_ref": "github.creds"},
metadata={**_fake_transport_metadata(), "auth_ref": "github.creds"},
)
assert service.load_connection_auth(connection) == AuthRecord(
@@ -264,9 +278,9 @@ async def test_upstream_transport_live_diagnostics_report_missing_auth_ref(
connections = ConnectionRegistry()
connection = ConnectionConfig(
id="github.work",
server="demo",
server="github",
account="work",
metadata={"auth_ref": "github.creds"},
metadata={**_fake_transport_metadata(), "auth_ref": "github.creds"},
)
connections.register(connection)
transport = UpstreamTransportService(
tests/wf_mcp/test_auth.py
+23 -13
View File
@@ -15,6 +15,8 @@ from wf_mcp.auth import (
from wf_mcp.models import AuthRecord as McpAuthRecord
from wf_mcp.models import ConnectionConfig
from wf_mcp.storage import FileStore
from wf_sources_mcp.connections import McpSourceConnection
from wf_sources_mcp.transports import StdioSourceTransport
def test_mcp_auth_from_neutral_preserves_scheme_and_payload() -> None:
@@ -125,22 +127,23 @@ def test_file_store_legacy_auth_methods_still_work(tmp_path: Path) -> None:
def test_auth_ref_for_connection_returns_string_only() -> None:
assert (
auth_ref_for_connection(
ConnectionConfig(
McpSourceConnection(
id="github.work",
server="github",
provider="github",
account="work",
metadata={"auth_ref": "github.creds"},
transport=StdioSourceTransport(command="placeholder"),
auth_ref="github.creds",
)
)
== "github.creds"
)
assert (
auth_ref_for_connection(
ConnectionConfig(
McpSourceConnection(
id="github.work",
server="github",
provider="github",
account="work",
metadata={"auth_ref": 123},
transport=StdioSourceTransport(command="placeholder"),
)
)
is None
@@ -148,11 +151,12 @@ def test_auth_ref_for_connection_returns_string_only() -> None:
def test_connection_auth_diagnostic_reports_missing_auth_ref() -> None:
connection = ConnectionConfig(
connection = McpSourceConnection(
id="github.work",
server="github",
provider="github",
account="work",
metadata={"auth_ref": "github.creds"},
transport=StdioSourceTransport(command="placeholder"),
auth_ref="github.creds",
)
diagnostic = connection_auth_diagnostic(
@@ -172,12 +176,18 @@ def test_connection_auth_diagnostic_reports_missing_auth_ref() -> None:
def test_connection_auth_diagnostic_ignores_absent_or_present_auth_ref() -> None:
no_ref = ConnectionConfig(id="github.work", server="github", account="work")
with_ref = ConnectionConfig(
no_ref = McpSourceConnection(
id="github.work",
server="github",
provider="github",
account="work",
metadata={"auth_ref": "github.creds"},
transport=StdioSourceTransport(command="placeholder"),
)
with_ref = McpSourceConnection(
id="github.work",
provider="github",
account="work",
transport=StdioSourceTransport(command="placeholder"),
auth_ref="github.creds",
)
auth = McpAuthRecord(
connection_id="github.creds",
tests/wf_mcp/test_sdk_adapter.py
+8 -2
View File
@@ -11,6 +11,7 @@ from wf_mcp.capabilities import DiscoveredTool
from wf_mcp.models import ConnectionConfig
from wf_mcp.sdk import McpSdkAdapter
from wf_mcp.storage import FileStore
from wf_sources_mcp.connections import mcp_source_connection_from_connection_config
from .test_support import (
everything_server_connection,
@@ -39,7 +40,7 @@ class _ToolsOnlyAdapter:
raise McpError(ErrorData(code=-32601, message="Method not found"))
async def get_connection_metadata(self, connection, auth):
return {"server": connection.server}
return {"server": getattr(connection, "provider", getattr(connection, "server", None))}
async def read_resource(self, connection, auth, uri):
raise NotImplementedError
@@ -133,9 +134,12 @@ def test_mcp_sdk_adapter_lists_and_calls_stdio_tool() -> None:
adapter = McpSdkAdapter()
try:
source_connection = mcp_source_connection_from_connection_config(
service.connections.get("fixture.personal")
)
result = asyncio.run(
adapter.call_tool(
connection=service.connections.get("fixture.personal"),
connection=source_connection,
auth=None,
tool_name="echo_tool",
payload={"text": "hello"},
@@ -184,6 +188,7 @@ def test_refresh_catalog_keeps_tools_when_optional_lists_are_unsupported() -> No
id="tools_only.personal",
server="tools_only",
account="personal",
metadata={"transport": "stdio", "command": "fake-tools-only"},
)
)
service.register_adapter("tools_only", _ToolsOnlyAdapter())
@@ -205,6 +210,7 @@ def test_refresh_catalog_unwraps_taskgroup_method_not_found() -> None:
id="wrapped_tools_only.personal",
server="wrapped_tools_only",
account="personal",
metadata={"transport": "stdio", "command": "fake-tools-only"},
)
)
service.register_adapter("wrapped_tools_only", _WrappedToolsOnlyAdapter())
tests/wf_mcp/test_stateful_runtime.py
+6 -3
View File
@@ -16,6 +16,8 @@ from wf_mcp.runtime import McpRuntimePool, PersistentMcpSession
from wf_mcp.runtime.factory import PersistentSessionFactory
from wf_mcp.sdk import ToolCallResult
from wf_mcp.workflow import wrap_discovered_tool
from wf_sources_mcp.connections import McpSourceConnection
from wf_sources_mcp.transports import StdioSourceTransport
@dataclass(slots=True)
@@ -27,7 +29,7 @@ class FakeStatefulExecutor:
async def call_tool(
self,
connection: ConnectionConfig,
connection,
auth: AuthRecord | None,
tool_name: str,
payload: dict[str, Any],
@@ -120,10 +122,11 @@ def _tool(name: str) -> DiscoveredTool:
def test_generated_workflow_specs_share_injected_tool_executor() -> None:
"""Generated NodeSpecs use the injected executor, not a baked-in adapter."""
connection = ConnectionConfig(
connection = McpSourceConnection(
id="playwright.default",
server="playwright",
provider="playwright",
account="default",
transport=StdioSourceTransport(command="placeholder"),
)
executor = FakeStatefulExecutor()
navigate = wrap_discovered_tool(
tests/wf_mcp/test_support.py
+21 -21
View File
@@ -112,8 +112,8 @@ def everything_server_connection() -> ConnectionConfig | None:
class FakeAdapter:
async def list_tools(
self,
connection: ConnectionConfig,
auth: AuthRecord | None,
connection,
auth,
) -> list[DiscoveredTool]:
return [
DiscoveredTool(
@@ -145,8 +145,8 @@ class FakeAdapter:
async def list_resources(
self,
connection: ConnectionConfig,
auth: AuthRecord | None,
connection,
auth,
) -> list[DiscoveredResource]:
return [
DiscoveredResource(
@@ -161,8 +161,8 @@ class FakeAdapter:
async def list_prompts(
self,
connection: ConnectionConfig,
auth: AuthRecord | None,
connection,
auth,
) -> list[DiscoveredPrompt]:
return [
DiscoveredPrompt(
@@ -182,19 +182,19 @@ class FakeAdapter:
async def get_connection_metadata(
self,
connection: ConnectionConfig,
auth: AuthRecord | None,
connection,
auth,
) -> dict[str, Any]:
return {
"server": connection.server,
"server": getattr(connection, "provider", getattr(connection, "server", None)),
"account": connection.account,
"auth_scheme": auth.scheme if auth is not None else None,
}
async def read_resource(
self,
connection: ConnectionConfig,
auth: AuthRecord | None,
connection,
auth,
uri: str,
) -> dict[str, Any]:
if uri != "demo://docs/welcome":
@@ -211,8 +211,8 @@ class FakeAdapter:
async def get_prompt(
self,
connection: ConnectionConfig,
auth: AuthRecord | None,
connection,
auth,
prompt_name: str,
arguments: dict[str, str] | None = None,
) -> dict[str, Any]:
@@ -234,8 +234,8 @@ class FakeAdapter:
async def invoke_method(
self,
connection: ConnectionConfig,
auth: AuthRecord | None,
connection,
auth,
method: str,
params: dict[str, Any] | None = None,
) -> dict[str, Any]:
@@ -247,8 +247,8 @@ class FakeAdapter:
async def send_notification(
self,
connection: ConnectionConfig,
auth: AuthRecord | None,
connection,
auth,
method: str,
params: dict[str, Any] | None = None,
) -> None:
@@ -256,8 +256,8 @@ class FakeAdapter:
async def call_tool(
self,
connection: ConnectionConfig,
auth: AuthRecord | None,
connection,
auth,
tool_name: str,
payload: dict[str, Any],
) -> ToolCallResult:
@@ -272,8 +272,8 @@ class FakeAdapter:
class FailingDiscoveryAdapter(FakeAdapter):
async def list_tools(
self,
connection: ConnectionConfig,
auth: AuthRecord | None,
connection,
auth,
) -> list[DiscoveredTool]:
raise PermissionError("Access is denied")
tests/wf_mcp/test_workflow_wrappers.py
+8 -4
View File
@@ -10,6 +10,8 @@ from wf_mcp.models import AuthRecord, ConnectionConfig
from wf_mcp.runtime import ToolExecutor
from wf_mcp.sdk import ToolCallResult
from wf_mcp.workflow import wrap_discovered_tool
from wf_sources_mcp.connections import McpSourceConnection
from wf_sources_mcp.transports import StdioSourceTransport
class RecordingAdapter:
@@ -54,10 +56,11 @@ class TextContentAdapter:
def test_discovered_tool_wrapper_omits_unset_optional_arguments() -> None:
adapter = RecordingAdapter()
spec = wrap_discovered_tool(
connection=ConnectionConfig(
connection=McpSourceConnection(
id="playwright.default",
server="playwright",
provider="playwright",
account="default",
transport=StdioSourceTransport(command="placeholder"),
),
auth=None,
executor=cast(ToolExecutor, adapter),
@@ -92,10 +95,11 @@ def test_discovered_tool_wrapper_omits_unset_optional_arguments() -> None:
def test_discovered_tool_wrapper_preserves_raw_mcp_content_output() -> None:
spec = wrap_discovered_tool(
connection=ConnectionConfig(
connection=McpSourceConnection(
id="everything.default",
server="everything",
provider="everything",
account="default",
transport=StdioSourceTransport(command="placeholder"),
),
auth=None,
executor=cast(ToolExecutor, TextContentAdapter()),
tests/wf_mcp/workflow_surface/conftest.py
+19 -19
View File
@@ -64,8 +64,8 @@ class ContentOnlyOutputAdapter:
async def list_tools(
self,
connection: ConnectionConfig,
auth: AuthRecord | None,
connection,
auth,
) -> list[DiscoveredTool]:
return [
DiscoveredTool(
@@ -87,29 +87,29 @@ class ContentOnlyOutputAdapter:
async def list_resources(
self,
connection: ConnectionConfig,
auth: AuthRecord | None,
connection,
auth,
) -> list[Any]:
return []
async def list_prompts(
self,
connection: ConnectionConfig,
auth: AuthRecord | None,
connection,
auth,
) -> list[Any]:
return []
async def get_connection_metadata(
self,
connection: ConnectionConfig,
auth: AuthRecord | None,
connection,
auth,
) -> dict[str, Any]:
return {"server": connection.server}
return {"server": getattr(connection, "provider", getattr(connection, "server", None))}
async def call_tool(
self,
connection: ConnectionConfig,
auth: AuthRecord | None,
connection,
auth,
tool_name: str,
payload: dict[str, Any],
) -> ToolCallResult:
@@ -121,16 +121,16 @@ class ContentOnlyOutputAdapter:
async def read_resource(
self,
connection: ConnectionConfig,
auth: AuthRecord | None,
connection,
auth,
uri: str,
) -> dict[str, Any]:
raise KeyError(uri)
async def get_prompt(
self,
connection: ConnectionConfig,
auth: AuthRecord | None,
connection,
auth,
prompt_name: str,
arguments: dict[str, str] | None = None,
) -> dict[str, Any]:
@@ -138,8 +138,8 @@ class ContentOnlyOutputAdapter:
async def invoke_method(
self,
connection: ConnectionConfig,
auth: AuthRecord | None,
connection,
auth,
method: str,
params: dict[str, Any] | None = None,
) -> dict[str, Any]:
@@ -147,8 +147,8 @@ class ContentOnlyOutputAdapter:
async def send_notification(
self,
connection: ConnectionConfig,
auth: AuthRecord | None,
connection,
auth,
method: str,
params: dict[str, Any] | None = None,
) -> None:
tests/wf_sources_mcp/test_connections.py
+200
View File
@@ -0,0 +1,200 @@
from typing import Protocol
import pytest
from pydantic import TypeAdapter
from wf_sources_mcp.auth import auth_ref_for_connection
from wf_sources_mcp.connections import (
McpSourceConnection,
mcp_source_connection_from_connection_config,
mcp_source_connection_from_registry_entry,
)
from wf_sources_mcp.ids import (
CONNECTION_ID_PATTERN,
RESERVED_CONNECTION_IDS,
parse_connection_id,
)
from wf_sources_mcp.sdk import BackendAdapter, ToolExecutor
from wf_sources_mcp.source_registry import McpSourceRegistryEntry
from wf_sources_mcp.transports import (
HttpSourceTransport,
SourceTransport,
StdioSourceTransport,
)
def test_stdio_source_transport_is_typed() -> None:
transport = StdioSourceTransport(
command="uvx",
args=("mcp-server",),
env={"TOKEN": "x"},
)
assert transport.kind == "stdio"
assert transport.command == "uvx"
assert transport.args == ("mcp-server",)
assert transport.env == {"TOKEN": "x"}
def test_http_source_transport_is_typed() -> None:
transport = HttpSourceTransport(url="http://127.0.0.1:8000/mcp")
assert transport.kind == "http"
assert str(transport.url) == "http://127.0.0.1:8000/mcp"
def test_source_transport_discriminated_union_parses() -> None:
adapter = TypeAdapter(SourceTransport)
transport = adapter.validate_python(
{"kind": "stdio", "command": "pnpx", "args": ["-y", "server"]}
)
assert isinstance(transport, StdioSourceTransport)
assert transport.args == ("-y", "server")
def test_parse_connection_id_splits_provider_and_account() -> None:
assert parse_connection_id("github.work") == ("github", "work")
@pytest.mark.parametrize(
"source_id",
["github", ".github.work", "github.", "github/work", "github work"],
)
def test_parse_connection_id_rejects_unsafe_or_unqualified_ids(source_id: str) -> None:
with pytest.raises(ValueError):
parse_connection_id(source_id)
def test_reserved_connection_ids_are_source_provider_constants() -> None:
assert "wf.admin" in RESERVED_CONNECTION_IDS
assert "wf.mcp" in RESERVED_CONNECTION_IDS
assert CONNECTION_ID_PATTERN.startswith("^")
def test_mcp_source_connection_from_registry_entry() -> None:
entry = McpSourceRegistryEntry.model_validate(
{
"id": "github.work",
"provider": "github",
"account": "work",
"profile": "engineering",
"transport": {
"kind": "stdio",
"command": "uvx",
"args": ["github-mcp"],
"env": {"A": "B"},
},
"auth_ref": "github.token",
"metadata": {"team": "platform"},
}
)
connection = mcp_source_connection_from_registry_entry(entry)
assert connection == McpSourceConnection(
id="github.work",
provider="github",
account="work",
enabled=True,
profile="engineering",
transport=StdioSourceTransport(
command="uvx",
args=("github-mcp",),
env={"A": "B"},
),
auth_ref="github.token",
metadata={"team": "platform"},
)
def test_mcp_source_connection_from_legacy_connection_config_stdio() -> None:
from wf_mcp.broker.models import ConnectionConfig
legacy = ConnectionConfig(
id="github.work",
server="github",
account="work",
enabled=False,
metadata={
"transport": "stdio",
"command": "uvx",
"args": ["github-mcp"],
"env": {"A": "B"},
"auth_ref": "github.token",
"profile": "engineering",
"source_registry": True,
"team": "platform",
},
)
connection = mcp_source_connection_from_connection_config(legacy)
assert connection.id == "github.work"
assert connection.provider == "github"
assert connection.account == "work"
assert connection.enabled is False
assert connection.profile == "engineering"
assert connection.auth_ref == "github.token"
assert connection.metadata == {"source_registry": True, "team": "platform"}
assert isinstance(connection.transport, StdioSourceTransport)
assert connection.transport.command == "uvx"
assert connection.transport.args == ("github-mcp",)
def test_mcp_source_connection_from_legacy_connection_config_http() -> None:
from wf_mcp.broker.models import ConnectionConfig
legacy = ConnectionConfig(
id="github.work",
server="github",
account="work",
metadata={
"transport": "streamable_http",
"url": "http://127.0.0.1:8000/mcp",
"headers": {"X-Test": "yes"},
},
)
connection = mcp_source_connection_from_connection_config(legacy)
assert isinstance(connection.transport, HttpSourceTransport)
assert str(connection.transport.url) == "http://127.0.0.1:8000/mcp"
assert connection.transport.headers == {"X-Test": "yes"}
def test_mcp_source_connection_rejects_missing_legacy_transport() -> None:
from wf_mcp.broker.models import ConnectionConfig
legacy = ConnectionConfig(
id="github.work",
server="github",
account="work",
metadata={},
)
with pytest.raises(ValueError, match="requires metadata.transport"):
mcp_source_connection_from_connection_config(legacy)
class _ConnectionLike(Protocol):
id: str
auth_ref: str | None
def test_auth_ref_for_typed_mcp_source_connection() -> None:
connection = McpSourceConnection(
id="github.work",
provider="github",
account="work",
transport=StdioSourceTransport(command="uvx"),
auth_ref="github.token",
)
assert auth_ref_for_connection(connection) == "github.token"
def test_sdk_protocols_are_importable_without_broker_connection_config() -> None:
assert BackendAdapter is not None
assert ToolExecutor is not None
tests/wf_sources_mcp/test_sdk_protocols.py
+16 -5
View File
@@ -3,16 +3,17 @@ from __future__ import annotations
from dataclasses import is_dataclass
from typing import cast
from wf_mcp.broker.models import ConnectionConfig
from wf_sources_mcp.auth import AuthRecord
from wf_sources_mcp.catalog import DiscoveredTool
from wf_sources_mcp.connections import McpSourceConnection
from wf_sources_mcp.sdk import BackendAdapter, ToolCallResult, ToolExecutor
from wf_sources_mcp.transports import StdioSourceTransport
class EchoAdapter:
async def list_tools(
self,
connection: ConnectionConfig,
connection: McpSourceConnection,
auth: AuthRecord | None,
) -> list[DiscoveredTool]:
return [
@@ -27,7 +28,7 @@ class EchoAdapter:
async def call_tool(
self,
connection: ConnectionConfig,
connection: McpSourceConnection,
auth: AuthRecord | None,
tool_name: str,
payload: dict[str, object],
@@ -46,7 +47,12 @@ def test_tool_call_result_is_slots_dataclass_with_empty_defaults() -> None:
async def test_backend_adapter_protocol_can_describe_tool_listing() -> None:
adapter = cast(BackendAdapter, EchoAdapter())
tools = await adapter.list_tools(
ConnectionConfig(id="demo.default", server="demo", account="default"),
McpSourceConnection(
id="demo.default",
provider="demo",
account="default",
transport=StdioSourceTransport(command="echo"),
),
None,
)
@@ -56,7 +62,12 @@ async def test_backend_adapter_protocol_can_describe_tool_listing() -> None:
async def test_tool_executor_protocol_can_describe_tool_calls() -> None:
executor = cast(ToolExecutor, EchoAdapter())
result = await executor.call_tool(
ConnectionConfig(id="demo.default", server="demo", account="default"),
McpSourceConnection(
id="demo.default",
provider="demo",
account="default",
transport=StdioSourceTransport(command="echo"),
),
None,
"echo",
{"message": "hello"},