refactor: route mcp runtime operations through client facade

This commit is contained in:
lda
2026-06-07 18:30:54 +07:00 Verified
parent 40d06734f1
commit efeb8c7cfb
6 changed files with 94 additions and 48 deletions
+4
View File
@@ -242,6 +242,10 @@ implementation state.
`wf_sources_mcp.client`. The one-shot SDK adapter delegates MCP operation `wf_sources_mcp.client`. The one-shot SDK adapter delegates MCP operation
calls and conversion through the facade; persistent runtime remains calls and conversion through the facade; persistent runtime remains
tool-call-only until a separate owner-task routing slice. tool-call-only until a separate owner-task routing slice.
- Completed: persistent MCP runtime owner now uses a generic explicit
operation queue with request metadata and `McpSourceClient` execution.
Public runtime remains tool-call-only; `operation` strings are diagnostics
labels, not dispatch.
- Auth/source secrets boundary: keep registry desired state separate from - Auth/source secrets boundary: keep registry desired state separate from
upstream credentials, and surface missing auth as validation diagnostics. upstream credentials, and surface missing auth as validation diagnostics.
The contract is now specified in The contract is now specified in
@@ -111,7 +111,10 @@ First slices should move leaf modules only and leave `wf_mcp` re-export shims:
`wf_sources_mcp.client`. `McpSdkAdapter` now delegates operation handling to `wf_sources_mcp.client`. `McpSdkAdapter` now delegates operation handling to
this facade. Persistent runtime still exposes only `call_tool`; expanding it this facade. Persistent runtime still exposes only `call_tool`; expanding it
requires a separate owner-task request routing slice. requires a separate owner-task request routing slice.
10. Upstream transport/discovery/session services. 10. Complete: persistent MCP runtime owner now routes explicit callables through
a generic operation queue with request metadata. The runtime still exposes
only `call_tool`; non-tool methods require a separate public-surface slice.
11. 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`,
+51 -20
View File
@@ -1,18 +1,25 @@
from __future__ import annotations from __future__ import annotations
import asyncio import asyncio
import time
from collections.abc import Awaitable, Callable
from contextlib import AsyncExitStack from contextlib import AsyncExitStack
from dataclasses import dataclass, field from dataclasses import dataclass, field
from itertools import count
from typing import Any, Generic, TypeVar
from mcp.client.session import ClientSession from mcp.client.session import ClientSession
from mcp.types import CallToolResult
from wf_sources_mcp.auth import AuthRecord from wf_sources_mcp.auth import AuthRecord
from wf_sources_mcp.client import 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 .session import PersistentMcpSession from .session import PersistentMcpSession
T = TypeVar("T")
ClientOperation = Callable[[McpSourceClient], Awaitable[T]]
@dataclass(slots=True) @dataclass(slots=True)
class PersistentSessionFactory: class PersistentSessionFactory:
@@ -49,12 +56,19 @@ class PersistentSessionFactory:
@dataclass(slots=True) @dataclass(slots=True)
class _ToolCallRequest: class _ClientOperationRequest(Generic[T]):
"""One request submitted to the task that owns the MCP transport.""" """One explicit operation submitted to the MCP transport owner task.
tool_name: str `operation` is metadata for diagnostics/tracing only. Execution uses `run`;
payload: dict[str, object] do not dispatch with `getattr(client, operation)`.
result: asyncio.Future[CallToolResult] """
operation: str
connection_id: str
sequence: int
submitted_at: float
run: ClientOperation[T]
result: asyncio.Future[T]
@dataclass(slots=True) @dataclass(slots=True)
@@ -71,9 +85,10 @@ class _SessionOwner:
factory: PersistentSessionFactory factory: PersistentSessionFactory
connection: McpSourceConnection connection: McpSourceConnection
auth: AuthRecord | None auth: AuthRecord | None
_requests: asyncio.Queue[_ToolCallRequest | None] = field( _requests: asyncio.Queue[_ClientOperationRequest[Any] | None] = field(
default_factory=asyncio.Queue default_factory=asyncio.Queue
) )
_sequence: count = field(default_factory=lambda: count(1))
_task: asyncio.Task[None] | None = None _task: asyncio.Task[None] | None = None
async def start(self) -> None: async def start(self) -> None:
@@ -85,21 +100,29 @@ class _SessionOwner:
) )
await ready await ready
async def call_tool( async def submit(
self, self,
tool_name: str, operation: str,
payload: dict[str, object], run: ClientOperation[T],
) -> CallToolResult: ) -> T:
"""Submit a call and fail promptly if its transport owner exits.""" """Submit an explicit client operation to the MCP owner task."""
task = self._task task = self._task
if task is None: if task is None:
raise RuntimeError("persistent MCP session is not started") raise RuntimeError("persistent MCP session is not started")
if task.done(): if task.done():
await task await task
raise RuntimeError("persistent MCP session stopped unexpectedly") raise RuntimeError("persistent MCP session stopped unexpectedly")
result = asyncio.get_running_loop().create_future()
result: asyncio.Future[T] = asyncio.get_running_loop().create_future()
await self._requests.put( await self._requests.put(
_ToolCallRequest(tool_name=tool_name, payload=payload, result=result) _ClientOperationRequest(
operation=operation,
connection_id=self.connection.id,
sequence=next(self._sequence),
submitted_at=time.monotonic(),
run=run,
result=result,
)
) )
done, _pending = await asyncio.wait( done, _pending = await asyncio.wait(
{result, task}, return_when=asyncio.FIRST_COMPLETED {result, task}, return_when=asyncio.FIRST_COMPLETED
@@ -109,6 +132,17 @@ class _SessionOwner:
await task await task
raise RuntimeError("persistent MCP session stopped unexpectedly") raise RuntimeError("persistent MCP session stopped unexpectedly")
async def call_tool(
self,
tool_name: str,
payload: dict[str, object],
) -> ToolCallResult:
"""Submit a tool call through the generic owner-task operation queue."""
return await self.submit(
operation="call_tool",
run=lambda client: client.call_tool(tool_name, payload),
)
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
@@ -126,15 +160,14 @@ class _SessionOwner:
session = await self.factory._create_with_stack( session = await self.factory._create_with_stack(
stack, self.connection, self.auth stack, self.connection, self.auth
) )
client = McpSourceClient(session=session, connection=self.connection)
ready.set_result(None) ready.set_result(None)
while True: while True:
request = await self._requests.get() request = await self._requests.get()
if request is None: if request is None:
return return
try: try:
response = await session.call_tool( response = await request.run(client)
request.tool_name, request.payload
)
except Exception as exc: except Exception as exc:
request.result.set_exception(exc) request.result.set_exception(exc)
else: else:
@@ -143,8 +176,6 @@ class _SessionOwner:
if not ready.done(): if not ready.done():
ready.set_exception(exc) ready.set_exception(exc)
return return
# Calls already queued behind the failing request cannot otherwise
# observe that their sole transport owner has exited.
while not self._requests.empty(): while not self._requests.empty():
pending = self._requests.get_nowait() pending = self._requests.get_nowait()
if pending is not None and not pending.result.done(): if pending is not None and not pending.result.done():
+4 -6
View File
@@ -5,14 +5,13 @@ from dataclasses import dataclass
from typing import Any from typing import Any
from mcp.client.session import ClientSession from mcp.client.session import ClientSession
from mcp.types import CallToolResult
from wf_sources_mcp.auth import AuthRecord from wf_sources_mcp.auth import AuthRecord
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
RawToolCaller = Callable[[str, dict[str, Any]], Awaitable[CallToolResult]] RawToolCaller = Callable[[str, dict[str, Any]], Awaitable[ToolCallResult]]
@dataclass(slots=True) @dataclass(slots=True)
@@ -39,12 +38,11 @@ class PersistentMcpSession:
self, tool_name: str, payload: dict[str, Any] self, tool_name: str, payload: dict[str, Any]
) -> ToolCallResult: ) -> ToolCallResult:
if self.call_callback is not None: if self.call_callback is not None:
result = await self.call_callback(tool_name, payload) return await self.call_callback(tool_name, payload)
elif self.client is not None: if self.client is not None:
result = await self.client.call_tool(tool_name, payload) result = await self.client.call_tool(tool_name, payload)
else:
raise RuntimeError("persistent MCP session has no tool call transport")
return tool_result_to_call_result(result) return tool_result_to_call_result(result)
raise RuntimeError("persistent MCP session has no tool call 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."""
+30 -20
View File
@@ -5,7 +5,8 @@ from typing import Any
import pytest import pytest
from mcp.client.session import ClientSession from mcp.client.session import ClientSession
from mcp.types import CallToolResult, TextContent from mcp.types import CallToolResult as RawCallToolResult
from mcp.types import TextContent
from wf_sources_mcp.auth import AuthRecord from wf_sources_mcp.auth import AuthRecord
from wf_sources_mcp.connections import McpSourceConnection from wf_sources_mcp.connections import McpSourceConnection
@@ -15,6 +16,7 @@ from wf_sources_mcp.runtime import (
connection_runtime_fingerprint, connection_runtime_fingerprint,
) )
from wf_sources_mcp.runtime.factory import PersistentSessionFactory from wf_sources_mcp.runtime.factory import PersistentSessionFactory
from wf_sources_mcp.sdk import ToolCallResult
from wf_sources_mcp.transports import StdioSourceTransport from wf_sources_mcp.transports import StdioSourceTransport
@@ -28,14 +30,11 @@ def _connection() -> McpSourceConnection:
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_persistent_session_call_callback_normalizes_tool_result() -> None: async def test_persistent_session_call_callback_returns_canonical_result() -> None:
async def call_tool(tool_name: str, payload: dict[str, Any]) -> CallToolResult: async def call_tool(tool_name: str, payload: dict[str, Any]) -> ToolCallResult:
assert tool_name == "echo" assert tool_name == "echo"
assert payload == {"text": "hi"} assert payload == {"text": "hi"}
return CallToolResult( return ToolCallResult(outcome="ok", output={"echoed": "hi"})
content=[TextContent(type="text", text="ok")],
structuredContent={"echoed": "hi"},
)
session = PersistentMcpSession( session = PersistentMcpSession(
connection=_connection(), connection=_connection(),
@@ -60,32 +59,30 @@ async def test_persistent_session_raises_without_transport() -> None:
class _FakeFactory(PersistentSessionFactory): class _FakeFactory(PersistentSessionFactory):
def __init__(self) -> None: def __init__(self) -> None:
self.calls: list[tuple[str, dict[str, object]]] = [] self.calls: list[tuple[str, dict[str, object]]] = []
self.closed = False self.created_connections: list[McpSourceConnection] = []
async def _call_tool( async def _call_tool(
self, tool_name: str, payload: dict[str, object] self, tool_name: str, payload: dict[str, object]
) -> CallToolResult: ) -> RawCallToolResult:
self.calls.append((tool_name, payload)) self.calls.append((tool_name, payload))
return CallToolResult( return RawCallToolResult(
content=[TextContent(type="text", text="ok")], content=[TextContent(type="text", text="ok")],
structuredContent={"echoed": payload["text"]}, structuredContent={"echoed": payload["text"]},
) )
async def _close(self) -> None:
self.closed = True
async def _create_with_stack( async def _create_with_stack(
self, self,
stack: AsyncExitStack, stack: AsyncExitStack,
connection: McpSourceConnection, connection: McpSourceConnection,
auth: AuthRecord | None, auth: AuthRecord | None,
) -> ClientSession: ) -> ClientSession:
self.created_connections.append(connection)
factory = self factory = self
class _FakeClient: class _FakeClient:
async def call_tool( async def call_tool(
self, tool_name: str, payload: dict[str, object] self, tool_name: str, payload: dict[str, object]
) -> CallToolResult: ) -> RawCallToolResult:
return await factory._call_tool(tool_name, payload) return await factory._call_tool(tool_name, payload)
return _FakeClient() # type: ignore[return-value] return _FakeClient() # type: ignore[return-value]
@@ -94,7 +91,8 @@ class _FakeFactory(PersistentSessionFactory):
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_persistent_session_factory_serializes_tool_calls() -> None: async def test_persistent_session_factory_serializes_tool_calls() -> None:
factory = _FakeFactory() factory = _FakeFactory()
session = await factory.create(_connection(), None) connection = _connection()
session = await factory.create(connection, None)
first = await session.call_tool("echo", {"text": "one"}) first = await session.call_tool("echo", {"text": "one"})
second = await session.call_tool("echo", {"text": "two"}) second = await session.call_tool("echo", {"text": "two"})
@@ -102,6 +100,7 @@ async def test_persistent_session_factory_serializes_tool_calls() -> None:
assert first.output == {"echoed": "one"} assert first.output == {"echoed": "one"}
assert second.output == {"echoed": "two"} assert second.output == {"echoed": "two"}
assert factory.created_connections == [connection]
assert factory.calls == [ assert factory.calls == [
("echo", {"text": "one"}), ("echo", {"text": "one"}),
("echo", {"text": "two"}), ("echo", {"text": "two"}),
@@ -117,11 +116,8 @@ async def test_runtime_pool_reuses_unchanged_connection() -> None:
) -> PersistentMcpSession: ) -> PersistentMcpSession:
created.append(connection) created.append(connection)
async def _call(tool_name: str, payload: dict[str, Any]) -> CallToolResult: async def _call(tool_name: str, payload: dict[str, Any]) -> ToolCallResult:
return CallToolResult( return ToolCallResult(outcome="ok", output={"echoed": payload["text"]})
content=[TextContent(type="text", text="ok")],
structuredContent={"echoed": payload["text"]},
)
return PersistentMcpSession( return PersistentMcpSession(
connection=connection, connection=connection,
@@ -150,3 +146,17 @@ def test_runtime_fingerprint_changes_when_transport_changes() -> None:
assert connection_runtime_fingerprint(original) != connection_runtime_fingerprint( assert connection_runtime_fingerprint(original) != connection_runtime_fingerprint(
changed changed
) )
def test_persistent_session_public_runtime_is_tool_call_only() -> None:
public_operations = {
name
for name in dir(PersistentMcpSession)
if not name.startswith("_") and callable(getattr(PersistentMcpSession, name))
}
assert "call_tool" in public_operations
assert "read_resource" not in public_operations
assert "get_prompt" not in public_operations
assert "invoke_method" not in public_operations
assert "send_notification" not in public_operations