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
|
`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`,
|
||||||
|
|||||||
@@ -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():
|
||||||
|
|||||||
@@ -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."""
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
Reference in New Issue
Block a user