refactor: move mcp runtime to wf_sources_mcp
This commit is contained in:
@@ -0,0 +1,10 @@
|
||||
from .factory import PersistentSessionFactory
|
||||
from .pool import McpRuntimePool, connection_runtime_fingerprint
|
||||
from .session import PersistentMcpSession
|
||||
|
||||
__all__ = [
|
||||
"McpRuntimePool",
|
||||
"PersistentMcpSession",
|
||||
"PersistentSessionFactory",
|
||||
"connection_runtime_fingerprint",
|
||||
]
|
||||
@@ -0,0 +1,152 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from contextlib import AsyncExitStack
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
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.connections import McpSourceConnection
|
||||
|
||||
from .session import PersistentMcpSession
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class PersistentSessionFactory:
|
||||
"""Create initialized persistent MCP sessions for configured connections.
|
||||
|
||||
Input connection metadata must describe either stdio transport
|
||||
(`command`, optional `args`/`env`/`cwd`) or streamable HTTP transport
|
||||
(`url`). The returned session owns its transport stack and closes it through
|
||||
the `PersistentMcpSession.close_callback`.
|
||||
"""
|
||||
|
||||
async def create(
|
||||
self,
|
||||
connection: McpSourceConnection,
|
||||
auth: AuthRecord | None,
|
||||
) -> PersistentMcpSession:
|
||||
owner = _SessionOwner(factory=self, connection=connection, auth=auth)
|
||||
await owner.start()
|
||||
return PersistentMcpSession(
|
||||
connection=connection,
|
||||
auth=auth,
|
||||
call_callback=owner.call_tool,
|
||||
close_callback=owner.close,
|
||||
)
|
||||
|
||||
async def _create_with_stack(
|
||||
self,
|
||||
stack: AsyncExitStack,
|
||||
connection: McpSourceConnection,
|
||||
auth: AuthRecord | None,
|
||||
) -> ClientSession:
|
||||
session = await stack.enter_async_context(open_mcp_session(connection, auth))
|
||||
return session
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class _ToolCallRequest:
|
||||
"""One request submitted to the task that owns the MCP transport."""
|
||||
|
||||
tool_name: str
|
||||
payload: dict[str, object]
|
||||
result: asyncio.Future[CallToolResult]
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class _SessionOwner:
|
||||
"""Run one MCP client session entirely inside its owning asyncio task.
|
||||
|
||||
MCP SDK transports open AnyIO cancel scopes. Entering a transport in one
|
||||
inbound MCP request and reusing it from another causes
|
||||
`ClosedResourceError`/cancel-scope ownership failures. This actor keeps
|
||||
transport creation, calls, and cleanup in one stable task while the public
|
||||
workflow surface submits requests through a queue.
|
||||
"""
|
||||
|
||||
factory: PersistentSessionFactory
|
||||
connection: McpSourceConnection
|
||||
auth: AuthRecord | None
|
||||
_requests: asyncio.Queue[_ToolCallRequest | None] = field(
|
||||
default_factory=asyncio.Queue
|
||||
)
|
||||
_task: asyncio.Task[None] | None = None
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Start the owner task and wait until its MCP session is initialized."""
|
||||
ready = asyncio.get_running_loop().create_future()
|
||||
self._task = asyncio.create_task(
|
||||
self._run(ready),
|
||||
name=f"wf-mcp-session:{self.connection.id}",
|
||||
)
|
||||
await ready
|
||||
|
||||
async def call_tool(
|
||||
self,
|
||||
tool_name: str,
|
||||
payload: dict[str, object],
|
||||
) -> CallToolResult:
|
||||
"""Submit a call and fail promptly if its transport owner exits."""
|
||||
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()
|
||||
await self._requests.put(
|
||||
_ToolCallRequest(tool_name=tool_name, payload=payload, result=result)
|
||||
)
|
||||
done, _pending = await asyncio.wait(
|
||||
{result, task}, return_when=asyncio.FIRST_COMPLETED
|
||||
)
|
||||
if result in done:
|
||||
return result.result()
|
||||
await task
|
||||
raise RuntimeError("persistent MCP session stopped unexpectedly")
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Ask the owner task to close the MCP transport in its own scope."""
|
||||
task = self._task
|
||||
if task is None:
|
||||
return
|
||||
if not task.done():
|
||||
await self._requests.put(None)
|
||||
await task
|
||||
self._task = None
|
||||
|
||||
async def _run(self, ready: asyncio.Future[None]) -> None:
|
||||
"""Own the complete MCP transport lifecycle and serialized call loop."""
|
||||
try:
|
||||
async with AsyncExitStack() as stack:
|
||||
session = await self.factory._create_with_stack(
|
||||
stack, self.connection, self.auth
|
||||
)
|
||||
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
|
||||
)
|
||||
except Exception as exc:
|
||||
request.result.set_exception(exc)
|
||||
else:
|
||||
request.result.set_result(response)
|
||||
except BaseException as exc:
|
||||
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():
|
||||
pending.result.set_exception(exc)
|
||||
raise
|
||||
@@ -0,0 +1,96 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from inspect import isawaitable
|
||||
from typing import Any, cast
|
||||
|
||||
from wf_sources_mcp.auth import AuthRecord
|
||||
from wf_sources_mcp.connections import McpSourceConnection
|
||||
from wf_sources_mcp.sdk import ToolCallResult
|
||||
|
||||
from .session import PersistentMcpSession
|
||||
|
||||
SessionFactory = Callable[
|
||||
[McpSourceConnection, AuthRecord | None],
|
||||
PersistentMcpSession | Awaitable[PersistentMcpSession],
|
||||
]
|
||||
|
||||
|
||||
def connection_runtime_fingerprint(
|
||||
connection: McpSourceConnection,
|
||||
auth: AuthRecord | None = None,
|
||||
) -> str:
|
||||
"""Return the connection identity that decides MCP runtime reuse.
|
||||
|
||||
This is intentionally transport/auth level, not catalog level. Tool list
|
||||
refreshes should not restart a browser-like MCP session, but changing the
|
||||
command, URL, account, or auth payload must create a fresh session.
|
||||
"""
|
||||
|
||||
return json.dumps(
|
||||
{
|
||||
"connection": asdict(connection),
|
||||
"auth": asdict(auth) if auth is not None else None,
|
||||
},
|
||||
sort_keys=True,
|
||||
separators=(",", ":"),
|
||||
default=str,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class McpRuntimePool:
|
||||
"""Cache one persistent MCP runtime per unchanged connection fingerprint.
|
||||
|
||||
Callers provide full `McpSourceConnection` and optional `AuthRecord` on
|
||||
every call. The pool decides whether that identity still maps to the
|
||||
existing upstream MCP session. If command, URL, account, or auth changes,
|
||||
the old session is closed and replaced.
|
||||
"""
|
||||
|
||||
session_factory: SessionFactory
|
||||
_sessions: dict[str, tuple[str, PersistentMcpSession]] = field(default_factory=dict)
|
||||
|
||||
async def get_session(
|
||||
self,
|
||||
connection: McpSourceConnection,
|
||||
auth: AuthRecord | None,
|
||||
) -> PersistentMcpSession:
|
||||
fingerprint = connection_runtime_fingerprint(connection, auth)
|
||||
current = self._sessions.get(connection.id)
|
||||
if current is not None and current[0] == fingerprint:
|
||||
return current[1]
|
||||
if current is not None:
|
||||
await current[1].close()
|
||||
|
||||
created = self.session_factory(connection, auth)
|
||||
if isawaitable(created):
|
||||
session = await created
|
||||
else:
|
||||
session = cast(PersistentMcpSession, created)
|
||||
self._sessions[connection.id] = (fingerprint, session)
|
||||
return session
|
||||
|
||||
async def call_tool(
|
||||
self,
|
||||
connection: McpSourceConnection,
|
||||
auth: AuthRecord | None,
|
||||
tool_name: str,
|
||||
payload: dict[str, Any],
|
||||
) -> ToolCallResult:
|
||||
session = await self.get_session(connection, auth)
|
||||
return await session.call_tool(tool_name, payload)
|
||||
|
||||
async def close_connection(self, connection_id: str) -> None:
|
||||
current = self._sessions.pop(connection_id, None)
|
||||
if current is not None:
|
||||
await current[1].close()
|
||||
|
||||
async def close_all(self) -> None:
|
||||
"""Close all live runtimes; useful for server shutdown and tests."""
|
||||
sessions = list(self._sessions.values())
|
||||
self._sessions.clear()
|
||||
for _fingerprint, session in sessions:
|
||||
await session.close()
|
||||
@@ -0,0 +1,48 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Awaitable, Callable
|
||||
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]]
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class PersistentMcpSession:
|
||||
"""Long-lived MCP execution handle for one configured connection.
|
||||
|
||||
Production sessions use `call_callback` because MCP transports are entered
|
||||
inside an AnyIO cancel scope and must be called and closed by that same
|
||||
owner task. `client` remains available for simple injected/fake sessions in
|
||||
tests. `call_tool()` always normalizes SDK results for workflow nodes.
|
||||
"""
|
||||
|
||||
connection: McpSourceConnection
|
||||
auth: AuthRecord | None
|
||||
client: ClientSession | None = None
|
||||
call_callback: RawToolCaller | None = None
|
||||
close_callback: Callable[[], Awaitable[None]] | None = None
|
||||
|
||||
async def call_tool(
|
||||
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:
|
||||
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)
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Close the transport/session stack owned by the runtime factory."""
|
||||
if self.close_callback is not None:
|
||||
await self.close_callback()
|
||||
Reference in New Issue
Block a user