refactor: route mcp runtime operations through client facade
This commit is contained in:
@@ -242,6 +242,10 @@ implementation state.
|
||||
`wf_sources_mcp.client`. The one-shot SDK adapter delegates MCP operation
|
||||
calls and conversion through the facade; persistent runtime remains
|
||||
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
|
||||
upstream credentials, and surface missing auth as validation diagnostics.
|
||||
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
|
||||
this facade. Persistent runtime still exposes only `call_tool`; expanding it
|
||||
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
|
||||
does not depend on `wf_mcp.workflow_surface`, `wf_mcp.admin_surface`,
|
||||
|
||||
@@ -1,18 +1,25 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable
|
||||
from contextlib import AsyncExitStack
|
||||
from dataclasses import dataclass, field
|
||||
from itertools import count
|
||||
from typing import Any, Generic, TypeVar
|
||||
|
||||
from mcp.client.session import ClientSession
|
||||
from mcp.types import CallToolResult
|
||||
|
||||
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.sdk import ToolCallResult
|
||||
|
||||
from .session import PersistentMcpSession
|
||||
|
||||
T = TypeVar("T")
|
||||
ClientOperation = Callable[[McpSourceClient], Awaitable[T]]
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class PersistentSessionFactory:
|
||||
@@ -49,12 +56,19 @@ class PersistentSessionFactory:
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class _ToolCallRequest:
|
||||
"""One request submitted to the task that owns the MCP transport."""
|
||||
class _ClientOperationRequest(Generic[T]):
|
||||
"""One explicit operation submitted to the MCP transport owner task.
|
||||
|
||||
tool_name: str
|
||||
payload: dict[str, object]
|
||||
result: asyncio.Future[CallToolResult]
|
||||
`operation` is metadata for diagnostics/tracing only. Execution uses `run`;
|
||||
do not dispatch with `getattr(client, operation)`.
|
||||
"""
|
||||
|
||||
operation: str
|
||||
connection_id: str
|
||||
sequence: int
|
||||
submitted_at: float
|
||||
run: ClientOperation[T]
|
||||
result: asyncio.Future[T]
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
@@ -71,9 +85,10 @@ class _SessionOwner:
|
||||
factory: PersistentSessionFactory
|
||||
connection: McpSourceConnection
|
||||
auth: AuthRecord | None
|
||||
_requests: asyncio.Queue[_ToolCallRequest | None] = field(
|
||||
_requests: asyncio.Queue[_ClientOperationRequest[Any] | None] = field(
|
||||
default_factory=asyncio.Queue
|
||||
)
|
||||
_sequence: count = field(default_factory=lambda: count(1))
|
||||
_task: asyncio.Task[None] | None = None
|
||||
|
||||
async def start(self) -> None:
|
||||
@@ -85,21 +100,29 @@ class _SessionOwner:
|
||||
)
|
||||
await ready
|
||||
|
||||
async def call_tool(
|
||||
async def submit(
|
||||
self,
|
||||
tool_name: str,
|
||||
payload: dict[str, object],
|
||||
) -> CallToolResult:
|
||||
"""Submit a call and fail promptly if its transport owner exits."""
|
||||
operation: str,
|
||||
run: ClientOperation[T],
|
||||
) -> T:
|
||||
"""Submit an explicit client operation to the MCP owner task."""
|
||||
task = self._task
|
||||
if task is None:
|
||||
raise RuntimeError("persistent MCP session is not started")
|
||||
if task.done():
|
||||
await task
|
||||
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(
|
||||
_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(
|
||||
{result, task}, return_when=asyncio.FIRST_COMPLETED
|
||||
@@ -109,6 +132,17 @@ class _SessionOwner:
|
||||
await task
|
||||
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:
|
||||
"""Ask the owner task to close the MCP transport in its own scope."""
|
||||
task = self._task
|
||||
@@ -126,15 +160,14 @@ class _SessionOwner:
|
||||
session = await self.factory._create_with_stack(
|
||||
stack, self.connection, self.auth
|
||||
)
|
||||
client = McpSourceClient(session=session, connection=self.connection)
|
||||
ready.set_result(None)
|
||||
while True:
|
||||
request = await self._requests.get()
|
||||
if request is None:
|
||||
return
|
||||
try:
|
||||
response = await session.call_tool(
|
||||
request.tool_name, request.payload
|
||||
)
|
||||
response = await request.run(client)
|
||||
except Exception as exc:
|
||||
request.result.set_exception(exc)
|
||||
else:
|
||||
@@ -143,8 +176,6 @@ class _SessionOwner:
|
||||
if not ready.done():
|
||||
ready.set_exception(exc)
|
||||
return
|
||||
# Calls already queued behind the failing request cannot otherwise
|
||||
# observe that their sole transport owner has exited.
|
||||
while not self._requests.empty():
|
||||
pending = self._requests.get_nowait()
|
||||
if pending is not None and not pending.result.done():
|
||||
|
||||
@@ -5,14 +5,13 @@ from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from mcp.client.session import ClientSession
|
||||
from mcp.types import CallToolResult
|
||||
|
||||
from wf_sources_mcp.auth import AuthRecord
|
||||
from wf_sources_mcp.connections import McpSourceConnection
|
||||
from wf_sources_mcp.sdk import ToolCallResult
|
||||
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)
|
||||
@@ -39,12 +38,11 @@ class PersistentMcpSession:
|
||||
self, tool_name: str, payload: dict[str, Any]
|
||||
) -> ToolCallResult:
|
||||
if self.call_callback is not None:
|
||||
result = await self.call_callback(tool_name, payload)
|
||||
elif self.client is not None:
|
||||
return await self.call_callback(tool_name, payload)
|
||||
if self.client is not None:
|
||||
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:
|
||||
"""Close the transport/session stack owned by the runtime factory."""
|
||||
|
||||
@@ -5,7 +5,8 @@ from typing import Any
|
||||
|
||||
import pytest
|
||||
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.connections import McpSourceConnection
|
||||
@@ -15,6 +16,7 @@ from wf_sources_mcp.runtime import (
|
||||
connection_runtime_fingerprint,
|
||||
)
|
||||
from wf_sources_mcp.runtime.factory import PersistentSessionFactory
|
||||
from wf_sources_mcp.sdk import ToolCallResult
|
||||
from wf_sources_mcp.transports import StdioSourceTransport
|
||||
|
||||
|
||||
@@ -28,14 +30,11 @@ def _connection() -> McpSourceConnection:
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_persistent_session_call_callback_normalizes_tool_result() -> None:
|
||||
async def call_tool(tool_name: str, payload: dict[str, Any]) -> CallToolResult:
|
||||
async def test_persistent_session_call_callback_returns_canonical_result() -> None:
|
||||
async def call_tool(tool_name: str, payload: dict[str, Any]) -> ToolCallResult:
|
||||
assert tool_name == "echo"
|
||||
assert payload == {"text": "hi"}
|
||||
return CallToolResult(
|
||||
content=[TextContent(type="text", text="ok")],
|
||||
structuredContent={"echoed": "hi"},
|
||||
)
|
||||
return ToolCallResult(outcome="ok", output={"echoed": "hi"})
|
||||
|
||||
session = PersistentMcpSession(
|
||||
connection=_connection(),
|
||||
@@ -60,32 +59,30 @@ async def test_persistent_session_raises_without_transport() -> None:
|
||||
class _FakeFactory(PersistentSessionFactory):
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[tuple[str, dict[str, object]]] = []
|
||||
self.closed = False
|
||||
self.created_connections: list[McpSourceConnection] = []
|
||||
|
||||
async def _call_tool(
|
||||
self, tool_name: str, payload: dict[str, object]
|
||||
) -> CallToolResult:
|
||||
) -> RawCallToolResult:
|
||||
self.calls.append((tool_name, payload))
|
||||
return CallToolResult(
|
||||
return RawCallToolResult(
|
||||
content=[TextContent(type="text", text="ok")],
|
||||
structuredContent={"echoed": payload["text"]},
|
||||
)
|
||||
|
||||
async def _close(self) -> None:
|
||||
self.closed = True
|
||||
|
||||
async def _create_with_stack(
|
||||
self,
|
||||
stack: AsyncExitStack,
|
||||
connection: McpSourceConnection,
|
||||
auth: AuthRecord | None,
|
||||
) -> ClientSession:
|
||||
self.created_connections.append(connection)
|
||||
factory = self
|
||||
|
||||
class _FakeClient:
|
||||
async def call_tool(
|
||||
self, tool_name: str, payload: dict[str, object]
|
||||
) -> CallToolResult:
|
||||
) -> RawCallToolResult:
|
||||
return await factory._call_tool(tool_name, payload)
|
||||
|
||||
return _FakeClient() # type: ignore[return-value]
|
||||
@@ -94,7 +91,8 @@ class _FakeFactory(PersistentSessionFactory):
|
||||
@pytest.mark.asyncio
|
||||
async def test_persistent_session_factory_serializes_tool_calls() -> None:
|
||||
factory = _FakeFactory()
|
||||
session = await factory.create(_connection(), None)
|
||||
connection = _connection()
|
||||
session = await factory.create(connection, None)
|
||||
|
||||
first = await session.call_tool("echo", {"text": "one"})
|
||||
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 second.output == {"echoed": "two"}
|
||||
assert factory.created_connections == [connection]
|
||||
assert factory.calls == [
|
||||
("echo", {"text": "one"}),
|
||||
("echo", {"text": "two"}),
|
||||
@@ -117,11 +116,8 @@ async def test_runtime_pool_reuses_unchanged_connection() -> None:
|
||||
) -> PersistentMcpSession:
|
||||
created.append(connection)
|
||||
|
||||
async def _call(tool_name: str, payload: dict[str, Any]) -> CallToolResult:
|
||||
return CallToolResult(
|
||||
content=[TextContent(type="text", text="ok")],
|
||||
structuredContent={"echoed": payload["text"]},
|
||||
)
|
||||
async def _call(tool_name: str, payload: dict[str, Any]) -> ToolCallResult:
|
||||
return ToolCallResult(outcome="ok", output={"echoed": payload["text"]})
|
||||
|
||||
return PersistentMcpSession(
|
||||
connection=connection,
|
||||
@@ -150,3 +146,17 @@ def test_runtime_fingerprint_changes_when_transport_changes() -> None:
|
||||
assert connection_runtime_fingerprint(original) != connection_runtime_fingerprint(
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user