refactor: route all mcp operations through runtime pool

This commit is contained in:
lda
2026-06-08 18:29:22 +07:00 Verified
parent 929103219e
commit 0a8cc4190f
14 changed files with 961 additions and 30 deletions
+4
View File
@@ -260,6 +260,10 @@ implementation state.
resources, and prompts, and can route session-scoped `list_resources` and resources, and prompts, and can route session-scoped `list_resources` and
`list_prompts` through the owner task. Catalog refresh still uses one-shot `list_prompts` through the owner task. Catalog refresh still uses one-shot
adapter policy. adapter policy.
- Completed: persistent MCP runtime now implements the full upstream MCP
operation surface used by the one-shot adapter. Broker upstream operations
prefer the shared runtime pool when configured, with one-shot adapters
retained as fallback.
- Completed: MCP source catalog aggregation helpers (`CombinedCatalog` and - Completed: MCP source catalog aggregation helpers (`CombinedCatalog` and
`snapshot_from_specs`) now live in `wf_sources_mcp.catalog`; the old `snapshot_from_specs`) now live in `wf_sources_mcp.catalog`; the old
`wf_mcp.broker.catalog` path is a compatibility shim. `wf_mcp.broker.catalog` path is a compatibility shim.
@@ -157,7 +157,11 @@ First slices should move leaf modules only and leave `wf_mcp` re-export shims:
23. Complete: broker DTO construction removed from `wf_sources_mcp`. 23. Complete: broker DTO construction removed from `wf_sources_mcp`.
`wf_sources_mcp` accepts legacy-shaped inputs structurally, while `wf_sources_mcp` accepts legacy-shaped inputs structurally, while
`wf_mcp.source_registry` owns helpers that construct `ConnectionConfig`. `wf_mcp.source_registry` owns helpers that construct `ConnectionConfig`.
24. Upstream transport/discovery/session services. 24. Complete: `McpRuntimePool` implements the full MCP operation surface
(`list_tools`, resources, prompts, tools, raw methods, notifications, and
local metadata). Broker upstream operations prefer the persistent runtime
when configured and fall back to one-shot adapters.
25. Upstream transport/discovery/session services.
Each slice should add import-direction tests so the new source-provider package Each slice should add import-direction tests so the new source-provider package
does not depend on `wf_mcp.workflow_surface`, `wf_mcp.admin_surface`, does not depend on `wf_mcp.workflow_surface`, `wf_mcp.admin_surface`,
+35 -12
View File
@@ -158,7 +158,9 @@ class UpstreamTransportService:
) )
else: else:
adapter = require_adapter(connection, self.adapters) adapter = require_adapter(connection, self.adapters)
result = await adapter.get_prompt(source_connection, auth, local_name, arguments) result = await adapter.get_prompt(
source_connection, auth, local_name, arguments
)
self.event_sink( self.event_sink(
make_event( make_event(
"prompt_get_completed", "prompt_get_completed",
@@ -176,7 +178,6 @@ class UpstreamTransportService:
*, *,
params: dict[str, Any] | None = None, params: dict[str, Any] | None = None,
) -> dict[str, Any]: ) -> dict[str, Any]:
adapter = require_adapter(connection, self.adapters)
auth = self.load_connection_auth(connection) auth = self.load_connection_auth(connection)
# Compatibility boundary: broker callers still pass ConnectionConfig. # Compatibility boundary: broker callers still pass ConnectionConfig.
source_connection = mcp_source_connection_from_connection_config(connection) source_connection = mcp_source_connection_from_connection_config(connection)
@@ -188,7 +189,15 @@ class UpstreamTransportService:
payload={"params": params or {}}, payload={"params": params or {}},
) )
) )
result = await adapter.invoke_method(source_connection, auth, method, params) if self.stateful_runtime is not None:
result = await self.stateful_runtime.invoke_method(
source_connection, auth, method, params
)
else:
adapter = require_adapter(connection, self.adapters)
result = await adapter.invoke_method(
source_connection, auth, method, params
)
self.event_sink( self.event_sink(
make_event( make_event(
"raw_method_completed", "raw_method_completed",
@@ -206,7 +215,6 @@ class UpstreamTransportService:
*, *,
params: dict[str, Any] | None = None, params: dict[str, Any] | None = None,
) -> None: ) -> None:
adapter = require_adapter(connection, self.adapters)
auth = self.load_connection_auth(connection) auth = self.load_connection_auth(connection)
# Compatibility boundary: broker callers still pass ConnectionConfig. # Compatibility boundary: broker callers still pass ConnectionConfig.
source_connection = mcp_source_connection_from_connection_config(connection) source_connection = mcp_source_connection_from_connection_config(connection)
@@ -218,7 +226,13 @@ class UpstreamTransportService:
payload={"params": params or {}}, payload={"params": params or {}},
) )
) )
await adapter.send_notification(source_connection, auth, method, params) if self.stateful_runtime is not None:
await self.stateful_runtime.send_notification(
source_connection, auth, method, params
)
else:
adapter = require_adapter(connection, self.adapters)
await adapter.send_notification(source_connection, auth, method, params)
self.event_sink( self.event_sink(
make_event( make_event(
"raw_notification_completed", "raw_notification_completed",
@@ -246,12 +260,15 @@ class UpstreamTransportService:
) )
) )
try: try:
adapter = require_adapter(connection, self.adapters)
source_connection = mcp_source_connection_from_connection_config(connection) source_connection = mcp_source_connection_from_connection_config(connection)
if self.stateful_runtime is not None:
operations: BackendAdapter = self.stateful_runtime
else:
operations = require_adapter(connection, self.adapters)
capabilities = await discover_connection_capabilities( capabilities = await discover_connection_capabilities(
connection=source_connection, connection=source_connection,
auth=auth, auth=auth,
adapter=adapter, adapter=operations,
) )
specs = specs_from_discovered_tools( specs = specs_from_discovered_tools(
connection=connection, connection=connection,
@@ -351,12 +368,18 @@ class UpstreamTransportService:
diagnostics.append(auth_diagnostic) diagnostics.append(auth_diagnostic)
continue continue
try: try:
adapter = require_adapter(connection, self.adapters)
auth = self.load_connection_auth(connection) auth = self.load_connection_auth(connection)
await asyncio.wait_for( if self.stateful_runtime is not None:
adapter.list_tools(source_connection, auth), await asyncio.wait_for(
timeout=LIVE_SOURCE_CHECK_TIMEOUT_SECONDS, self.stateful_runtime.list_tools(source_connection, auth),
) timeout=LIVE_SOURCE_CHECK_TIMEOUT_SECONDS,
)
else:
adapter = require_adapter(connection, self.adapters)
await asyncio.wait_for(
adapter.list_tools(source_connection, auth),
timeout=LIVE_SOURCE_CHECK_TIMEOUT_SECONDS,
)
except _LIVE_SOURCE_CHECK_FAILURES as exc: except _LIVE_SOURCE_CHECK_FAILURES as exc:
diagnostics.append( diagnostics.append(
_source_unreachable_diagnostic( _source_unreachable_diagnostic(
+42 -1
View File
@@ -11,7 +11,7 @@ from typing import Any, Generic, TypeVar
from mcp.client.session import ClientSession from mcp.client.session import ClientSession
from wf_sources_mcp.auth import AuthRecord from wf_sources_mcp.auth import AuthRecord
from wf_sources_mcp.catalog import DiscoveredPrompt, DiscoveredResource from wf_sources_mcp.catalog import DiscoveredPrompt, DiscoveredResource, DiscoveredTool
from wf_sources_mcp.client import McpSourceClient, open_mcp_session from wf_sources_mcp.client import McpSourceClient, open_mcp_session
from wf_sources_mcp.connections import McpSourceConnection from wf_sources_mcp.connections import McpSourceConnection
from wf_sources_mcp.sdk import ToolCallResult from wf_sources_mcp.sdk import ToolCallResult
@@ -47,6 +47,10 @@ class PersistentSessionFactory:
get_prompt_callback=owner.get_prompt, get_prompt_callback=owner.get_prompt,
list_resources_callback=owner.list_resources, list_resources_callback=owner.list_resources,
list_prompts_callback=owner.list_prompts, list_prompts_callback=owner.list_prompts,
list_tools_callback=owner.list_tools,
get_connection_metadata_callback=owner.get_connection_metadata,
invoke_method_callback=owner.invoke_method,
send_notification_callback=owner.send_notification,
close_callback=owner.close, close_callback=owner.close,
) )
@@ -180,6 +184,43 @@ class _SessionOwner:
run=lambda client: client.list_prompts(), run=lambda client: client.list_prompts(),
) )
async def list_tools(self) -> list[DiscoveredTool]:
"""Submit tool listing through the generic owner-task operation queue."""
return await self.submit(
operation="list_tools",
run=lambda client: client.list_tools(),
)
async def get_connection_metadata(self) -> dict[str, Any]:
"""Return connection metadata computed locally without an upstream call."""
transport = self.connection.transport
return {
"server": self.connection.provider,
"transport": transport.kind if transport is not None else None,
}
async def invoke_method(
self,
method: str,
params: dict[str, Any] | None = None,
) -> dict[str, Any]:
"""Submit a raw method invocation through the generic owner-task operation queue."""
return await self.submit(
operation="invoke_method",
run=lambda client: client.invoke_method(method, params),
)
async def send_notification(
self,
method: str,
params: dict[str, Any] | None = None,
) -> None:
"""Submit a notification send through the generic owner-task operation queue."""
await self.submit(
operation="send_notification",
run=lambda client: client.send_notification(method, params),
)
async def close(self) -> None: async def close(self) -> None:
"""Ask the owner task to close the MCP transport in its own scope.""" """Ask the owner task to close the MCP transport in its own scope."""
task = self._task task = self._task
+37 -1
View File
@@ -7,7 +7,7 @@ from inspect import isawaitable
from typing import Any, cast from typing import Any, cast
from wf_sources_mcp.auth import AuthRecord from wf_sources_mcp.auth import AuthRecord
from wf_sources_mcp.catalog import DiscoveredPrompt, DiscoveredResource from wf_sources_mcp.catalog import DiscoveredPrompt, DiscoveredResource, DiscoveredTool
from wf_sources_mcp.connections import McpSourceConnection from wf_sources_mcp.connections import McpSourceConnection
from wf_sources_mcp.sdk import ToolCallResult from wf_sources_mcp.sdk import ToolCallResult
@@ -119,6 +119,42 @@ class McpRuntimePool:
session = await self.get_session(connection, auth) session = await self.get_session(connection, auth)
return await session.list_prompts() return await session.list_prompts()
async def list_tools(
self,
connection: McpSourceConnection,
auth: AuthRecord | None,
) -> list[DiscoveredTool]:
session = await self.get_session(connection, auth)
return await session.list_tools()
async def get_connection_metadata(
self,
connection: McpSourceConnection,
auth: AuthRecord | None,
) -> dict[str, Any]:
session = await self.get_session(connection, auth)
return await session.get_connection_metadata()
async def invoke_method(
self,
connection: McpSourceConnection,
auth: AuthRecord | None,
method: str,
params: dict[str, Any] | None = None,
) -> dict[str, Any]:
session = await self.get_session(connection, auth)
return await session.invoke_method(method, params)
async def send_notification(
self,
connection: McpSourceConnection,
auth: AuthRecord | None,
method: str,
params: dict[str, Any] | None = None,
) -> None:
session = await self.get_session(connection, auth)
await session.send_notification(method, params)
async def close_connection(self, connection_id: str) -> None: async def close_connection(self, connection_id: str) -> None:
current = self._sessions.pop(connection_id, None) current = self._sessions.pop(connection_id, None)
if current is not None: if current is not None:
+67 -1
View File
@@ -8,7 +8,7 @@ from mcp.client.session import ClientSession
from pydantic import AnyUrl from pydantic import AnyUrl
from wf_sources_mcp.auth import AuthRecord from wf_sources_mcp.auth import AuthRecord
from wf_sources_mcp.catalog import DiscoveredPrompt, DiscoveredResource from wf_sources_mcp.catalog import DiscoveredPrompt, DiscoveredResource, DiscoveredTool
from wf_sources_mcp.connections import McpSourceConnection from wf_sources_mcp.connections import McpSourceConnection
from wf_sources_mcp.sdk import ToolCallResult from wf_sources_mcp.sdk import ToolCallResult
from wf_sources_mcp.sdk.converters import tool_result_to_call_result from wf_sources_mcp.sdk.converters import tool_result_to_call_result
@@ -21,6 +21,10 @@ RawPromptGetter = Callable[
] ]
RawResourceLister = Callable[[], Awaitable[list[DiscoveredResource]]] RawResourceLister = Callable[[], Awaitable[list[DiscoveredResource]]]
RawPromptLister = Callable[[], Awaitable[list[DiscoveredPrompt]]] RawPromptLister = Callable[[], Awaitable[list[DiscoveredPrompt]]]
RawToolLister = Callable[[], Awaitable[list[DiscoveredTool]]]
RawMetadataGetter = Callable[[], Awaitable[dict[str, Any]]]
RawMethodInvoker = Callable[[str, dict[str, Any] | None], Awaitable[dict[str, Any]]]
RawNotificationSender = Callable[[str, dict[str, Any] | None], Awaitable[None]]
@dataclass(slots=True) @dataclass(slots=True)
@@ -41,6 +45,10 @@ class PersistentMcpSession:
get_prompt_callback: RawPromptGetter | None = None get_prompt_callback: RawPromptGetter | None = None
list_resources_callback: RawResourceLister | None = None list_resources_callback: RawResourceLister | None = None
list_prompts_callback: RawPromptLister | None = None list_prompts_callback: RawPromptLister | None = None
list_tools_callback: RawToolLister | None = None
get_connection_metadata_callback: RawMetadataGetter | None = None
invoke_method_callback: RawMethodInvoker | None = None
send_notification_callback: RawNotificationSender | None = None
close_callback: Callable[[], Awaitable[None]] | None = None close_callback: Callable[[], Awaitable[None]] | None = None
async def call_tool( async def call_tool(
@@ -97,6 +105,64 @@ class PersistentMcpSession:
return [prompt_to_discovered(prompt) for prompt in result.prompts] return [prompt_to_discovered(prompt) for prompt in result.prompts]
raise RuntimeError("persistent MCP session has no prompt list transport") raise RuntimeError("persistent MCP session has no prompt list transport")
async def list_tools(self) -> list[DiscoveredTool]:
"""List MCP tools through the owner task or injected session."""
if self.list_tools_callback is not None:
return await self.list_tools_callback()
if self.client is not None:
from wf_sources_mcp.sdk.converters import tool_to_discovered
result = await self.client.list_tools()
return [tool_to_discovered(tool) for tool in result.tools]
raise RuntimeError("persistent MCP session has no tools list transport")
async def get_connection_metadata(self) -> dict[str, Any]:
"""Return connection metadata from callback or local connection info."""
if self.get_connection_metadata_callback is not None:
return await self.get_connection_metadata_callback()
transport = self.connection.transport
return {
"server": self.connection.provider,
"transport": transport.kind if transport is not None else None,
}
async def invoke_method(
self,
method: str,
params: dict[str, Any] | None = None,
) -> dict[str, Any]:
"""Invoke a raw MCP method through the owner task or injected session."""
if self.invoke_method_callback is not None:
return await self.invoke_method_callback(method, params)
if self.client is not None:
from mcp import ClientResult
from mcp.types import ClientRequest
result = await self.client.send_request(
ClientRequest.model_validate({"method": method, "params": params}),
ClientResult,
)
return result.model_dump(by_alias=True, mode="json", exclude_none=True)
raise RuntimeError("persistent MCP session has no method invoke transport")
async def send_notification(
self,
method: str,
params: dict[str, Any] | None = None,
) -> None:
"""Send an MCP notification through the owner task or injected session."""
if self.send_notification_callback is not None:
await self.send_notification_callback(method, params)
return
if self.client is not None:
from mcp.types import ClientNotification
await self.client.send_notification(
ClientNotification.model_validate({"method": method, "params": params})
)
return
raise RuntimeError("persistent MCP session has no notification send transport")
async def close(self) -> None: async def close(self) -> None:
"""Close the transport/session stack owned by the runtime factory.""" """Close the transport/session stack owned by the runtime factory."""
if self.close_callback is not None: if self.close_callback is not None:
+2
View File
@@ -10,6 +10,7 @@ from .converters import (
) )
from .protocols import ( from .protocols import (
BackendAdapter, BackendAdapter,
McpSourceOperations,
PromptRuntime, PromptRuntime,
ResourceRuntime, ResourceRuntime,
StatefulMcpRuntime, StatefulMcpRuntime,
@@ -21,6 +22,7 @@ from .protocols import (
__all__ = [ __all__ = [
"BackendAdapter", "BackendAdapter",
"McpSdkAdapter", "McpSdkAdapter",
"McpSourceOperations",
"PromptRuntime", "PromptRuntime",
"ResourceRuntime", "ResourceRuntime",
"StatefulMcpRuntime", "StatefulMcpRuntime",
+10 -3
View File
@@ -17,7 +17,9 @@ class ToolCallResult:
meta: dict[str, Any] = field(default_factory=dict) meta: dict[str, Any] = field(default_factory=dict)
class BackendAdapter(Protocol): class McpSourceOperations(Protocol):
"""Full MCP operation surface shared by one-shot adapters and persistent runtimes."""
async def list_tools( async def list_tools(
self, self,
connection: McpSourceConnection, connection: McpSourceConnection,
@@ -82,6 +84,10 @@ class BackendAdapter(Protocol):
) -> ToolCallResult: ... ) -> ToolCallResult: ...
class BackendAdapter(McpSourceOperations, Protocol):
"""One-shot or adapter-style MCP operation executor."""
class ToolRuntime(Protocol): class ToolRuntime(Protocol):
"""Runtime boundary for executing MCP tools from workflow nodes.""" """Runtime boundary for executing MCP tools from workflow nodes."""
@@ -133,8 +139,8 @@ class PromptRuntime(Protocol):
) -> dict[str, Any]: ... ) -> dict[str, Any]: ...
class StatefulMcpRuntime(ToolRuntime, ResourceRuntime, PromptRuntime, Protocol): class StatefulMcpRuntime(McpSourceOperations, Protocol):
"""Stateful execution/read/list boundary for configured MCP sources. """Persistent MCP operation executor for configured sources.
Implementations keep source session state across calls. Catalog refresh may Implementations keep source session state across calls. Catalog refresh may
still use one-shot adapters by policy. still use one-shot adapters by policy.
@@ -143,6 +149,7 @@ class StatefulMcpRuntime(ToolRuntime, ResourceRuntime, PromptRuntime, Protocol):
__all__ = [ __all__ = [
"BackendAdapter", "BackendAdapter",
"McpSourceOperations",
"PromptRuntime", "PromptRuntime",
"ResourceRuntime", "ResourceRuntime",
"StatefulMcpRuntime", "StatefulMcpRuntime",
+14 -5
View File
@@ -47,11 +47,20 @@ _TRANSPORT_METADATA_KEYS = {
class LegacyConnectionConfigLike(Protocol): class LegacyConnectionConfigLike(Protocol):
"""Structural shape needed from legacy broker connection configs.""" """Structural shape needed from legacy broker connection configs."""
id: str @property
server: str def id(self) -> str: ...
account: str
enabled: bool @property
metadata: Mapping[str, object] def server(self) -> str: ...
@property
def account(self) -> str: ...
@property
def enabled(self) -> bool: ...
@property
def metadata(self) -> Mapping[str, object]: ...
class McpSourceRegistryEntry(SourceRegistryBaseModel): class McpSourceRegistryEntry(SourceRegistryBaseModel):
@@ -181,6 +181,59 @@ class _StatefulRuntime:
] ]
} }
async def list_tools(self, connection, auth):
from wf_sources_mcp.catalog import DiscoveredTool
return [
DiscoveredTool(
name="echo_tool",
title="Echo Tool",
description="Echo text back",
input_schema={"type": "object"},
output_schema={"type": "object"},
)
]
async def list_resources(self, connection, auth):
from wf_sources_mcp.catalog import DiscoveredResource
return [
DiscoveredResource(
uri="demo://docs/welcome",
name="resource.welcome",
title="Welcome Resource",
description="Welcome resource",
mime_type="text/plain",
)
]
async def list_prompts(self, connection, auth):
from wf_sources_mcp.catalog import DiscoveredPrompt
return [
DiscoveredPrompt(
name="prompt.summarize",
title="Summarize Prompt",
description="Summarize text",
arguments=[
{
"name": "text",
"required": True,
"description": "Text to summarize",
}
],
)
]
async def get_connection_metadata(self, connection, auth):
return {"server": connection.server, "transport": "stdio"}
async def invoke_method(self, connection, auth, method, params=None):
raise AssertionError("not used by content access tests")
async def send_notification(self, connection, auth, method, params=None):
raise AssertionError("not used by content access tests")
async def test_content_access_uses_stateful_runtime_for_upstream_content() -> None: async def test_content_access_uses_stateful_runtime_for_upstream_content() -> None:
runtime = _StatefulRuntime() runtime = _StatefulRuntime()
+210 -3
View File
@@ -11,6 +11,7 @@ from wf_mcp.events import McpEvent
from wf_mcp.models import AuthRecord, CatalogSnapshot, ConnectionConfig from wf_mcp.models import AuthRecord, CatalogSnapshot, ConnectionConfig
from wf_mcp.storage import FileAuthStore, FileCatalogStore, FileStore from wf_mcp.storage import FileAuthStore, FileCatalogStore, FileStore
from wf_platform import CapabilityBuckets, CapabilitySource, SourcePermissions from wf_platform import CapabilityBuckets, CapabilitySource, SourcePermissions
from wf_sources_mcp.catalog import DiscoveredTool
from ..test_support import FakeAdapter, local_temp_root from ..test_support import FakeAdapter, local_temp_root
from ..workflow_surface.conftest import echo_artifact from ..workflow_surface.conftest import echo_artifact
@@ -357,6 +358,11 @@ class _StatefulRuntime:
def __init__(self) -> None: def __init__(self) -> None:
self.resources: list[tuple[str, str]] = [] self.resources: list[tuple[str, str]] = []
self.prompts: list[tuple[str, str, dict[str, str] | None]] = [] self.prompts: list[tuple[str, str, dict[str, str] | None]] = []
self.tools_called: list[tuple[str, str]] = []
self.methods_invoked: list[tuple[str, str, str, dict[str, object] | None]] = []
self.notifications_sent: list[
tuple[str, str, str, dict[str, object] | None]
] = []
async def call_tool(self, connection, auth, tool_name, payload): async def call_tool(self, connection, auth, tool_name, payload):
raise AssertionError("not used by these tests") raise AssertionError("not used by these tests")
@@ -382,6 +388,38 @@ class _StatefulRuntime:
] ]
} }
async def list_tools(self, connection, auth):
self.tools_called.append((connection.id, connection.provider))
return [
DiscoveredTool(
name="stateful_tool",
title="Stateful Tool",
description="A stateful tool",
input_schema={"type": "object"},
output_schema={"type": "object"},
)
]
async def list_resources(self, connection, auth):
return []
async def list_prompts(self, connection, auth):
return []
async def get_connection_metadata(self, connection, auth):
return {"server": connection.provider, "transport": "stdio"}
async def invoke_method(self, connection, auth, method, params=None):
self.methods_invoked.append(
(connection.id, connection.provider, method, params)
)
return {"echoed": (params or {}).get("text", "")}
async def send_notification(self, connection, auth, method, params=None):
self.notifications_sent.append(
(connection.id, connection.provider, method, params)
)
class _ExplodingContentAdapter(FakeAdapter): class _ExplodingContentAdapter(FakeAdapter):
async def read_resource(self, connection, auth, uri): async def read_resource(self, connection, auth, uri):
@@ -451,10 +489,179 @@ async def test_upstream_transport_prefers_stateful_runtime_for_prompts(
) )
assert result["messages"][0]["content"]["text"] == "stateful prompt" assert result["messages"][0]["content"]["text"] == "stateful prompt"
assert runtime.prompts == [ assert runtime.prompts == [("demo.personal", "prompt.summarize", {"text": "hello"})]
("demo.personal", "prompt.summarize", {"text": "hello"})
]
assert [event.kind for event in events] == [ assert [event.kind for event in events] == [
"prompt_get_started", "prompt_get_started",
"prompt_get_completed", "prompt_get_completed",
] ]
async def test_upstream_transport_prefers_stateful_runtime_for_invoke_method(
tmp_path: Path,
) -> None:
events: list[McpEvent] = []
runtime = _StatefulRuntime()
transport = UpstreamTransportService(
auth_store=FileStore(tmp_path),
catalog_store=FileStore(tmp_path),
event_sink=events.append,
stateful_runtime=runtime,
)
transport.register_adapter("demo", _ExplodingContentAdapter())
connection = ConnectionConfig(
id="demo.personal",
server="demo",
account="personal",
metadata=_fake_transport_metadata(),
)
result = await transport.invoke_method(
connection,
"demo.echo",
params={"text": "hello"},
)
assert result["echoed"] == "hello"
assert runtime.methods_invoked == [
("demo.personal", "demo", "demo.echo", {"text": "hello"})
]
assert [event.kind for event in events] == [
"raw_method_started",
"raw_method_completed",
]
async def test_upstream_transport_prefers_stateful_runtime_for_send_notification(
tmp_path: Path,
) -> None:
events: list[McpEvent] = []
runtime = _StatefulRuntime()
transport = UpstreamTransportService(
auth_store=FileStore(tmp_path),
catalog_store=FileStore(tmp_path),
event_sink=events.append,
stateful_runtime=runtime,
)
transport.register_adapter("demo", _ExplodingContentAdapter())
connection = ConnectionConfig(
id="demo.personal",
server="demo",
account="personal",
metadata=_fake_transport_metadata(),
)
await transport.send_notification(
connection,
"notifications/test",
params={"data": "value"},
)
assert runtime.notifications_sent == [
("demo.personal", "demo", "notifications/test", {"data": "value"})
]
assert [event.kind for event in events] == [
"raw_notification_started",
"raw_notification_completed",
]
async def test_upstream_transport_prefers_stateful_runtime_for_catalog_refresh(
tmp_path: Path,
) -> None:
events: list[McpEvent] = []
store = FileStore(tmp_path)
connections = ConnectionRegistry()
connection = ConnectionConfig(
id="demo.personal",
server="demo",
account="personal",
metadata=_fake_transport_metadata(),
)
connections.register(connection)
runtime = _StatefulRuntime()
transport = UpstreamTransportService(
auth_store=store,
catalog_store=store,
event_sink=events.append,
stateful_runtime=runtime,
)
transport.register_adapter("demo", _ExplodingContentAdapter())
source_catalog = SourceCatalogService(
store=store,
connection_lookup=connections.get,
connection_list_enabled=connections.list_enabled,
connection_list_all=connections.list_all,
tool_executor_for=transport.tool_executor_for,
load_auth=transport.load_connection_auth,
emit_event=events.append,
)
source_catalog.hydrate_connection_source_from_snapshot(connection)
await transport.refresh_connection_catalog(
connection,
source_catalog=source_catalog,
record_catalog_change_events=lambda source_id, snapshot, reason: None,
)
assert runtime.tools_called == [("demo.personal", "demo")]
assert "catalog_refresh_started" in [event.kind for event in events]
assert "catalog_refresh_completed" in [event.kind for event in events]
class _ExplodingAdapterForDiagnostics(FakeAdapter):
async def list_tools(self, connection, auth):
raise AssertionError("adapter list_tools should not be used in diagnostics")
async def test_upstream_transport_prefers_stateful_runtime_for_deployment_diagnostics(
tmp_path: Path,
) -> None:
runtime = _StatefulRuntime()
connections = ConnectionRegistry()
connection = ConnectionConfig(
id="demo.personal",
server="demo",
account="personal",
metadata=_fake_transport_metadata(),
)
connections.register(connection)
transport = UpstreamTransportService(
auth_store=FileStore(tmp_path),
catalog_store=FileStore(tmp_path),
event_sink=lambda event: None,
stateful_runtime=runtime,
)
transport.register_adapter("demo", _ExplodingAdapterForDiagnostics())
source_catalog = SourceCatalogService(
store=transport.catalog_store,
connection_lookup=connections.get,
connection_list_enabled=connections.list_enabled,
connection_list_all=connections.list_all,
tool_executor_for=transport.tool_executor_for,
load_auth=transport.load_connection_auth,
emit_event=lambda event: None,
)
source_catalog.register_capability_source(
CapabilitySource(
id="demo.personal",
kind="connection",
permissions=SourcePermissions(calls_upstream=True),
capabilities=CapabilityBuckets(),
)
)
artifact = echo_artifact()
deployment = WorkflowDeployment(
id="echo.personal",
artifact_id="echo",
artifact_version=1,
bindings=[{"logical_source": "demo", "concrete_source": "demo.personal"}],
)
diagnostics = await transport.deployment_diagnostics(
deployment=deployment,
artifacts=[artifact],
source_catalog=source_catalog,
)
assert diagnostics == []
assert runtime.tools_called == [("demo.personal", "demo")]
+329 -2
View File
@@ -4,14 +4,19 @@ from contextlib import AsyncExitStack
from typing import Any from typing import Any
import pytest import pytest
from mcp import ClientResult
from mcp.client.session import ClientSession from mcp.client.session import ClientSession
from mcp.types import CallToolResult as RawCallToolResult from mcp.types import CallToolResult as RawCallToolResult
from mcp.types import ( from mcp.types import (
ClientNotification,
ClientRequest,
ListPromptsResult, ListPromptsResult,
ListResourcesResult, ListResourcesResult,
ListToolsResult,
Prompt, Prompt,
Resource, Resource,
TextContent, TextContent,
Tool,
) )
from pydantic import AnyUrl from pydantic import AnyUrl
@@ -167,6 +172,44 @@ class _FakeFactory(PersistentSessionFactory):
] ]
) )
async def list_tools(self) -> ListToolsResult:
return ListToolsResult(
tools=[
Tool(
name="tool.runtime",
title="Runtime Tool",
description="Runtime-scoped tool.",
inputSchema={"type": "object"},
)
]
)
async def get_connection_metadata(self) -> dict[str, object]:
return {"server": "demo", "transport": "stdio"}
async def invoke_method(
self,
method: str,
params: dict[str, object] | None = None,
) -> dict[str, object]:
return {"echoed": (params or {}).get("text", "")}
async def send_notification(
self,
method: str,
params: dict[str, object] | None = None,
) -> None:
return None
async def send_request(
self,
request: ClientRequest,
result_type: type[ClientResult],
) -> ClientResult:
return ClientResult.model_validate(
{"jsonrpc": "2.0", "id": 1, "result": {}}
)
return _FakeClient() # type: ignore[return-value] return _FakeClient() # type: ignore[return-value]
@@ -242,8 +285,10 @@ def test_persistent_session_public_runtime_exposes_safe_read_operations() -> Non
assert "get_prompt" in public_operations assert "get_prompt" in public_operations
assert "list_resources" in public_operations assert "list_resources" in public_operations
assert "list_prompts" in public_operations assert "list_prompts" in public_operations
assert "invoke_method" not in public_operations assert "list_tools" in public_operations
assert "send_notification" not in public_operations assert "get_connection_metadata" in public_operations
assert "invoke_method" in public_operations
assert "send_notification" in public_operations
@pytest.mark.asyncio @pytest.mark.asyncio
@@ -347,6 +392,237 @@ async def test_persistent_session_factory_routes_resource_and_prompt_lists() ->
assert prompts[0].name == "prompt.runtime" assert prompts[0].name == "prompt.runtime"
@pytest.mark.asyncio
async def test_persistent_session_factory_routes_list_tools_through_owner() -> None:
factory = _FakeFactory()
session = await factory.create(_connection(), None)
tools = await session.list_tools()
await session.close()
assert tools[0].name == "tool.runtime"
assert tools[0].description == "Runtime-scoped tool."
@pytest.mark.asyncio
async def test_persistent_session_factory_routes_metadata_through_owner() -> None:
factory = _FakeFactory()
session = await factory.create(_connection(), None)
metadata = await session.get_connection_metadata()
await session.close()
assert metadata["server"] == "demo"
assert metadata["transport"] == "stdio"
@pytest.mark.asyncio
async def test_persistent_session_factory_routes_invoke_method_through_owner() -> None:
factory = _FakeFactory()
session = await factory.create(_connection(), None)
result = await session.invoke_method("ping")
await session.close()
assert isinstance(result, dict)
@pytest.mark.asyncio
async def test_persistent_session_factory_routes_send_notification_through_owner() -> (
None
):
factory = _FakeFactory()
session = await factory.create(_connection(), None)
await session.send_notification("notifications/initialized")
await session.close()
@pytest.mark.asyncio
async def test_persistent_session_list_tools_callback() -> None:
async def list_tools_cb() -> list:
from wf_sources_mcp.catalog import DiscoveredTool
return [
DiscoveredTool(
name="cb_tool",
title=None,
description="Callback tool",
input_schema={"type": "object"},
output_schema={"type": "object"},
)
]
session = PersistentMcpSession(
connection=_connection(),
auth=None,
list_tools_callback=list_tools_cb,
)
tools = await session.list_tools()
assert tools[0].name == "cb_tool"
@pytest.mark.asyncio
async def test_persistent_session_get_connection_metadata_callback() -> None:
async def metadata_cb() -> dict[str, object]:
return {"server": "cb_server", "transport": "http"}
session = PersistentMcpSession(
connection=_connection(),
auth=None,
get_connection_metadata_callback=metadata_cb,
)
metadata = await session.get_connection_metadata()
assert metadata["server"] == "cb_server"
@pytest.mark.asyncio
async def test_persistent_session_invoke_method_callback() -> None:
async def invoke_cb(
method: str, params: dict[str, object] | None
) -> dict[str, object]:
return {"method": method, "params": params}
session = PersistentMcpSession(
connection=_connection(),
auth=None,
invoke_method_callback=invoke_cb,
)
result = await session.invoke_method("test.method", {"key": "val"})
assert result["method"] == "test.method"
assert result["params"] == {"key": "val"}
@pytest.mark.asyncio
async def test_persistent_session_send_notification_callback() -> None:
sent: list[tuple[str, dict[str, object] | None]] = []
async def notify_cb(method: str, params: dict[str, object] | None) -> None:
sent.append((method, params))
session = PersistentMcpSession(
connection=_connection(),
auth=None,
send_notification_callback=notify_cb,
)
await session.send_notification("test.event", {"data": 1})
assert sent == [("test.event", {"data": 1})]
@pytest.mark.asyncio
async def test_persistent_session_list_tools_client_fallback() -> None:
from mcp.types import ListToolsResult, Tool
class _MinimalClient:
async def list_tools(self) -> ListToolsResult:
return ListToolsResult(
tools=[
Tool(
name="client_tool",
description="Client tool",
inputSchema={"type": "object"},
)
]
)
session = PersistentMcpSession(
connection=_connection(),
auth=None,
client=_MinimalClient(), # type: ignore[arg-type]
)
tools = await session.list_tools()
assert tools[0].name == "client_tool"
@pytest.mark.asyncio
async def test_persistent_session_invoke_method_client_fallback() -> None:
from mcp import ClientResult
class _MinimalClient:
async def send_request(
self,
request: ClientRequest,
result_type: type[ClientResult],
) -> ClientResult:
return ClientResult.model_validate(
{"jsonrpc": "2.0", "id": 1, "result": {"tools": []}}
)
session = PersistentMcpSession(
connection=_connection(),
auth=None,
client=_MinimalClient(), # type: ignore[arg-type]
)
result = await session.invoke_method("tools/list")
assert result["result"]["tools"] == []
@pytest.mark.asyncio
async def test_persistent_session_send_notification_client_fallback() -> None:
sent: list[ClientNotification] = []
class _MinimalClient:
async def send_notification(self, notification: ClientNotification) -> None:
sent.append(notification)
session = PersistentMcpSession(
connection=_connection(),
auth=None,
client=_MinimalClient(), # type: ignore[arg-type]
)
await session.send_notification("notifications/initialized")
assert len(sent) == 1
@pytest.mark.asyncio
async def test_persistent_session_raises_without_tools_transport() -> None:
session = PersistentMcpSession(connection=_connection(), auth=None)
with pytest.raises(RuntimeError, match="no tools list transport"):
await session.list_tools()
@pytest.mark.asyncio
async def test_persistent_session_raises_without_invoke_transport() -> None:
session = PersistentMcpSession(connection=_connection(), auth=None)
with pytest.raises(RuntimeError, match="no method invoke transport"):
await session.invoke_method("test.ping")
@pytest.mark.asyncio
async def test_persistent_session_raises_without_notification_transport() -> None:
session = PersistentMcpSession(connection=_connection(), auth=None)
with pytest.raises(RuntimeError, match="no notification send transport"):
await session.send_notification("test.event")
@pytest.mark.asyncio
async def test_persistent_session_get_connection_metadata_local_fallback() -> None:
session = PersistentMcpSession(connection=_connection(), auth=None)
metadata = await session.get_connection_metadata()
assert metadata["server"] == "demo"
assert metadata["transport"] == "stdio"
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_runtime_pool_reuses_session_for_resource_and_prompt_lists() -> None: async def test_runtime_pool_reuses_session_for_resource_and_prompt_lists() -> None:
factory = _FakeFactory() factory = _FakeFactory()
@@ -362,6 +638,57 @@ async def test_runtime_pool_reuses_session_for_resource_and_prompt_lists() -> No
assert factory.created_connections == [connection] assert factory.created_connections == [connection]
@pytest.mark.asyncio
async def test_runtime_pool_reuses_session_for_list_tools() -> None:
factory = _FakeFactory()
pool = McpRuntimePool(factory.create)
connection = _connection()
tools = await pool.list_tools(connection, None)
await pool.close_all()
assert tools[0].name == "tool.runtime"
assert factory.created_connections == [connection]
@pytest.mark.asyncio
async def test_runtime_pool_reuses_session_for_invoke_method() -> None:
factory = _FakeFactory()
pool = McpRuntimePool(factory.create)
connection = _connection()
result = await pool.invoke_method(connection, None, "ping")
await pool.close_all()
assert isinstance(result, dict)
assert factory.created_connections == [connection]
@pytest.mark.asyncio
async def test_runtime_pool_reuses_session_for_send_notification() -> None:
factory = _FakeFactory()
pool = McpRuntimePool(factory.create)
connection = _connection()
await pool.send_notification(connection, None, "notifications/initialized")
await pool.close_all()
assert factory.created_connections == [connection]
@pytest.mark.asyncio
async def test_runtime_pool_reuses_session_for_metadata() -> None:
factory = _FakeFactory()
pool = McpRuntimePool(factory.create)
connection = _connection()
metadata = await pool.get_connection_metadata(connection, None)
await pool.close_all()
assert metadata["server"] == "demo"
assert factory.created_connections == [connection]
def test_runtime_pool_satisfies_stateful_protocol_static_shape() -> None: def test_runtime_pool_satisfies_stateful_protocol_static_shape() -> None:
from wf_sources_mcp.sdk import ( from wf_sources_mcp.sdk import (
PromptRuntime, PromptRuntime,
+153 -1
View File
@@ -4,7 +4,7 @@ from dataclasses import is_dataclass
from typing import cast from typing import cast
from wf_sources_mcp.auth import AuthRecord from wf_sources_mcp.auth import AuthRecord
from wf_sources_mcp.catalog import DiscoveredTool from wf_sources_mcp.catalog import DiscoveredPrompt, DiscoveredResource, DiscoveredTool
from wf_sources_mcp.connections import McpSourceConnection from wf_sources_mcp.connections import McpSourceConnection
from wf_sources_mcp.sdk import BackendAdapter, ToolCallResult, ToolExecutor from wf_sources_mcp.sdk import BackendAdapter, ToolCallResult, ToolExecutor
from wf_sources_mcp.transports import StdioSourceTransport from wf_sources_mcp.transports import StdioSourceTransport
@@ -36,6 +36,90 @@ class EchoAdapter:
return ToolCallResult(outcome="ok", output={"echoed": payload}) return ToolCallResult(outcome="ok", output={"echoed": payload})
class _FullSurfaceAdapter:
"""Implements every MCP operation for protocol conformance tests."""
async def list_tools(
self,
connection: McpSourceConnection,
auth: AuthRecord | None,
) -> list[DiscoveredTool]:
return [
DiscoveredTool(
name="echo",
title=None,
description="Echo",
input_schema={"type": "object"},
output_schema={"type": "object"},
)
]
async def list_resources(
self,
connection: McpSourceConnection,
auth: AuthRecord | None,
) -> list[DiscoveredResource]:
return []
async def list_prompts(
self,
connection: McpSourceConnection,
auth: AuthRecord | None,
) -> list[DiscoveredPrompt]:
return []
async def get_connection_metadata(
self,
connection: McpSourceConnection,
auth: AuthRecord | None,
) -> dict[str, object]:
return {"server": "demo"}
async def read_resource(
self,
connection: McpSourceConnection,
auth: AuthRecord | None,
uri: str,
) -> dict[str, object]:
return {"contents": []}
async def get_prompt(
self,
connection: McpSourceConnection,
auth: AuthRecord | None,
prompt_name: str,
arguments: dict[str, str] | None = None,
) -> dict[str, object]:
return {"messages": []}
async def invoke_method(
self,
connection: McpSourceConnection,
auth: AuthRecord | None,
method: str,
params: dict[str, object] | None = None,
) -> dict[str, object]:
return {}
async def send_notification(
self,
connection: McpSourceConnection,
auth: AuthRecord | None,
method: str,
params: dict[str, object] | None = None,
) -> None:
return None
async def call_tool(
self,
connection: McpSourceConnection,
auth: AuthRecord | None,
tool_name: str,
payload: dict[str, object],
) -> ToolCallResult:
return ToolCallResult(outcome="ok", output={"echoed": payload})
def test_tool_call_result_is_slots_dataclass_with_empty_defaults() -> None: def test_tool_call_result_is_slots_dataclass_with_empty_defaults() -> None:
result = ToolCallResult(outcome="ok") result = ToolCallResult(outcome="ok")
@@ -77,6 +161,72 @@ async def test_tool_executor_protocol_can_describe_tool_calls() -> None:
assert result.output == {"echoed": {"message": "hello"}} assert result.output == {"echoed": {"message": "hello"}}
async def test_backend_adapter_protocol_full_operation_surface() -> None:
adapter = cast(BackendAdapter, _FullSurfaceAdapter())
conn = McpSourceConnection(
id="demo.default",
provider="demo",
account="default",
transport=StdioSourceTransport(command="echo"),
)
tools = await adapter.list_tools(conn, None)
resources = await adapter.list_resources(conn, None)
prompts = await adapter.list_prompts(conn, None)
metadata = await adapter.get_connection_metadata(conn, None)
read_result = await adapter.read_resource(conn, None, "test://x")
prompt_result = await adapter.get_prompt(conn, None, "prompt.summarize")
invoke_result = await adapter.invoke_method(conn, None, "ping")
await adapter.send_notification(conn, None, "test.notify")
call_result = await adapter.call_tool(conn, None, "echo", {"text": "hi"})
assert tools[0].name == "echo"
assert resources == []
assert prompts == []
assert metadata["server"] == "demo"
assert read_result == {"contents": []}
assert prompt_result == {"messages": []}
assert invoke_result == {}
assert call_result.outcome == "ok"
async def test_stateful_mcp_runtime_protocol_full_operation_surface() -> None:
from wf_sources_mcp.sdk import StatefulMcpRuntime
runtime = cast(StatefulMcpRuntime, _FullSurfaceAdapter())
conn = McpSourceConnection(
id="demo.default",
provider="demo",
account="default",
transport=StdioSourceTransport(command="echo"),
)
tools = await runtime.list_tools(conn, None)
resources = await runtime.list_resources(conn, None)
prompts = await runtime.list_prompts(conn, None)
metadata = await runtime.get_connection_metadata(conn, None)
read_result = await runtime.read_resource(conn, None, "test://x")
prompt_result = await runtime.get_prompt(conn, None, "prompt.summarize")
invoke_result = await runtime.invoke_method(conn, None, "ping")
await runtime.send_notification(conn, None, "test.notify")
call_result = await runtime.call_tool(conn, None, "echo", {"text": "hi"})
assert tools[0].name == "echo"
assert resources == []
assert prompts == []
assert metadata["server"] == "demo"
assert read_result == {"contents": []}
assert prompt_result == {"messages": []}
assert invoke_result == {}
assert call_result.outcome == "ok"
def test_mcp_source_operations_protocol_shape() -> None:
from wf_sources_mcp.sdk import McpSourceOperations
assert McpSourceOperations.__name__ == "McpSourceOperations"
def test_stateful_mcp_runtime_protocol_shape() -> None: def test_stateful_mcp_runtime_protocol_shape() -> None:
from wf_sources_mcp.sdk import StatefulMcpRuntime, ToolExecutor from wf_sources_mcp.sdk import StatefulMcpRuntime, ToolExecutor
@@ -86,6 +236,7 @@ def test_stateful_mcp_runtime_protocol_shape() -> None:
def test_stateful_runtime_protocol_slices_export() -> None: def test_stateful_runtime_protocol_slices_export() -> None:
from wf_sources_mcp.sdk import ( from wf_sources_mcp.sdk import (
McpSourceOperations,
PromptRuntime, PromptRuntime,
ResourceRuntime, ResourceRuntime,
StatefulMcpRuntime, StatefulMcpRuntime,
@@ -97,4 +248,5 @@ def test_stateful_runtime_protocol_slices_export() -> None:
assert ResourceRuntime.__name__ == "ResourceRuntime" assert ResourceRuntime.__name__ == "ResourceRuntime"
assert PromptRuntime.__name__ == "PromptRuntime" assert PromptRuntime.__name__ == "PromptRuntime"
assert ToolExecutor.__name__ == "ToolExecutor" assert ToolExecutor.__name__ == "ToolExecutor"
assert McpSourceOperations.__name__ == "McpSourceOperations"
assert StatefulMcpRuntime.__name__ == "StatefulMcpRuntime" assert StatefulMcpRuntime.__name__ == "StatefulMcpRuntime"