refactor: move mcp runtime to wf_sources_mcp
This commit is contained in:
@@ -227,11 +227,13 @@ implementation state.
|
||||
[2026-06-07 MCP source connection seam](./historical/superpowers/plans/2026-06-07-mcp-source-connection-seam.md).
|
||||
- Completed: shared MCP session opener exists in `wf_sources_mcp.client`.
|
||||
One-shot adapter (`McpSdkAdapter`) and persistent runtime
|
||||
(`PersistentSessionFactory`) both use it. Runtime files remain in
|
||||
`wf_mcp` for compatibility. Next slice can move `PersistentSessionFactory`,
|
||||
`PersistentMcpSession`, and `McpRuntimePool` to
|
||||
`wf_sources_mcp.runtime`. The completed plan was
|
||||
[2026-06-07 MCP client session opener](./historical/superpowers/plans/2026-06-07-mcp-client-session-opener.md).
|
||||
(`PersistentSessionFactory`) both use it.
|
||||
- Completed: persistent MCP runtime moved to `wf_sources_mcp.runtime`.
|
||||
`PersistentMcpSession`, `PersistentSessionFactory`, `McpRuntimePool`,
|
||||
and `connection_runtime_fingerprint` are now canonical in
|
||||
`wf_sources_mcp.runtime`; `wf_mcp.runtime.*` are compatibility shims.
|
||||
Runtime remains tool-call-only. The completed plan was
|
||||
[2026-06-07 MCP runtime package move](./historical/superpowers/plans/2026-06-07-mcp-runtime-package-move.md).
|
||||
- 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
|
||||
|
||||
@@ -0,0 +1,588 @@
|
||||
# MCP Runtime Package Move Implementation Plan
|
||||
|
||||
> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.
|
||||
|
||||
**Goal:** Move persistent MCP runtime ownership from `wf_mcp.runtime` to `wf_sources_mcp.runtime` while preserving old imports as compatibility shims.
|
||||
|
||||
**Architecture:** The typed `McpSourceConnection` seam and shared `open_mcp_session()` now exist. This slice makes `wf_sources_mcp.runtime` canonical for persistent MCP sessions, pool reuse, and connection fingerprinting. `wf_mcp.runtime.*` should become thin re-export shims only; behavior should not change and persistent runtime remains tool-call-only.
|
||||
|
||||
**Tech Stack:** Python 3.14, dataclasses, asyncio actor/queue pattern, MCP Python SDK `ClientSession`, pytest, ruff, basedpyright.
|
||||
|
||||
---
|
||||
|
||||
## Current State
|
||||
|
||||
Canonical source-provider code already exists:
|
||||
|
||||
- `wf_sources_mcp.connections.McpSourceConnection`
|
||||
- `wf_sources_mcp.client.open_mcp_session`
|
||||
- `wf_sources_mcp.sdk.ToolCallResult`
|
||||
- `wf_sources_mcp.sdk.converters.tool_result_to_call_result`
|
||||
|
||||
Old runtime files still live in `wf_mcp`:
|
||||
|
||||
- `src/wf_mcp/runtime/factory.py`
|
||||
- `src/wf_mcp/runtime/session.py`
|
||||
- `src/wf_mcp/runtime/pool.py`
|
||||
|
||||
The current `McpRuntimePool` has temporary compatibility glue:
|
||||
|
||||
```text
|
||||
McpSourceConnection -> _legacy_connection_config() -> PersistentSessionFactory
|
||||
```
|
||||
|
||||
After this plan, that back-conversion should disappear. The canonical runtime factory should accept `McpSourceConnection` directly.
|
||||
|
||||
---
|
||||
|
||||
## Non-Goals
|
||||
|
||||
- Do not broaden persistent runtime beyond `call_tool`.
|
||||
- Do not add persistent `read_resource`, `get_prompt`, `invoke_method`, or `send_notification`.
|
||||
- Do not move `McpSdkAdapter`.
|
||||
- Do not touch MCP proxy/frontend transport.
|
||||
- Do not change workflow runtime semantics.
|
||||
- Do not change on-disk auth/catalog/source registry formats.
|
||||
|
||||
---
|
||||
|
||||
## Target File Structure
|
||||
|
||||
Create:
|
||||
|
||||
- `src/wf_sources_mcp/runtime/__init__.py`
|
||||
- `src/wf_sources_mcp/runtime/session.py`
|
||||
- `src/wf_sources_mcp/runtime/factory.py`
|
||||
- `src/wf_sources_mcp/runtime/pool.py`
|
||||
- `tests/wf_sources_mcp/test_runtime.py`
|
||||
|
||||
Modify:
|
||||
|
||||
- `src/wf_mcp/runtime/__init__.py` -> re-export shim
|
||||
- `src/wf_mcp/runtime/session.py` -> re-export shim
|
||||
- `src/wf_mcp/runtime/factory.py` -> re-export shim
|
||||
- `src/wf_mcp/runtime/pool.py` -> re-export shim
|
||||
- `src/wf_mcp/broker/config.py` -> canonical import from `wf_sources_mcp.runtime`
|
||||
- any other production imports found by `rg 'wf_mcp\\.runtime' src`
|
||||
- `tests/wf_mcp/test_compat_imports.py` -> shim identity tests
|
||||
- `docs/current_roadmap.md`
|
||||
- `docs/superpowers/specs/2026-06-03-long-lived-workflow-api-boundary.md`
|
||||
|
||||
---
|
||||
|
||||
## Task 1: Create Canonical Runtime Session
|
||||
|
||||
**Files:**
|
||||
- Create: `src/wf_sources_mcp/runtime/session.py`
|
||||
- Test: `tests/wf_sources_mcp/test_runtime.py`
|
||||
|
||||
- [ ] **Step 1: Add session tests**
|
||||
|
||||
Create `tests/wf_sources_mcp/test_runtime.py` with:
|
||||
|
||||
```python
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from mcp.types import CallToolResult, TextContent
|
||||
|
||||
from wf_sources_mcp.auth import AuthRecord
|
||||
from wf_sources_mcp.connections import McpSourceConnection
|
||||
from wf_sources_mcp.runtime import PersistentMcpSession
|
||||
from wf_sources_mcp.transports import StdioSourceTransport
|
||||
|
||||
|
||||
def _connection() -> McpSourceConnection:
|
||||
return McpSourceConnection(
|
||||
id="demo.personal",
|
||||
provider="demo",
|
||||
account="personal",
|
||||
transport=StdioSourceTransport(command="fake"),
|
||||
)
|
||||
|
||||
|
||||
@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:
|
||||
assert tool_name == "echo"
|
||||
assert payload == {"text": "hi"}
|
||||
return CallToolResult(
|
||||
content=[TextContent(type="text", text="ok")],
|
||||
structuredContent={"echoed": "hi"},
|
||||
)
|
||||
|
||||
session = PersistentMcpSession(
|
||||
connection=_connection(),
|
||||
auth=AuthRecord(connection_id="demo.personal", scheme="none"),
|
||||
call_callback=call_tool,
|
||||
)
|
||||
|
||||
result = await session.call_tool("echo", {"text": "hi"})
|
||||
|
||||
assert result.outcome == "ok"
|
||||
assert result.output == {"echoed": "hi"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_persistent_session_raises_without_transport() -> None:
|
||||
session = PersistentMcpSession(connection=_connection(), auth=None)
|
||||
|
||||
with pytest.raises(RuntimeError, match="no tool call transport"):
|
||||
await session.call_tool("echo", {})
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Run failing tests**
|
||||
|
||||
Run:
|
||||
|
||||
```bash
|
||||
uv run pytest tests/wf_sources_mcp/test_runtime.py -q
|
||||
```
|
||||
|
||||
Expected: fail because `wf_sources_mcp.runtime` does not exist.
|
||||
|
||||
- [ ] **Step 3: Implement canonical session**
|
||||
|
||||
Create `src/wf_sources_mcp/runtime/session.py` by moving the implementation from `src/wf_mcp/runtime/session.py`, but change imports/types:
|
||||
|
||||
- Import `AuthRecord` from `wf_sources_mcp.auth`.
|
||||
- Import `McpSourceConnection` from `wf_sources_mcp.connections`.
|
||||
- Import `ToolCallResult` and `tool_result_to_call_result` from `wf_sources_mcp`.
|
||||
- `PersistentMcpSession.connection` must be `McpSourceConnection`, not `ConnectionConfig`.
|
||||
|
||||
Keep:
|
||||
|
||||
- `RawToolCaller`
|
||||
- `client` injection path
|
||||
- `call_callback` path
|
||||
- `close_callback`
|
||||
- error message `"persistent MCP session has no tool call transport"`
|
||||
|
||||
- [ ] **Step 4: Run session tests**
|
||||
|
||||
Run:
|
||||
|
||||
```bash
|
||||
uv run pytest tests/wf_sources_mcp/test_runtime.py -q
|
||||
uv run basedpyright --level error src/wf_sources_mcp/runtime
|
||||
```
|
||||
|
||||
Expected: pass.
|
||||
|
||||
---
|
||||
|
||||
## Task 2: Create Canonical Runtime Factory
|
||||
|
||||
**Files:**
|
||||
- Create: `src/wf_sources_mcp/runtime/factory.py`
|
||||
- Modify: `tests/wf_sources_mcp/test_runtime.py`
|
||||
|
||||
- [ ] **Step 1: Add factory owner-task tests**
|
||||
|
||||
Append to `tests/wf_sources_mcp/test_runtime.py`:
|
||||
|
||||
```python
|
||||
from wf_sources_mcp.runtime.factory import PersistentSessionFactory
|
||||
|
||||
|
||||
class _FakeFactory(PersistentSessionFactory):
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[tuple[str, dict[str, object]]] = []
|
||||
self.closed = False
|
||||
|
||||
async def _call_tool(self, tool_name: str, payload: dict[str, object]):
|
||||
self.calls.append((tool_name, payload))
|
||||
return CallToolResult(
|
||||
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, connection, auth):
|
||||
class _FakeClient:
|
||||
async def call_tool(self, tool_name, payload):
|
||||
return await self_factory._call_tool(tool_name, payload)
|
||||
|
||||
self_factory = self
|
||||
return _FakeClient()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_persistent_session_factory_serializes_tool_calls() -> None:
|
||||
factory = _FakeFactory()
|
||||
session = await factory.create(_connection(), None)
|
||||
|
||||
first = await session.call_tool("echo", {"text": "one"})
|
||||
second = await session.call_tool("echo", {"text": "two"})
|
||||
await session.close()
|
||||
|
||||
assert first.output == {"echoed": "one"}
|
||||
assert second.output == {"echoed": "two"}
|
||||
assert factory.calls == [
|
||||
("echo", {"text": "one"}),
|
||||
("echo", {"text": "two"}),
|
||||
]
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Implement factory**
|
||||
|
||||
Create `src/wf_sources_mcp/runtime/factory.py` by moving the implementation from `src/wf_mcp/runtime/factory.py`, but change types/imports:
|
||||
|
||||
- `PersistentSessionFactory.create(connection: McpSourceConnection, auth: AuthRecord | None)`.
|
||||
- `_SessionOwner.connection: McpSourceConnection`.
|
||||
- `_create_with_stack(stack, connection: McpSourceConnection, auth)` uses:
|
||||
|
||||
```python
|
||||
session = await stack.enter_async_context(open_mcp_session(connection, auth))
|
||||
return session
|
||||
```
|
||||
|
||||
- Remove any `ConnectionConfig` import.
|
||||
- Keep `_ToolCallRequest` and `_SessionOwner` private.
|
||||
- Keep actor/queue behavior unchanged.
|
||||
|
||||
- [ ] **Step 3: Run factory tests**
|
||||
|
||||
Run:
|
||||
|
||||
```bash
|
||||
uv run pytest tests/wf_sources_mcp/test_runtime.py -q
|
||||
uv run basedpyright --level error src/wf_sources_mcp/runtime
|
||||
```
|
||||
|
||||
Expected: pass.
|
||||
|
||||
---
|
||||
|
||||
## Task 3: Create Canonical Runtime Pool
|
||||
|
||||
**Files:**
|
||||
- Create: `src/wf_sources_mcp/runtime/pool.py`
|
||||
- Modify: `src/wf_sources_mcp/runtime/__init__.py`
|
||||
- Modify: `tests/wf_sources_mcp/test_runtime.py`
|
||||
|
||||
- [ ] **Step 1: Add pool tests**
|
||||
|
||||
Append to `tests/wf_sources_mcp/test_runtime.py`:
|
||||
|
||||
```python
|
||||
from wf_sources_mcp.runtime import McpRuntimePool, connection_runtime_fingerprint
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runtime_pool_reuses_unchanged_connection() -> None:
|
||||
created: list[McpSourceConnection] = []
|
||||
|
||||
async def create_session(connection: McpSourceConnection, auth: AuthRecord | None):
|
||||
created.append(connection)
|
||||
return PersistentMcpSession(
|
||||
connection=connection,
|
||||
auth=auth,
|
||||
call_callback=lambda tool_name, payload: CallToolResult(
|
||||
content=[TextContent(type="text", text="ok")],
|
||||
structuredContent={"echoed": payload["text"]},
|
||||
),
|
||||
)
|
||||
|
||||
pool = McpRuntimePool(session_factory=create_session)
|
||||
connection = _connection()
|
||||
|
||||
await pool.call_tool(connection, None, "echo", {"text": "one"})
|
||||
await pool.call_tool(connection, None, "echo", {"text": "two"})
|
||||
|
||||
assert created == [connection]
|
||||
|
||||
|
||||
def test_runtime_fingerprint_changes_when_transport_changes() -> None:
|
||||
original = _connection()
|
||||
changed = McpSourceConnection(
|
||||
id="demo.personal",
|
||||
provider="demo",
|
||||
account="personal",
|
||||
transport=StdioSourceTransport(command="changed"),
|
||||
)
|
||||
|
||||
assert connection_runtime_fingerprint(original) != connection_runtime_fingerprint(
|
||||
changed
|
||||
)
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Implement pool**
|
||||
|
||||
Create `src/wf_sources_mcp/runtime/pool.py` by moving the implementation from `src/wf_mcp/runtime/pool.py`, but make it canonical:
|
||||
|
||||
- `RuntimeConnection` should be `McpSourceConnection`.
|
||||
- `SessionFactory` should accept `McpSourceConnection`, not `ConnectionConfig`.
|
||||
- Remove `_legacy_connection_config`.
|
||||
- Remove imports from `wf_mcp`.
|
||||
- `connection_runtime_fingerprint` should accept `McpSourceConnection`.
|
||||
- Keep reuse/close behavior unchanged.
|
||||
|
||||
Create `src/wf_sources_mcp/runtime/__init__.py`:
|
||||
|
||||
```python
|
||||
from .factory import PersistentSessionFactory
|
||||
from .pool import McpRuntimePool, connection_runtime_fingerprint
|
||||
from .session import PersistentMcpSession
|
||||
|
||||
__all__ = [
|
||||
"McpRuntimePool",
|
||||
"PersistentMcpSession",
|
||||
"PersistentSessionFactory",
|
||||
"connection_runtime_fingerprint",
|
||||
]
|
||||
```
|
||||
|
||||
- [ ] **Step 3: Run runtime tests**
|
||||
|
||||
Run:
|
||||
|
||||
```bash
|
||||
uv run pytest tests/wf_sources_mcp/test_runtime.py -q
|
||||
uv run basedpyright --level error src/wf_sources_mcp/runtime
|
||||
```
|
||||
|
||||
Expected: pass.
|
||||
|
||||
---
|
||||
|
||||
## Task 4: Turn `wf_mcp.runtime` Into Compatibility Shims
|
||||
|
||||
**Files:**
|
||||
- Replace: `src/wf_mcp/runtime/session.py`
|
||||
- Replace: `src/wf_mcp/runtime/factory.py`
|
||||
- Replace: `src/wf_mcp/runtime/pool.py`
|
||||
- Modify: `src/wf_mcp/runtime/__init__.py`
|
||||
- Test: `tests/wf_mcp/test_compat_imports.py`
|
||||
|
||||
- [ ] **Step 1: Add shim identity tests**
|
||||
|
||||
In `tests/wf_mcp/test_compat_imports.py`, add:
|
||||
|
||||
```python
|
||||
def test_runtime_shims_reexport_wf_sources_mcp_runtime() -> None:
|
||||
from wf_mcp.runtime import (
|
||||
McpRuntimePool as OldMcpRuntimePool,
|
||||
PersistentMcpSession as OldPersistentMcpSession,
|
||||
PersistentSessionFactory as OldPersistentSessionFactory,
|
||||
connection_runtime_fingerprint as old_connection_runtime_fingerprint,
|
||||
)
|
||||
from wf_sources_mcp.runtime import (
|
||||
McpRuntimePool,
|
||||
PersistentMcpSession,
|
||||
PersistentSessionFactory,
|
||||
connection_runtime_fingerprint,
|
||||
)
|
||||
|
||||
assert OldMcpRuntimePool is McpRuntimePool
|
||||
assert OldPersistentMcpSession is PersistentMcpSession
|
||||
assert OldPersistentSessionFactory is PersistentSessionFactory
|
||||
assert old_connection_runtime_fingerprint is connection_runtime_fingerprint
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Replace old runtime files with shims**
|
||||
|
||||
`src/wf_mcp/runtime/session.py`:
|
||||
|
||||
```python
|
||||
"""Compatibility shim for the canonical MCP source runtime session."""
|
||||
|
||||
from wf_sources_mcp.runtime.session import PersistentMcpSession, RawToolCaller
|
||||
|
||||
__all__ = ["PersistentMcpSession", "RawToolCaller"]
|
||||
```
|
||||
|
||||
`src/wf_mcp/runtime/factory.py`:
|
||||
|
||||
```python
|
||||
"""Compatibility shim for the canonical MCP source runtime factory."""
|
||||
|
||||
from wf_sources_mcp.runtime.factory import PersistentSessionFactory
|
||||
|
||||
__all__ = ["PersistentSessionFactory"]
|
||||
```
|
||||
|
||||
`src/wf_mcp/runtime/pool.py`:
|
||||
|
||||
```python
|
||||
"""Compatibility shim for the canonical MCP source runtime pool."""
|
||||
|
||||
from wf_sources_mcp.runtime.pool import (
|
||||
McpRuntimePool,
|
||||
SessionFactory,
|
||||
connection_runtime_fingerprint,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"McpRuntimePool",
|
||||
"SessionFactory",
|
||||
"connection_runtime_fingerprint",
|
||||
]
|
||||
```
|
||||
|
||||
Keep `src/wf_mcp/runtime/protocols.py` as-is if it already shims `ToolExecutor`.
|
||||
|
||||
Update `src/wf_mcp/runtime/__init__.py` to re-export from `wf_sources_mcp.runtime` plus `ToolExecutor`:
|
||||
|
||||
```python
|
||||
from wf_sources_mcp.runtime import (
|
||||
McpRuntimePool,
|
||||
PersistentMcpSession,
|
||||
PersistentSessionFactory,
|
||||
connection_runtime_fingerprint,
|
||||
)
|
||||
|
||||
from .protocols import ToolExecutor
|
||||
|
||||
__all__ = [
|
||||
"McpRuntimePool",
|
||||
"PersistentMcpSession",
|
||||
"PersistentSessionFactory",
|
||||
"ToolExecutor",
|
||||
"connection_runtime_fingerprint",
|
||||
]
|
||||
```
|
||||
|
||||
- [ ] **Step 3: Run shim tests**
|
||||
|
||||
Run:
|
||||
|
||||
```bash
|
||||
uv run pytest tests/wf_mcp/test_compat_imports.py tests/wf_sources_mcp/test_runtime.py -q
|
||||
```
|
||||
|
||||
Expected: pass.
|
||||
|
||||
---
|
||||
|
||||
## Task 5: Update Production Imports To Canonical Runtime
|
||||
|
||||
**Files:**
|
||||
- Modify: `src/wf_mcp/broker/config.py`
|
||||
- Search all source files
|
||||
|
||||
- [ ] **Step 1: Find old runtime imports**
|
||||
|
||||
Run:
|
||||
|
||||
```bash
|
||||
rg -n 'wf_mcp\.runtime|from \.\.runtime|from \.runtime' src tests
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Update production imports**
|
||||
|
||||
For production code outside shim files, import canonical runtime from `wf_sources_mcp.runtime`.
|
||||
|
||||
Likely file:
|
||||
|
||||
`src/wf_mcp/broker/config.py`
|
||||
|
||||
Replace:
|
||||
|
||||
```python
|
||||
from wf_mcp.runtime import McpRuntimePool, PersistentSessionFactory
|
||||
```
|
||||
|
||||
or relative equivalents with:
|
||||
|
||||
```python
|
||||
from wf_sources_mcp.runtime import McpRuntimePool, PersistentSessionFactory
|
||||
```
|
||||
|
||||
Do not update tests that intentionally verify compatibility shims.
|
||||
|
||||
- [ ] **Step 3: Run focused production tests**
|
||||
|
||||
Run:
|
||||
|
||||
```bash
|
||||
uv run pytest tests/wf_mcp/test_stateful_runtime.py tests/wf_mcp/server/test_config.py::test_server_reuses_real_upstream_session_across_workflow_requests -q
|
||||
uv run basedpyright --level error src
|
||||
```
|
||||
|
||||
Expected: pass.
|
||||
|
||||
---
|
||||
|
||||
## Task 6: Preserve Existing Stateful Runtime Tests
|
||||
|
||||
**Files:**
|
||||
- Modify: `tests/wf_mcp/test_stateful_runtime.py` only if required
|
||||
- Test: `tests/wf_mcp/test_stateful_runtime.py`
|
||||
|
||||
- [ ] **Step 1: Run existing tests unchanged first**
|
||||
|
||||
Run:
|
||||
|
||||
```bash
|
||||
uv run pytest tests/wf_mcp/test_stateful_runtime.py -q
|
||||
```
|
||||
|
||||
Expected: should pass through shims. If it fails only because helper subclasses still type `ConnectionConfig`, update the tests to import canonical runtime but keep behavior assertions unchanged.
|
||||
|
||||
- [ ] **Step 2: Do not weaken behavior assertions**
|
||||
|
||||
The following behavior must remain tested:
|
||||
|
||||
- pool reuses unchanged connection fingerprint
|
||||
- pool replaces changed connection fingerprint
|
||||
- session owner serializes calls through one owner task
|
||||
- closed sessions are closed via callback
|
||||
- crashing factory surfaces errors to queued calls
|
||||
|
||||
If any of these tests need edits, preserve the same assertions and explain why in final report.
|
||||
|
||||
---
|
||||
|
||||
## Task 7: Documentation And Verification
|
||||
|
||||
**Files:**
|
||||
- Modify: `docs/current_roadmap.md`
|
||||
- Modify: `docs/superpowers/specs/2026-06-03-long-lived-workflow-api-boundary.md`
|
||||
|
||||
- [ ] **Step 1: Update docs**
|
||||
|
||||
Roadmap/spec should say:
|
||||
|
||||
- persistent MCP runtime moved to `wf_sources_mcp.runtime`
|
||||
- `wf_mcp.runtime.*` are compatibility shims
|
||||
- runtime remains tool-call-only
|
||||
- next slice is moving `McpSdkAdapter` to `wf_sources_mcp.sdk.adapter`
|
||||
|
||||
- [ ] **Step 2: Final verification**
|
||||
|
||||
Run:
|
||||
|
||||
```bash
|
||||
uv run pytest tests/wf_sources_mcp tests/wf_mcp/test_stateful_runtime.py tests/wf_mcp/test_compat_imports.py tests/wf_mcp/server/test_config.py::test_server_reuses_real_upstream_session_across_workflow_requests -q
|
||||
uv run ruff check src tests
|
||||
uv run basedpyright --level error src
|
||||
git diff --check
|
||||
```
|
||||
|
||||
Expected:
|
||||
|
||||
- focused tests pass
|
||||
- ruff passes
|
||||
- basedpyright has 0 errors
|
||||
- no whitespace errors
|
||||
|
||||
If `ruff check src tests` finds unrelated pre-existing errors, do not fix unrelated files in this slice. Report exact files/errors and run `ruff check` on changed files instead.
|
||||
|
||||
---
|
||||
|
||||
## Final Report Requirements
|
||||
|
||||
The final report must state:
|
||||
|
||||
- runtime files moved or shimmed
|
||||
- no behavior expansion beyond `call_tool`
|
||||
- `wf_mcp.runtime` compatibility status
|
||||
- whether any tests had to change and why
|
||||
- exact verification commands and outputs
|
||||
|
||||
@@ -95,11 +95,14 @@ First slices should move leaf modules only and leave `wf_mcp` re-export shims:
|
||||
5. Complete: MCP SDK conversion helpers moved to `wf_sources_mcp.sdk.converters`, with `wf_mcp.sdk.converters` retained as a shim.
|
||||
6. Complete: shared MCP session opener in `wf_sources_mcp.client`. One-shot
|
||||
adapter (`McpSdkAdapter`) and persistent runtime
|
||||
(`PersistentSessionFactory`) both use `open_mcp_session`. Runtime files
|
||||
remain in `wf_mcp` for compatibility; next slice moves
|
||||
`PersistentSessionFactory`, `PersistentMcpSession`, and `McpRuntimePool`
|
||||
to `wf_sources_mcp.runtime`.
|
||||
7. Upstream transport/discovery/session services.
|
||||
(`PersistentSessionFactory`) both use `open_mcp_session`.
|
||||
7. Complete: persistent MCP runtime (`PersistentMcpSession`,
|
||||
`PersistentSessionFactory`, `McpRuntimePool`,
|
||||
`connection_runtime_fingerprint`) moved to `wf_sources_mcp.runtime`, with
|
||||
`wf_mcp.runtime.*` retained as compatibility shims. Runtime remains
|
||||
tool-call-only. Next slice is moving `McpSdkAdapter` to
|
||||
`wf_sources_mcp.sdk.adapter`.
|
||||
8. 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`,
|
||||
|
||||
@@ -6,6 +6,7 @@ from pathlib import Path
|
||||
from wf_api import file_workflow_stores
|
||||
from wf_config import WorkflowConfigFile
|
||||
from wf_config.models import FilesystemStoreConfig, McpSourceConfig, ServerConfig
|
||||
from wf_sources_mcp.runtime import McpRuntimePool, PersistentSessionFactory
|
||||
from wf_sources_mcp.source_registry import (
|
||||
FileSourceRegistryStore,
|
||||
workflow_mcp_source_to_connection_config,
|
||||
@@ -14,7 +15,6 @@ from wf_sources_mcp.storage import FileAuthStore, FileCatalogStore, FileStore
|
||||
|
||||
from ..control import BrokerConfigFile, ConnectionConfigFile
|
||||
from ..models import BrokerConfig
|
||||
from ..runtime import McpRuntimePool, PersistentSessionFactory
|
||||
from ..sdk import McpSdkAdapter
|
||||
from .models import BrokerStoreRoots
|
||||
from .service import WfMcpService
|
||||
|
||||
@@ -1,7 +1,11 @@
|
||||
from .factory import PersistentSessionFactory
|
||||
from .pool import McpRuntimePool, connection_runtime_fingerprint
|
||||
from wf_sources_mcp.runtime import (
|
||||
McpRuntimePool,
|
||||
PersistentMcpSession,
|
||||
PersistentSessionFactory,
|
||||
connection_runtime_fingerprint,
|
||||
)
|
||||
|
||||
from .protocols import ToolExecutor
|
||||
from .session import PersistentMcpSession
|
||||
|
||||
__all__ = [
|
||||
"McpRuntimePool",
|
||||
|
||||
@@ -1,156 +1,5 @@
|
||||
from __future__ import annotations
|
||||
"""Compatibility shim for the canonical MCP source runtime factory."""
|
||||
|
||||
import asyncio
|
||||
from contextlib import AsyncExitStack
|
||||
from dataclasses import dataclass, field
|
||||
from wf_sources_mcp.runtime.factory import PersistentSessionFactory
|
||||
|
||||
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 mcp_source_connection_from_connection_config
|
||||
|
||||
from ..models import ConnectionConfig
|
||||
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: ConnectionConfig,
|
||||
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: ConnectionConfig,
|
||||
auth: AuthRecord | None,
|
||||
) -> ClientSession:
|
||||
source_connection = mcp_source_connection_from_connection_config(connection)
|
||||
session = await stack.enter_async_context(
|
||||
open_mcp_session(source_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: ConnectionConfig
|
||||
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
|
||||
__all__ = ["PersistentSessionFactory"]
|
||||
|
||||
+10
-141
@@ -1,144 +1,13 @@
|
||||
from __future__ import annotations
|
||||
"""Compatibility shim for the canonical MCP source runtime pool."""
|
||||
|
||||
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.runtime.pool import (
|
||||
McpRuntimePool,
|
||||
SessionFactory,
|
||||
connection_runtime_fingerprint,
|
||||
)
|
||||
|
||||
from wf_sources_mcp.connections import McpSourceConnection
|
||||
from wf_sources_mcp.sdk import ToolCallResult
|
||||
from wf_sources_mcp.transports import HttpSourceTransport, StdioSourceTransport
|
||||
|
||||
from ..auth import AuthRecord
|
||||
from ..models import ConnectionConfig
|
||||
from .session import PersistentMcpSession
|
||||
|
||||
RuntimeConnection = ConnectionConfig | McpSourceConnection
|
||||
SessionFactory = Callable[
|
||||
[ConnectionConfig, AuthRecord | None],
|
||||
PersistentMcpSession | Awaitable[PersistentMcpSession],
|
||||
__all__ = [
|
||||
"McpRuntimePool",
|
||||
"SessionFactory",
|
||||
"connection_runtime_fingerprint",
|
||||
]
|
||||
|
||||
|
||||
def connection_runtime_fingerprint(
|
||||
connection: RuntimeConnection,
|
||||
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 `ConnectionConfig` 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: RuntimeConnection,
|
||||
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()
|
||||
|
||||
# Compatibility boundary: wrappers now pass McpSourceConnection, while
|
||||
# PersistentSessionFactory still consumes the legacy broker DTO until
|
||||
# the shared opener/runtime move lands.
|
||||
created = self.session_factory(_legacy_connection_config(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: RuntimeConnection,
|
||||
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()
|
||||
|
||||
|
||||
def _legacy_connection_config(connection: RuntimeConnection) -> ConnectionConfig:
|
||||
if isinstance(connection, ConnectionConfig):
|
||||
return connection
|
||||
|
||||
metadata = dict(connection.metadata)
|
||||
transport = connection.transport
|
||||
if isinstance(transport, StdioSourceTransport):
|
||||
metadata.update(
|
||||
{
|
||||
"transport": "stdio",
|
||||
"command": transport.command,
|
||||
"args": list(transport.args),
|
||||
"env": dict(transport.env),
|
||||
}
|
||||
)
|
||||
if transport.cwd is not None:
|
||||
metadata["cwd"] = transport.cwd
|
||||
elif isinstance(transport, HttpSourceTransport):
|
||||
metadata.update(
|
||||
{
|
||||
"transport": "streamable_http",
|
||||
"url": str(transport.url),
|
||||
"headers": dict(transport.headers),
|
||||
}
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"connection {connection.id!r} requires metadata.transport")
|
||||
|
||||
if connection.profile is not None:
|
||||
metadata["profile"] = connection.profile
|
||||
if connection.auth_ref is not None:
|
||||
metadata["auth_ref"] = connection.auth_ref
|
||||
|
||||
return ConnectionConfig(
|
||||
id=connection.id,
|
||||
server=connection.provider,
|
||||
account=connection.account,
|
||||
enabled=connection.enabled,
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
@@ -1,49 +1,5 @@
|
||||
from __future__ import annotations
|
||||
"""Compatibility shim for the canonical MCP source runtime session."""
|
||||
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
from wf_sources_mcp.runtime.session import PersistentMcpSession, RawToolCaller
|
||||
|
||||
from mcp.client.session import ClientSession
|
||||
from mcp.types import CallToolResult
|
||||
|
||||
from wf_sources_mcp.sdk import ToolCallResult
|
||||
from wf_sources_mcp.sdk.converters import tool_result_to_call_result
|
||||
|
||||
from ..auth import AuthRecord
|
||||
from ..models import ConnectionConfig
|
||||
|
||||
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: ConnectionConfig
|
||||
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()
|
||||
__all__ = ["PersistentMcpSession", "RawToolCaller"]
|
||||
|
||||
@@ -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()
|
||||
@@ -147,3 +147,29 @@ def test_wf_mcp_sdk_converter_shim_reexports_wf_sources_mcp_converters() -> None
|
||||
assert compat_tool_result is tool_result_to_call_result
|
||||
assert compat_tool_to_discovered is tool_to_discovered
|
||||
assert compat_output_schema is workflow_output_schema_from_mcp_tool_schema
|
||||
|
||||
|
||||
def test_runtime_shims_reexport_wf_sources_mcp_runtime() -> None:
|
||||
from wf_mcp.runtime import (
|
||||
McpRuntimePool as OldMcpRuntimePool,
|
||||
)
|
||||
from wf_mcp.runtime import (
|
||||
PersistentMcpSession as OldPersistentMcpSession,
|
||||
)
|
||||
from wf_mcp.runtime import (
|
||||
PersistentSessionFactory as OldPersistentSessionFactory,
|
||||
)
|
||||
from wf_mcp.runtime import (
|
||||
connection_runtime_fingerprint as old_connection_runtime_fingerprint,
|
||||
)
|
||||
from wf_sources_mcp.runtime import (
|
||||
McpRuntimePool,
|
||||
PersistentMcpSession,
|
||||
PersistentSessionFactory,
|
||||
connection_runtime_fingerprint,
|
||||
)
|
||||
|
||||
assert OldMcpRuntimePool is McpRuntimePool
|
||||
assert OldPersistentMcpSession is PersistentMcpSession
|
||||
assert OldPersistentSessionFactory is PersistentSessionFactory
|
||||
assert old_connection_runtime_fingerprint is connection_runtime_fingerprint
|
||||
|
||||
@@ -11,7 +11,7 @@ from mcp.types import CallToolResult
|
||||
from wf_authoring import build_async_registry
|
||||
from wf_core import RuntimeContext
|
||||
from wf_mcp.capabilities import DiscoveredTool
|
||||
from wf_mcp.models import AuthRecord, ConnectionConfig
|
||||
from wf_mcp.models import AuthRecord
|
||||
from wf_mcp.runtime import McpRuntimePool, PersistentMcpSession
|
||||
from wf_mcp.runtime.factory import PersistentSessionFactory
|
||||
from wf_mcp.sdk import ToolCallResult
|
||||
@@ -102,7 +102,7 @@ class CrashingSessionFactory(PersistentSessionFactory):
|
||||
async def _create_with_stack(
|
||||
self,
|
||||
stack: AsyncExitStack,
|
||||
connection: ConnectionConfig,
|
||||
connection: McpSourceConnection,
|
||||
auth: AuthRecord | None,
|
||||
) -> ClientSession:
|
||||
return cast(ClientSession, self.client)
|
||||
@@ -161,20 +161,16 @@ def test_generated_workflow_specs_share_injected_tool_executor() -> None:
|
||||
|
||||
|
||||
def test_runtime_pool_reuses_stateful_session_for_same_connection() -> None:
|
||||
connection = ConnectionConfig(
|
||||
connection = McpSourceConnection(
|
||||
id="playwright.default",
|
||||
server="playwright",
|
||||
provider="playwright",
|
||||
account="default",
|
||||
metadata={
|
||||
"transport": "stdio",
|
||||
"command": "pnpx",
|
||||
"args": ["@playwright/mcp"],
|
||||
},
|
||||
transport=StdioSourceTransport(command="pnpx"),
|
||||
)
|
||||
created_clients: list[FakeStatefulClient] = []
|
||||
|
||||
async def factory(
|
||||
connection: ConnectionConfig,
|
||||
connection: McpSourceConnection,
|
||||
auth: AuthRecord | None,
|
||||
) -> PersistentMcpSession:
|
||||
client = FakeStatefulClient()
|
||||
@@ -199,22 +195,22 @@ def test_runtime_pool_reuses_stateful_session_for_same_connection() -> None:
|
||||
|
||||
|
||||
def test_runtime_pool_replaces_session_when_fingerprint_changes() -> None:
|
||||
original = ConnectionConfig(
|
||||
original = McpSourceConnection(
|
||||
id="playwright.default",
|
||||
server="playwright",
|
||||
provider="playwright",
|
||||
account="default",
|
||||
metadata={"transport": "stdio", "command": "pnpx", "args": ["old"]},
|
||||
transport=StdioSourceTransport(command="pnpx"),
|
||||
)
|
||||
changed = ConnectionConfig(
|
||||
changed = McpSourceConnection(
|
||||
id="playwright.default",
|
||||
server="playwright",
|
||||
provider="playwright",
|
||||
account="default",
|
||||
metadata={"transport": "stdio", "command": "pnpx", "args": ["new"]},
|
||||
transport=StdioSourceTransport(command="pnpx-new"),
|
||||
)
|
||||
created_clients: list[FakeStatefulClient] = []
|
||||
|
||||
def factory(
|
||||
connection: ConnectionConfig,
|
||||
connection: McpSourceConnection,
|
||||
auth: AuthRecord | None,
|
||||
) -> PersistentMcpSession:
|
||||
client = FakeStatefulClient()
|
||||
@@ -238,11 +234,10 @@ def test_runtime_pool_replaces_session_when_fingerprint_changes() -> None:
|
||||
|
||||
|
||||
def test_persistent_session_fails_inflight_and_queued_calls_if_owner_dies() -> None:
|
||||
connection = ConnectionConfig(
|
||||
connection = McpSourceConnection(
|
||||
id="failing.default",
|
||||
server="failing",
|
||||
provider="failing",
|
||||
account="default",
|
||||
metadata={},
|
||||
)
|
||||
|
||||
async def exercise() -> tuple[
|
||||
|
||||
@@ -0,0 +1,152 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from contextlib import AsyncExitStack
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from mcp.client.session import ClientSession
|
||||
from mcp.types import CallToolResult, TextContent
|
||||
|
||||
from wf_sources_mcp.auth import AuthRecord
|
||||
from wf_sources_mcp.connections import McpSourceConnection
|
||||
from wf_sources_mcp.runtime import (
|
||||
McpRuntimePool,
|
||||
PersistentMcpSession,
|
||||
connection_runtime_fingerprint,
|
||||
)
|
||||
from wf_sources_mcp.runtime.factory import PersistentSessionFactory
|
||||
from wf_sources_mcp.transports import StdioSourceTransport
|
||||
|
||||
|
||||
def _connection() -> McpSourceConnection:
|
||||
return McpSourceConnection(
|
||||
id="demo.personal",
|
||||
provider="demo",
|
||||
account="personal",
|
||||
transport=StdioSourceTransport(command="fake"),
|
||||
)
|
||||
|
||||
|
||||
@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:
|
||||
assert tool_name == "echo"
|
||||
assert payload == {"text": "hi"}
|
||||
return CallToolResult(
|
||||
content=[TextContent(type="text", text="ok")],
|
||||
structuredContent={"echoed": "hi"},
|
||||
)
|
||||
|
||||
session = PersistentMcpSession(
|
||||
connection=_connection(),
|
||||
auth=AuthRecord(connection_id="demo.personal", scheme="none"),
|
||||
call_callback=call_tool,
|
||||
)
|
||||
|
||||
result = await session.call_tool("echo", {"text": "hi"})
|
||||
|
||||
assert result.outcome == "ok"
|
||||
assert result.output == {"echoed": "hi"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_persistent_session_raises_without_transport() -> None:
|
||||
session = PersistentMcpSession(connection=_connection(), auth=None)
|
||||
|
||||
with pytest.raises(RuntimeError, match="no tool call transport"):
|
||||
await session.call_tool("echo", {})
|
||||
|
||||
|
||||
class _FakeFactory(PersistentSessionFactory):
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[tuple[str, dict[str, object]]] = []
|
||||
self.closed = False
|
||||
|
||||
async def _call_tool(
|
||||
self, tool_name: str, payload: dict[str, object]
|
||||
) -> CallToolResult:
|
||||
self.calls.append((tool_name, payload))
|
||||
return CallToolResult(
|
||||
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:
|
||||
factory = self
|
||||
|
||||
class _FakeClient:
|
||||
async def call_tool(
|
||||
self, tool_name: str, payload: dict[str, object]
|
||||
) -> CallToolResult:
|
||||
return await factory._call_tool(tool_name, payload)
|
||||
|
||||
return _FakeClient() # type: ignore[return-value]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_persistent_session_factory_serializes_tool_calls() -> None:
|
||||
factory = _FakeFactory()
|
||||
session = await factory.create(_connection(), None)
|
||||
|
||||
first = await session.call_tool("echo", {"text": "one"})
|
||||
second = await session.call_tool("echo", {"text": "two"})
|
||||
await session.close()
|
||||
|
||||
assert first.output == {"echoed": "one"}
|
||||
assert second.output == {"echoed": "two"}
|
||||
assert factory.calls == [
|
||||
("echo", {"text": "one"}),
|
||||
("echo", {"text": "two"}),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runtime_pool_reuses_unchanged_connection() -> None:
|
||||
created: list[McpSourceConnection] = []
|
||||
|
||||
async def create_session(
|
||||
connection: McpSourceConnection, auth: AuthRecord | 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"]},
|
||||
)
|
||||
|
||||
return PersistentMcpSession(
|
||||
connection=connection,
|
||||
auth=auth,
|
||||
call_callback=_call,
|
||||
)
|
||||
|
||||
pool = McpRuntimePool(session_factory=create_session)
|
||||
connection = _connection()
|
||||
|
||||
await pool.call_tool(connection, None, "echo", {"text": "one"})
|
||||
await pool.call_tool(connection, None, "echo", {"text": "two"})
|
||||
|
||||
assert created == [connection]
|
||||
|
||||
|
||||
def test_runtime_fingerprint_changes_when_transport_changes() -> None:
|
||||
original = _connection()
|
||||
changed = McpSourceConnection(
|
||||
id="demo.personal",
|
||||
provider="demo",
|
||||
account="personal",
|
||||
transport=StdioSourceTransport(command="changed"),
|
||||
)
|
||||
|
||||
assert connection_runtime_fingerprint(original) != connection_runtime_fingerprint(
|
||||
changed
|
||||
)
|
||||
Reference in New Issue
Block a user