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
src/wf_mcp/broker/discovery.py
+17 -5
View File
@@ -9,6 +9,10 @@ from mcp.types import METHOD_NOT_FOUND
from wf_authoring import NodeSpec
from wf_sources_mcp.catalog import DiscoveredPrompt, DiscoveredResource, DiscoveredTool
from wf_sources_mcp.connections import (
McpSourceConnection,
mcp_source_connection_from_connection_config,
)
from wf_sources_mcp.sdk import BackendAdapter, ToolExecutor
from ..auth import AuthRecord
@@ -34,14 +38,18 @@ async def discover_connection_capabilities(
auth: AuthRecord | None,
adapter: BackendAdapter,
) -> DiscoveredConnectionCapabilities:
tools = await adapter.list_tools(connection, auth)
# Compatibility boundary: broker callers still pass ConnectionConfig. Runtime
# internals use McpSourceConnection so the session code can move to
# wf_sources_mcp in a later slice.
source_connection = mcp_source_connection_from_connection_config(connection)
tools = await adapter.list_tools(source_connection, auth)
resources = await _list_optional_capabilities(
lambda: adapter.list_resources(connection, auth)
lambda: adapter.list_resources(source_connection, auth)
)
prompts = await _list_optional_capabilities(
lambda: adapter.list_prompts(connection, auth)
lambda: adapter.list_prompts(source_connection, auth)
)
metadata = await adapter.get_connection_metadata(connection, auth)
metadata = await adapter.get_connection_metadata(source_connection, auth)
return DiscoveredConnectionCapabilities(
tools=tools,
resources=resources,
@@ -77,9 +85,13 @@ def specs_from_discovered_tools(
tools: list[DiscoveredTool],
emit_event: Callable[[McpEvent], None] | None = None,
) -> list[NodeSpec[Any, Any]]:
# Compatibility boundary: broker callers still pass ConnectionConfig. Runtime
# internals use McpSourceConnection so the session code can move to
# wf_sources_mcp in a later slice.
source_connection = mcp_source_connection_from_connection_config(connection)
return [
wrap_discovered_tool(
connection=connection,
connection=source_connection,
auth=auth,
executor=executor,
tool=tool,
src/wf_mcp/broker/service/source_catalog.py
+4 -1
View File
@@ -23,6 +23,7 @@ from wf_sources_mcp.catalog import (
CatalogPromptEntry,
CatalogResourceEntry,
)
from wf_sources_mcp.connections import mcp_source_connection_from_connection_config
from wf_sources_mcp.sdk import ToolExecutor
from wf_sources_mcp.storage import CatalogStore
@@ -273,8 +274,10 @@ class SourceCatalogService:
async def invoke_tool(payload: BaseModel) -> NodeReturn[BaseModel]:
connection = self.connection_lookup(entry.connection_id)
auth = self.load_auth(connection)
# Compatibility boundary: broker callers still pass ConnectionConfig.
source_connection = mcp_source_connection_from_connection_config(connection)
result = await self.tool_executor_for(connection).call_tool(
connection,
source_connection,
auth,
entry.local_name,
payload.model_dump(exclude_unset=True),
src/wf_mcp/broker/service/source_registry_admin.py
+6 -1
View File
@@ -5,6 +5,7 @@ from dataclasses import dataclass, field
from typing import Any
from wf_api.source_registry_admin import WorkflowSourceRegistryMutationProvider
from wf_sources_mcp.connections import mcp_source_connection_from_connection_config
from wf_sources_mcp.source_registry import (
McpSourceRegistryEntry,
SourceRegistryFile,
@@ -162,8 +163,12 @@ class SourceRegistryAdminProvider(WorkflowSourceRegistryMutationProvider):
auth_diagnostics = []
if self.load_auth is not None:
for source_id in sorted(after):
# Compatibility boundary: broker callers still pass ConnectionConfig.
source_connection = mcp_source_connection_from_connection_config(
after[source_id]
)
diagnostic = connection_auth_diagnostic(
after[source_id],
source_connection,
load_auth_ref=self.load_auth,
)
if diagnostic is not None:
src/wf_mcp/broker/service/upstream_transport.py
+23 -6
View File
@@ -27,6 +27,10 @@ from wf_mcp.models import ConnectionConfig
from wf_mcp.shared.errors import error_payload
from wf_sources_mcp.auth import AuthRecord, connection_auth_diagnostic
from wf_sources_mcp.catalog.models import CatalogSnapshot
from wf_sources_mcp.connections import (
McpSourceConnection,
mcp_source_connection_from_connection_config,
)
from wf_sources_mcp.sdk import BackendAdapter, ToolExecutor
from wf_sources_mcp.storage import AuthStore, CatalogStore
@@ -74,6 +78,9 @@ class UpstreamTransportService:
old compatibility surface has no callers.
"""
# Compatibility boundary: broker callers still pass ConnectionConfig.
# Check legacy metadata for auth_ref first to avoid requiring transport
# metadata just for auth resolution.
auth_ref = connection.metadata.get("auth_ref")
if isinstance(auth_ref, str):
return self.load_auth(auth_ref)
@@ -98,6 +105,8 @@ class UpstreamTransportService:
) -> dict[str, Any]:
adapter = require_adapter(connection, self.adapters)
auth = self.load_connection_auth(connection)
# Compatibility boundary: broker callers still pass ConnectionConfig.
source_connection = mcp_source_connection_from_connection_config(connection)
self.event_sink(
make_event(
"resource_read_started",
@@ -106,7 +115,7 @@ class UpstreamTransportService:
payload={"uri": uri},
)
)
result = await adapter.read_resource(connection, auth, uri)
result = await adapter.read_resource(source_connection, auth, uri)
self.event_sink(
make_event(
"resource_read_completed",
@@ -126,6 +135,8 @@ class UpstreamTransportService:
) -> dict[str, Any]:
adapter = require_adapter(connection, self.adapters)
auth = self.load_connection_auth(connection)
# Compatibility boundary: broker callers still pass ConnectionConfig.
source_connection = mcp_source_connection_from_connection_config(connection)
self.event_sink(
make_event(
"prompt_get_started",
@@ -134,7 +145,7 @@ class UpstreamTransportService:
payload={"argument_keys": sorted((arguments or {}).keys())},
)
)
result = await adapter.get_prompt(connection, auth, local_name, arguments)
result = await adapter.get_prompt(source_connection, auth, local_name, arguments)
self.event_sink(
make_event(
"prompt_get_completed",
@@ -154,6 +165,8 @@ class UpstreamTransportService:
) -> dict[str, Any]:
adapter = require_adapter(connection, self.adapters)
auth = self.load_connection_auth(connection)
# Compatibility boundary: broker callers still pass ConnectionConfig.
source_connection = mcp_source_connection_from_connection_config(connection)
self.event_sink(
make_event(
"raw_method_started",
@@ -162,7 +175,7 @@ class UpstreamTransportService:
payload={"params": params or {}},
)
)
result = await adapter.invoke_method(connection, auth, method, params)
result = await adapter.invoke_method(source_connection, auth, method, params)
self.event_sink(
make_event(
"raw_method_completed",
@@ -182,6 +195,8 @@ class UpstreamTransportService:
) -> None:
adapter = require_adapter(connection, self.adapters)
auth = self.load_connection_auth(connection)
# Compatibility boundary: broker callers still pass ConnectionConfig.
source_connection = mcp_source_connection_from_connection_config(connection)
self.event_sink(
make_event(
"raw_notification_started",
@@ -190,7 +205,7 @@ class UpstreamTransportService:
payload={"params": params or {}},
)
)
await adapter.send_notification(connection, auth, method, params)
await adapter.send_notification(source_connection, auth, method, params)
self.event_sink(
make_event(
"raw_notification_completed",
@@ -309,8 +324,10 @@ class UpstreamTransportService:
)
)
continue
# Compatibility boundary: broker callers still pass ConnectionConfig.
source_connection = mcp_source_connection_from_connection_config(connection)
auth_diagnostic = connection_auth_diagnostic(
connection,
source_connection,
# The diagnostic helper passes the explicit auth_ref to this
# loader, matching load_connection_auth's auth_ref-first path.
load_auth_ref=self.load_auth,
@@ -323,7 +340,7 @@ class UpstreamTransportService:
adapter = require_adapter(connection, self.adapters)
auth = self.load_connection_auth(connection)
await asyncio.wait_for(
adapter.list_tools(connection, auth),
adapter.list_tools(source_connection, auth),
timeout=LIVE_SOURCE_CHECK_TIMEOUT_SECONDS,
)
except _LIVE_SOURCE_CHECK_FAILURES as exc:
src/wf_mcp/connections.py
+2 -20
View File
@@ -1,29 +1,11 @@
from __future__ import annotations
import re
from dataclasses import dataclass, field
from wf_sources_mcp.ids import CONNECTION_ID_PATTERN, parse_connection_id
from .models import ConnectionConfig
CONNECTION_ID_PATTERN = r"^[A-Za-z0-9_][A-Za-z0-9_.-]*$"
def parse_connection_id(connection_id: str) -> tuple[str, str]:
# Connection ids are logical source ids, but they also key persisted auth and
# catalog files. Keep this parser conservative so unsafe ids are rejected
# before they reach either registry or store boundaries.
if not re.fullmatch(CONNECTION_ID_PATTERN, connection_id):
raise ValueError(
"connection id must start with alphanumeric or underscore and contain "
"only [A-Za-z0-9_.-]"
)
if "." not in connection_id:
raise ValueError("connection id must look like '<server>.<account>'")
server, account = connection_id.split(".", 1)
if not server or not account:
raise ValueError("connection id must look like '<server>.<account>'")
return server, account
def qualify_node_name(connection_id: str, local_name: str) -> str:
parse_connection_id(connection_id)
src/wf_mcp/runtime/pool.py
+2 -2
View File
@@ -75,8 +75,8 @@ class McpRuntimePool:
async def call_tool(
self,
connection: ConnectionConfig,
auth: AuthRecord | None,
connection,
auth,
tool_name: str,
payload: dict[str, Any],
) -> ToolCallResult:
src/wf_mcp/sdk/adapter.py
+23 -28
View File
@@ -19,6 +19,7 @@ from pydantic import AnyUrl
from wf_sources_mcp.auth import AuthRecord, mcp_auth_env, mcp_auth_headers
from wf_sources_mcp.catalog import DiscoveredPrompt, DiscoveredResource, DiscoveredTool
from wf_sources_mcp.connections import McpSourceConnection
from wf_sources_mcp.sdk import BackendAdapter, ToolCallResult
from wf_sources_mcp.sdk.converters import (
prompt_to_discovered,
@@ -26,31 +27,26 @@ from wf_sources_mcp.sdk.converters import (
tool_result_to_call_result,
tool_to_discovered,
)
from ..models import ConnectionConfig
from wf_sources_mcp.transports import HttpSourceTransport, StdioSourceTransport
class McpSdkAdapter(BackendAdapter):
@asynccontextmanager
async def _session(
self,
connection: ConnectionConfig,
connection: McpSourceConnection,
auth: AuthRecord | None,
):
transport = connection.metadata.get("transport", "stdio")
if transport == "stdio":
command = connection.metadata["command"]
args = list(connection.metadata.get("args", []))
env = connection.metadata.get("env")
cwd = connection.metadata.get("cwd")
transport = connection.transport
if isinstance(transport, StdioSourceTransport):
auth_env = mcp_auth_env(auth)
env = dict(transport.env)
if auth_env:
env = {**(env or {}), **auth_env}
env = {**env, **auth_env}
params = StdioServerParameters(
command=command,
args=args,
command=transport.command,
args=list(transport.args),
env=env,
cwd=cwd,
)
async with stdio_client(params) as (read_stream, write_stream):
async with ClientSession(read_stream, write_stream) as session:
@@ -58,13 +54,12 @@ class McpSdkAdapter(BackendAdapter):
yield session
return
if transport == "streamable_http":
url = connection.metadata["url"]
if isinstance(transport, HttpSourceTransport):
headers = mcp_auth_headers(auth)
http_client = httpx.AsyncClient(headers=headers or None)
async with http_client:
async with streamable_http_client(
url,
str(transport.url),
http_client=http_client,
) as (read_stream, write_stream, _get_session_id):
async with ClientSession(read_stream, write_stream) as session:
@@ -72,11 +67,11 @@ class McpSdkAdapter(BackendAdapter):
yield session
return
raise ValueError(f"unsupported MCP transport {transport!r}")
raise ValueError(f"unsupported MCP transport {transport.kind!r}")
async def list_tools(
self,
connection: ConnectionConfig,
connection: McpSourceConnection,
auth: AuthRecord | None,
) -> list[DiscoveredTool]:
async with self._session(connection, auth) as session:
@@ -85,7 +80,7 @@ class McpSdkAdapter(BackendAdapter):
async def list_resources(
self,
connection: ConnectionConfig,
connection: McpSourceConnection,
auth: AuthRecord | None,
) -> list[DiscoveredResource]:
async with self._session(connection, auth) as session:
@@ -94,7 +89,7 @@ class McpSdkAdapter(BackendAdapter):
async def list_prompts(
self,
connection: ConnectionConfig,
connection: McpSourceConnection,
auth: AuthRecord | None,
) -> list[DiscoveredPrompt]:
async with self._session(connection, auth) as session:
@@ -103,17 +98,17 @@ class McpSdkAdapter(BackendAdapter):
async def get_connection_metadata(
self,
connection: ConnectionConfig,
connection: McpSourceConnection,
auth: AuthRecord | None,
) -> dict[str, Any]:
return {
"server": connection.server,
"transport": connection.metadata.get("transport", "stdio"),
"server": connection.provider,
"transport": connection.transport.kind,
}
async def read_resource(
self,
connection: ConnectionConfig,
connection: McpSourceConnection,
auth: AuthRecord | None,
uri: str,
) -> dict[str, Any]:
@@ -123,7 +118,7 @@ class McpSdkAdapter(BackendAdapter):
async def get_prompt(
self,
connection: ConnectionConfig,
connection: McpSourceConnection,
auth: AuthRecord | None,
prompt_name: str,
arguments: dict[str, str] | None = None,
@@ -134,7 +129,7 @@ class McpSdkAdapter(BackendAdapter):
async def invoke_method(
self,
connection: ConnectionConfig,
connection: McpSourceConnection,
auth: AuthRecord | None,
method: str,
params: dict[str, Any] | None = None,
@@ -148,7 +143,7 @@ class McpSdkAdapter(BackendAdapter):
async def send_notification(
self,
connection: ConnectionConfig,
connection: McpSourceConnection,
auth: AuthRecord | None,
method: str,
params: dict[str, Any] | None = None,
@@ -160,7 +155,7 @@ class McpSdkAdapter(BackendAdapter):
async def call_tool(
self,
connection: ConnectionConfig,
connection: McpSourceConnection,
auth: AuthRecord | None,
tool_name: str,
payload: dict[str, Any],
src/wf_mcp/shared/names.py
+2 -2
View File
@@ -21,9 +21,9 @@ if TYPE_CHECKING:
from fastmcp.resources.template import ResourceTemplate
from fastmcp.tools.base import Tool
from wf_sources_mcp.ids import RESERVED_CONNECTION_IDS
ADMIN_NAMESPACE = "wf.admin"
RESERVED_CONNECTION_IDS = frozenset({ADMIN_NAMESPACE, "wf.mcp"})
"""Source ids reserved by wf-mcp system capabilities."""
@dataclass(frozen=True, slots=True)
src/wf_mcp/workflow/wrappers.py
+3 -4
View File
@@ -9,12 +9,11 @@ from pydantic import BaseModel, ConfigDict, Field, create_model
from wf_authoring import NodeReturn, NodeSpec
from wf_core import RuntimeContext
from wf_mcp.broker.events import McpEvent, make_event
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 ToolExecutor
from ..auth import AuthRecord
from ..models import ConnectionConfig
_JSON_TYPE_MAP: dict[str, object] = {
"string": str,
"integer": int,
@@ -118,7 +117,7 @@ def _model_from_schema(name: str, schema: dict[str, Any]) -> type[BaseModel]:
def wrap_discovered_tool(
*,
connection: ConnectionConfig,
connection: McpSourceConnection,
auth: AuthRecord | None,
executor: ToolExecutor,
tool: DiscoveredTool,