stateful proxy

tests/wf_mcp/test_server.py

@@ -8,7 +8,7 @@ from typing import Any
 
 from mcp import types as mcp_types
 
-from wf_mcp.broker.config import load_broker_config
+from wf_mcp.broker.config import build_service_from_config, load_broker_config
 from wf_mcp.models import BrokerConfig, ConnectionConfig
 from wf_mcp.server import create_server_client
 
@@ -21,6 +21,29 @@ def _structured(result: Any) -> dict[str, Any]:
     return content
 
 
+def test_config_built_service_uses_persistent_tool_executor() -> None:
+    config = BrokerConfig(
+        store_root=local_temp_root() / "runtime_config_store",
+        connections=[
+            ConnectionConfig(
+                id="fixture.personal",
+                server="fixture",
+                account="personal",
+                metadata={
+                    "transport": "stdio",
+                    "command": sys.executable,
+                    "args": [fixture_server_path()],
+                },
+            )
+        ],
+    )
+
+    service = build_service_from_config(config)
+
+    assert service.adapters["fixture"].__class__.__name__ == "McpSdkAdapter"
+    assert service.tool_executor is not None
+
+
 async def _assert_safe_tool_maps(
     client: Any,
     *,

tests/wf_mcp/test_service.py

@@ -2,12 +2,15 @@ from __future__ import annotations
 
 import asyncio
 import shutil
+from typing import Any, cast
 
 from wf_artifacts import FileDraftWorkspaceStore
-from wf_authoring import NodeSpec
-from wf_core import END, NodeUse, RunStatus
+from wf_authoring import NodeSpec, build_async_registry
+from wf_core import END, NodeUse, RunStatus, RuntimeContext
 from wf_mcp.broker import WfMcpService
 from wf_mcp.models import AuthRecord, ConnectionConfig, RawWorkflowPlan
+from wf_mcp.runtime import ToolExecutor
+from wf_mcp.sdk import ToolCallResult
 from wf_mcp.shared.errors import error_payload
 from wf_mcp.storage import FileStore
 from wf_platform import (
@@ -1022,6 +1025,45 @@ def test_service_can_call_upstream_tool_directly() -> None:
     assert "tool_call_completed" in event_kinds
 
 
+def test_generated_specs_use_injected_tool_executor() -> None:
+    class RecordingExecutor:
+        def __init__(self) -> None:
+            self.payloads: list[dict[str, Any]] = []
+
+        async def call_tool(
+            self,
+            connection: ConnectionConfig,
+            auth: AuthRecord | None,
+            tool_name: str,
+            payload: dict[str, Any],
+        ) -> ToolCallResult:
+            self.payloads.append(payload)
+            return ToolCallResult(outcome="ok", output={"echoed": payload["text"]})
+
+    executor = RecordingExecutor()
+    service = WfMcpService(
+        store=FileStore(local_temp_root() / "injected_executor_store"),
+        tool_executor=cast(ToolExecutor, executor),
+    )
+    service.register_connection(
+        ConnectionConfig(id="demo.personal", server="demo", account="personal")
+    )
+    service.register_adapter("demo", FakeAdapter())
+
+    asyncio.run(service.refresh_connection_catalog("demo.personal"))
+    spec = service._get_qualified_spec("demo.personal.echo_tool")
+    handler = build_async_registry(spec)[spec.name]
+
+    async def run_node() -> dict[str, Any]:
+        return await handler({"text": "hello"}, RuntimeContext(current_node_id="echo"))
+
+    result = asyncio.run(run_node())
+
+    assert result["outcome"] == "ok"
+    assert result["output"]["echoed"] == "hello"
+    assert executor.payloads == [{"text": "hello"}]
+
+
 def test_service_records_catalog_refresh_failures() -> None:
     service = WfMcpService(store=FileStore(local_temp_root() / "refresh_fail_store"))
     service.register_connection(

tests/wf_mcp/test_stateful_runtime.py

@@ -8,6 +8,7 @@ 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.runtime import McpRuntimePool, PersistentMcpSession
 from wf_mcp.sdk import ToolCallResult
 from wf_mcp.workflow import wrap_discovered_tool
 
@@ -37,6 +38,27 @@ class FakeStatefulExecutor:
         raise KeyError(tool_name)
 
 
+@dataclass(slots=True)
+class FakeStatefulClient:
+    """Session-client fake with the same call shape as MCP SDK ClientSession."""
+
+    page_open: bool = False
+    closed: bool = False
+    calls: list[tuple[str, dict[str, Any]]] = field(default_factory=list)
+
+    async def call_tool(self, tool_name: str, payload: dict[str, Any]) -> ToolCallResult:
+        self.calls.append((tool_name, payload))
+        if tool_name == "browser_navigate":
+            self.page_open = True
+            return ToolCallResult(outcome="ok", output={"content": "opened"})
+        if tool_name == "browser_snapshot" and self.page_open:
+            return ToolCallResult(outcome="ok", output={"content": "snapshot"})
+        return ToolCallResult(outcome="error", output={"message": "No open page"})
+
+    async def close(self) -> None:
+        self.closed = True
+
+
 def _tool(name: str) -> DiscoveredTool:
     return DiscoveredTool(
         name=name,
@@ -86,3 +108,70 @@ def test_generated_workflow_specs_share_injected_tool_executor() -> None:
     assert result["outcome"] == "ok"
     assert result["output"]["content"] == "snapshot"
     assert executor.calls == [("browser_navigate", {}), ("browser_snapshot", {})]
+
+
+def test_runtime_pool_reuses_stateful_session_for_same_connection() -> None:
+    connection = ConnectionConfig(
+        id="playwright.default",
+        server="playwright",
+        account="default",
+        metadata={
+            "transport": "stdio",
+            "command": "pnpx",
+            "args": ["@playwright/mcp"],
+        },
+    )
+    created_clients: list[FakeStatefulClient] = []
+
+    async def factory(
+        connection: ConnectionConfig,
+        auth: AuthRecord | None,
+    ) -> PersistentMcpSession:
+        client = FakeStatefulClient()
+        created_clients.append(client)
+        return PersistentMcpSession(connection=connection, auth=auth, client=client)
+
+    async def run_calls() -> ToolCallResult:
+        pool = McpRuntimePool(factory)
+        await pool.call_tool(connection, None, "browser_navigate", {})
+        return await pool.call_tool(connection, None, "browser_snapshot", {})
+
+    result = asyncio.run(run_calls())
+
+    assert result.outcome == "ok"
+    assert result.output["content"] == "snapshot"
+    assert len(created_clients) == 1
+
+
+def test_runtime_pool_replaces_session_when_fingerprint_changes() -> None:
+    original = ConnectionConfig(
+        id="playwright.default",
+        server="playwright",
+        account="default",
+        metadata={"transport": "stdio", "command": "pnpx", "args": ["old"]},
+    )
+    changed = ConnectionConfig(
+        id="playwright.default",
+        server="playwright",
+        account="default",
+        metadata={"transport": "stdio", "command": "pnpx", "args": ["new"]},
+    )
+    created_clients: list[FakeStatefulClient] = []
+
+    def factory(
+        connection: ConnectionConfig,
+        auth: AuthRecord | None,
+    ) -> PersistentMcpSession:
+        client = FakeStatefulClient()
+        created_clients.append(client)
+        return PersistentMcpSession(connection=connection, auth=auth, client=client)
+
+    async def run_calls() -> None:
+        pool = McpRuntimePool(factory)
+        await pool.call_tool(original, None, "browser_navigate", {})
+        await pool.call_tool(changed, None, "browser_snapshot", {})
+
+    asyncio.run(run_calls())
+
+    assert len(created_clients) == 2
+    assert created_clients[0].closed is True