fix(examples): rewrite mcp-cluster-ops to use console SDK for cluster… (#278)

* fix(examples): rewrite mcp-cluster-ops to use console SDK for cluster routing The example MCP server was broken after the direct HTTP transport refactor — it used TurnstoneServer (single-node) for cluster ops that require TurnstoneConsole (cluster gateway). Rewrites dispatch flow to: route via console → SSE stream from node → cleanup via console. - Switch from TurnstoneServer to TurnstoneConsole for node listing and workstream routing (TURNSTONE_CONSOLE_URL replaces TURNSTONE_SERVER_URL) - Add proper workstream lifecycle: create via routing proxy, stream from node, close in finally block with leak-safe ws_id guard - Catch dispatch exceptions in run_on_node for structured JSON errors - Extract _extract_node_ids helper, remove dead n.get("id") fallback - Normalise _console_kwargs to always include token key - Rewrite tests against Console+Server mocks (36 → 44 tests) * fix(examples): paginate node listing and clarify auth in README Address Copilot review feedback on #278: - _list_nodes_sync now paginates via offset/limit loop so clusters with >100 nodes are fully discovered - README step 2 now mentions token passthrough for authenticated clusters - New test_paginates_large_clusters verifies multi-page fetch (45 tests)
2026-08-12 23:12:23 -06:00 · 2026-04-02 15:12:41 -07:00
parent 3cf9485169
commit 664d44c109
4 changed files with 319 additions and 118 deletions
@@ -1,10 +1,14 @@
 # MCP Cluster Ops

-An MCP server that exposes tools for executing commands across a [Turnstone](https://github.com/turnstonelabs/turnstone) cluster. Serves as a reference implementation for both MCP server patterns and Turnstone SDK usage.
+An MCP server that exposes tools for executing commands across a Turnstone cluster. Serves as a reference implementation for both MCP server patterns and Turnstone SDK usage.

 ## How it works

-This server uses Turnstone's SDK client (`TurnstoneServer`) to dispatch shell commands to specific nodes via HTTP. Remote agents execute the command and the raw bash output is captured directly from the `ToolResultEvent` stream — bypassing the costly "agent reads output → re-generates output as completion tokens" round-trip.
+This server uses the Turnstone console SDK (`TurnstoneConsole`) for node discovery and routing, and `TurnstoneServer` for per-node SSE streaming. The dispatch flow for each command is:
+
+1. **Route** — `TurnstoneConsole.route_create_workstream(target_node=..., auto_approve=True)` creates a workstream pinned to the target node via the console's hash-ring routing proxy, returning `ws_id` and `node_url`.
+2. **Execute** — `TurnstoneServer(node_url, token=...)` connects directly to the node's SSE stream using the same `TURNSTONE_API_TOKEN`. `send_and_wait(prompt, ws_id)` runs the command and the raw bash output is captured from the `ToolResultEvent` — bypassing the costly "agent reads output then re-generates output as completion tokens" round-trip.
+3. **Cleanup** — `TurnstoneConsole.route_close(ws_id)` closes the workstream.

 Multi-node dispatches run in parallel via `asyncio.gather`, so total wall time is bounded by the slowest node rather than the sum.

@@ -19,7 +23,7 @@ Multi-node dispatches run in parallel via `asyncio.gather`, so total wall time i

 ## Prerequisites

- A running Turnstone cluster (at least one `turnstone-server`)
+- A running Turnstone cluster with at least one `turnstone-server` and a `turnstone-console`
 - Python 3.11+

 ## Installation
@@ -35,8 +39,8 @@ pip install -e ./examples/mcp-cluster-ops

 | Variable | Default | Description |
 |----------|---------|-------------|
-| `TURNSTONE_SERVER_URL` | `http://localhost:8080` | Server URL |
-| `TURNSTONE_API_TOKEN` | _(none)_ | API token for authentication |
+| `TURNSTONE_CONSOLE_URL` | `http://localhost:8090` | Console URL for node discovery and routing |
+| `TURNSTONE_API_TOKEN` | _(none)_ | API token / JWT for authentication |
 | `MCP_CLUSTER_OPS_TIMEOUT` | `120` | Default command timeout (seconds, clamped 5-3600) |
 | `MCP_CLUSTER_OPS_MAX_OUTPUT` | `8192` | Max output bytes per node (0 = unlimited) |
 | `MCP_CLUSTER_OPS_MAX_NODES` | `32` | Max concurrent node dispatches |
@@ -51,7 +55,7 @@ pip install -e ./examples/mcp-cluster-ops
 command = "mcp-cluster-ops"

 [mcp.servers.cluster-ops.env]
-TURNSTONE_SERVER_URL = "http://turnstone.example.com:8080"
+TURNSTONE_CONSOLE_URL = "http://console.example.com:8090"
 ```

 **JSON** (via `--mcp-config`):
@@ -62,7 +66,7 @@ TURNSTONE_SERVER_URL = "http://turnstone.example.com:8080"
    "cluster-ops": {
      "command": "mcp-cluster-ops",
      "env": {
-        "TURNSTONE_SERVER_URL": "http://turnstone.example.com:8080"
+        "TURNSTONE_CONSOLE_URL": "http://console.example.com:8090"
      }
    }
  }
@@ -1,7 +1,8 @@
 """MCP server for Turnstone cluster operations.

 Exposes tools to execute commands on specific nodes in a Turnstone cluster.
-Uses the SDK client (``TurnstoneServer``) for direct node targeting via HTTP.
+Uses the SDK console client (``TurnstoneConsole``) for node discovery and
+routing, and ``TurnstoneServer`` for per-node SSE streaming.

 Usage::

@@ -14,12 +15,12 @@ Configure in ``~/.config/turnstone/config.toml``::
    command = "mcp-cluster-ops"

    [mcp.servers.cluster-ops.env]
-    TURNSTONE_SERVER_URL = "http://localhost:8080"
+    TURNSTONE_CONSOLE_URL = "http://localhost:8090"

 Environment variables
 ---------------------
-TURNSTONE_SERVER_URL        Server URL (default: http://localhost:8080)
-TURNSTONE_API_TOKEN         API token for authentication (default: none)
+TURNSTONE_CONSOLE_URL       Console URL (default: http://localhost:8090)
+TURNSTONE_API_TOKEN         API token / JWT for authentication (default: none)
 MCP_CLUSTER_OPS_TIMEOUT     Default command timeout in seconds (default: 120)
 MCP_CLUSTER_OPS_MAX_OUTPUT  Max output bytes per node (default: 8192, 0=unlimited)

@@ -43,7 +44,7 @@ from contextlib import asynccontextmanager
 from typing import TYPE_CHECKING, Any

 from mcp.server.fastmcp import Context, FastMCP
-from turnstone.sdk import TurnResult, TurnstoneServer
+from turnstone.sdk import TurnResult, TurnstoneConsole, TurnstoneServer

 if TYPE_CHECKING:
    from collections.abc import AsyncIterator
@@ -66,15 +67,12 @@ _MAX_TIMEOUT = 3600
 # ---------------------------------------------------------------------------


-def _server_kwargs() -> dict[str, Any]:
-    """Build TurnstoneServer connection kwargs from environment variables."""
-    kwargs: dict[str, Any] = {
-        "base_url": os.environ.get("TURNSTONE_SERVER_URL", "http://localhost:8080"),
+def _console_kwargs() -> dict[str, Any]:
+    """Build TurnstoneConsole connection kwargs from environment variables."""
+    return {
+        "base_url": os.environ.get("TURNSTONE_CONSOLE_URL", "http://localhost:8090"),
+        "token": os.environ.get("TURNSTONE_API_TOKEN", ""),
    }
-    token = os.environ.get("TURNSTONE_API_TOKEN")
-    if token:
-        kwargs["token"] = token
-    return kwargs


 def _exec_prompt(command: str) -> str:
@@ -141,6 +139,11 @@ def _validate_command(command: str) -> str | None:
    return None


+def _extract_node_ids(nodes: list[dict[str, Any]]) -> list[str]:
+    """Extract unique, non-empty node IDs from a list of node dicts."""
+    return list(dict.fromkeys(n["node_id"].strip() for n in nodes if n.get("node_id", "").strip()))
+
+
 def _format_node_result(
    node_id: str,
    result: TurnResult,
@@ -165,12 +168,12 @@ def _format_node_result(


 # ---------------------------------------------------------------------------
-# Core dispatch functions (testable with mocked TurnstoneServer)
+# Core dispatch functions (testable with mocked SDK clients)
 # ---------------------------------------------------------------------------


 def _exec_on_node_sync(
-    server_kw: dict[str, Any],
+    console_kw: dict[str, Any],
    node_id: str,
    command: str,
    timeout: float,
@@ -178,22 +181,35 @@ def _exec_on_node_sync(
    """Dispatch *command* to *node_id* and block until complete.

    Runs inside ``asyncio.to_thread`` so it does not block the event loop.
-    Each call creates its own ``TurnstoneServer`` client to avoid state
-    conflicts between concurrent dispatches.
+
+    Flow:
+      1. Create a workstream on the target node via the console routing proxy
+      2. Connect directly to the node's SSE stream to send + collect output
+      3. Close the workstream via the routing proxy
    """
    prompt = _exec_prompt(command)
-    with TurnstoneServer(**server_kw) as client:
-        result = client.send_and_wait(
-            message=prompt,
-            target_node=node_id,
-            auto_approve=True,
-            timeout=timeout,
-        )
+    ws_id = ""
+    with TurnstoneConsole(**console_kw) as console:
+        try:
+            route_resp = console.route_create_workstream(
+                target_node=node_id,
+                auto_approve=True,
+            )
+            ws_id = route_resp["ws_id"]
+            node_url: str = route_resp["node_url"]
+            with TurnstoneServer(
+                base_url=node_url,
+                token=console_kw["token"],
+            ) as server:
+                result = server.send_and_wait(prompt, ws_id, timeout=timeout)
+        finally:
+            if ws_id:
+                console.route_close(ws_id)
    return node_id, result


 async def _dispatch_parallel(
-    server_kw: dict[str, Any],
+    console_kw: dict[str, Any],
    node_ids: list[str],
    command: str,
    timeout: float,
@@ -204,7 +220,7 @@ async def _dispatch_parallel(
    Total wall time is bounded by the slowest node.
    """
    tasks = [
-        asyncio.to_thread(_exec_on_node_sync, server_kw, nid, command, timeout) for nid in node_ids
+        asyncio.to_thread(_exec_on_node_sync, console_kw, nid, command, timeout) for nid in node_ids
    ]
    outcomes = await asyncio.gather(*tasks, return_exceptions=True)

@@ -220,16 +236,22 @@ async def _dispatch_parallel(
    return results


-def _list_nodes_sync(server_kw: dict[str, Any]) -> list[dict[str, Any]]:
-    """List active cluster nodes (blocking)."""
-    with TurnstoneServer(**server_kw) as client:
-        nodes: list[dict[str, Any]] = client.list_nodes()
-        return nodes
+def _list_nodes_sync(console_kw: dict[str, Any]) -> list[dict[str, Any]]:
+    """List active cluster nodes (blocking), paginating if needed."""
+    page_size = 100
+    nodes: list[dict[str, Any]] = []
+    with TurnstoneConsole(**console_kw) as console:
+        while True:
+            resp = console.nodes(limit=page_size, offset=len(nodes))
+            nodes.extend(n.model_dump() for n in resp.nodes)
+            if len(nodes) >= resp.total or not resp.nodes:
+                break
+    return nodes


-async def _list_nodes_impl(server_kw: dict[str, Any]) -> list[dict[str, Any]]:
+async def _list_nodes_impl(console_kw: dict[str, Any]) -> list[dict[str, Any]]:
    """List active cluster nodes."""
-    return await asyncio.to_thread(_list_nodes_sync, server_kw)
+    return await asyncio.to_thread(_list_nodes_sync, console_kw)


 # ---------------------------------------------------------------------------
@@ -239,9 +261,9 @@ async def _list_nodes_impl(server_kw: dict[str, Any]) -> list[dict[str, Any]]:

@asynccontextmanager
 async def _lifespan(server: FastMCP[dict[str, Any]]) -> AsyncIterator[dict[str, Any]]:
-    """Lifespan context — stores server connection kwargs for tool handlers."""
-    kw = _server_kwargs()
-    yield {"server_kwargs": kw}
+    """Lifespan context — stores console connection kwargs for tool handlers."""
+    kw = _console_kwargs()
+    yield {"console_kwargs": kw}


 mcp = FastMCP(
@@ -263,8 +285,8 @@ async def list_nodes(ctx: Context[Any, Any, Any]) -> str:
    Call this before dispatching work to discover available node IDs.
    Returns a JSON array of node metadata objects.
    """
-    server_kw: dict[str, Any] = ctx.request_context.lifespan_context["server_kwargs"]
-    nodes = await _list_nodes_impl(server_kw)
+    console_kw: dict[str, Any] = ctx.request_context.lifespan_context["console_kwargs"]
+    nodes = await _list_nodes_impl(console_kw)
    return json.dumps(nodes, indent=2)


@@ -291,13 +313,16 @@ async def run_on_node(
    if cmd_err:
        return json.dumps({"error": cmd_err})

-    server_kw: dict[str, Any] = ctx.request_context.lifespan_context["server_kwargs"]
+    console_kw: dict[str, Any] = ctx.request_context.lifespan_context["console_kwargs"]
    max_output = _DEFAULT_MAX_OUTPUT

    log.info("run_on_node node=%s cmd=%r", node_id, command)
-    _, result = await asyncio.to_thread(
-        _exec_on_node_sync, server_kw, node_id, command, _clamp_timeout(timeout)
-    )
+    try:
+        _, result = await asyncio.to_thread(
+            _exec_on_node_sync, console_kw, node_id, command, _clamp_timeout(timeout)
+        )
+    except Exception as exc:
+        return json.dumps({"node": node_id, "ok": False, "error": str(exc)}, indent=2)
    formatted = _format_node_result(node_id, result, max_output)
    return json.dumps(formatted, indent=2)

@@ -323,7 +348,7 @@ async def run_on_nodes(
    if cmd_err:
        return json.dumps({"error": cmd_err})

-    server_kw: dict[str, Any] = ctx.request_context.lifespan_context["server_kwargs"]
+    console_kw: dict[str, Any] = ctx.request_context.lifespan_context["console_kwargs"]
    max_output = _DEFAULT_MAX_OUTPUT

    clean_ids = list(dict.fromkeys(nid.strip() for nid in node_ids if nid.strip()))
@@ -336,7 +361,7 @@ async def run_on_nodes(

    log.info("run_on_nodes nodes=%s cmd=%r", clean_ids, command)
    results = await _dispatch_parallel(
-        server_kw, clean_ids, command, _clamp_timeout(timeout), max_output
+        console_kw, clean_ids, command, _clamp_timeout(timeout), max_output
    )
    return json.dumps(results, indent=2)

@@ -361,18 +386,14 @@ async def run_on_all_nodes(
    if cmd_err:
        return json.dumps({"error": cmd_err})

-    server_kw: dict[str, Any] = ctx.request_context.lifespan_context["server_kwargs"]
+    console_kw: dict[str, Any] = ctx.request_context.lifespan_context["console_kwargs"]
    max_output = _DEFAULT_MAX_OUTPUT

-    nodes = await _list_nodes_impl(server_kw)
+    nodes = await _list_nodes_impl(console_kw)
    if not nodes:
        return json.dumps({"error": "No active nodes found in cluster"})

-    node_ids = list(
-        dict.fromkeys(
-            nid.strip() for n in nodes if (nid := n.get("node_id") or n.get("id")) and nid.strip()
-        )
-    )
+    node_ids = _extract_node_ids(nodes)
    if not node_ids:
        return json.dumps({"error": "No nodes with identifiable IDs found"})
    if len(node_ids) > _MAX_CONCURRENT_NODES:
@@ -381,7 +402,7 @@ async def run_on_all_nodes(
        )
    log.info("run_on_all_nodes nodes=%s cmd=%r", node_ids, command)
    results = await _dispatch_parallel(
-        server_kw, node_ids, command, _clamp_timeout(timeout), max_output
+        console_kw, node_ids, command, _clamp_timeout(timeout), max_output
    )
    return json.dumps(results, indent=2)

@@ -7,6 +7,7 @@ from turnstone.sdk import TurnResult
 from mcp_cluster_ops.server import (
    _clamp_timeout,
    _exec_prompt,
+    _extract_node_ids,
    _extract_output,
    _format_node_result,
    _truncate,
@@ -190,3 +191,53 @@ class TestClampTimeout:

    def test_negative(self):
        assert _clamp_timeout(-1) == 5.0
+
+
+# ---------------------------------------------------------------------------
+# _extract_node_ids
+# ---------------------------------------------------------------------------
+
+
+class TestExtractNodeIds:
+    def test_normal(self):
+        nodes = [
+            {"node_id": "a", "server_url": "http://a:8080"},
+            {"node_id": "b", "server_url": "http://b:8080"},
+        ]
+        assert _extract_node_ids(nodes) == ["a", "b"]
+
+    def test_deduplicates(self):
+        nodes = [
+            {"node_id": "a"},
+            {"node_id": "a"},
+            {"node_id": "b"},
+        ]
+        assert _extract_node_ids(nodes) == ["a", "b"]
+
+    def test_strips_whitespace(self):
+        nodes = [{"node_id": "  a  "}, {"node_id": "b "}]
+        assert _extract_node_ids(nodes) == ["a", "b"]
+
+    def test_skips_empty(self):
+        nodes = [
+            {"node_id": "a"},
+            {"node_id": ""},
+            {"node_id": "  "},
+            {"node_id": "b"},
+        ]
+        assert _extract_node_ids(nodes) == ["a", "b"]
+
+    def test_skips_missing_key(self):
+        nodes = [
+            {"node_id": "a"},
+            {"server_url": "http://orphan:8080"},
+            {"node_id": "b"},
+        ]
+        assert _extract_node_ids(nodes) == ["a", "b"]
+
+    def test_empty_list(self):
+        assert _extract_node_ids([]) == []
+
+    def test_all_empty_ids(self):
+        nodes = [{"node_id": ""}, {"node_id": "  "}]
+        assert _extract_node_ids(nodes) == []
@@ -1,11 +1,13 @@
-"""Tests for MCP tool handlers with mocked TurnstoneServer."""
+"""Tests for MCP tool handlers with mocked SDK clients."""

 from __future__ import annotations

 import asyncio
+import contextlib
 from typing import Any
 from unittest.mock import MagicMock, patch

+import pytest
 from turnstone.sdk import TurnResult

 from mcp_cluster_ops.server import (
@@ -14,6 +16,22 @@ from mcp_cluster_ops.server import (
    _list_nodes_impl,
 )

+_CONSOLE_KW: dict[str, Any] = {"base_url": "http://localhost:8090", "token": ""}
+_CONSOLE_KW_AUTH: dict[str, Any] = {"base_url": "http://localhost:8090", "token": "tok_test"}
+
+
+def _mock_console_ctx(mock_cls: MagicMock, mock_client: MagicMock) -> None:
+    """Wire up a TurnstoneConsole mock as a context manager."""
+    mock_cls.return_value.__enter__ = MagicMock(return_value=mock_client)
+    mock_cls.return_value.__exit__ = MagicMock(return_value=False)
+
+
+def _mock_server_ctx(mock_cls: MagicMock, mock_server: MagicMock) -> None:
+    """Wire up a TurnstoneServer mock as a context manager."""
+    mock_cls.return_value.__enter__ = MagicMock(return_value=mock_server)
+    mock_cls.return_value.__exit__ = MagicMock(return_value=False)
+
+
 # ---------------------------------------------------------------------------
 # _list_nodes_impl
 # ---------------------------------------------------------------------------
@@ -21,26 +39,71 @@ from mcp_cluster_ops.server import (

 class TestListNodesImpl:
    def test_returns_nodes(self):
-        nodes = [{"node_id": "a", "model": "gpt-5"}, {"node_id": "b", "model": "gpt-5"}]
-        with patch("mcp_cluster_ops.server.TurnstoneServer") as mock_cls:
-            mock_client = MagicMock()
-            mock_client.list_nodes.return_value = nodes
-            mock_cls.return_value.__enter__ = MagicMock(return_value=mock_client)
-            mock_cls.return_value.__exit__ = MagicMock(return_value=False)
+        mock_node_a = MagicMock()
+        mock_node_a.model_dump.return_value = {"node_id": "a", "server_url": "http://a:8080"}
+        mock_node_b = MagicMock()
+        mock_node_b.model_dump.return_value = {"node_id": "b", "server_url": "http://b:8080"}

-            result = asyncio.run(_list_nodes_impl({"host": "localhost"}))
-            assert result == nodes
+        mock_resp = MagicMock()
+        mock_resp.nodes = [mock_node_a, mock_node_b]
+        mock_resp.total = 2
+
+        with patch("mcp_cluster_ops.server.TurnstoneConsole") as mock_cls:
+            mock_client = MagicMock()
+            mock_client.nodes.return_value = mock_resp
+            _mock_console_ctx(mock_cls, mock_client)
+
+            result = asyncio.run(_list_nodes_impl(_CONSOLE_KW))
+            assert len(result) == 2
+            assert result[0]["node_id"] == "a"
+            assert result[1]["node_id"] == "b"

    def test_empty_cluster(self):
-        with patch("mcp_cluster_ops.server.TurnstoneServer") as mock_cls:
-            mock_client = MagicMock()
-            mock_client.list_nodes.return_value = []
-            mock_cls.return_value.__enter__ = MagicMock(return_value=mock_client)
-            mock_cls.return_value.__exit__ = MagicMock(return_value=False)
+        mock_resp = MagicMock()
+        mock_resp.nodes = []
+        mock_resp.total = 0

-            result = asyncio.run(_list_nodes_impl({"host": "localhost"}))
+        with patch("mcp_cluster_ops.server.TurnstoneConsole") as mock_cls:
+            mock_client = MagicMock()
+            mock_client.nodes.return_value = mock_resp
+            _mock_console_ctx(mock_cls, mock_client)
+
+            result = asyncio.run(_list_nodes_impl(_CONSOLE_KW))
            assert result == []

+    def test_paginates_large_clusters(self):
+        """Clusters with >100 nodes are fetched across multiple pages."""
+
+        def _make_node(nid: str) -> MagicMock:
+            m = MagicMock()
+            m.model_dump.return_value = {"node_id": nid}
+            return m
+
+        page1_nodes = [_make_node(f"n-{i}") for i in range(100)]
+        page2_nodes = [_make_node(f"n-{i}") for i in range(100, 150)]
+
+        page1_resp = MagicMock()
+        page1_resp.nodes = page1_nodes
+        page1_resp.total = 150
+
+        page2_resp = MagicMock()
+        page2_resp.nodes = page2_nodes
+        page2_resp.total = 150
+
+        with patch("mcp_cluster_ops.server.TurnstoneConsole") as mock_cls:
+            mock_client = MagicMock()
+            mock_client.nodes.side_effect = [page1_resp, page2_resp]
+            _mock_console_ctx(mock_cls, mock_client)
+
+            result = asyncio.run(_list_nodes_impl(_CONSOLE_KW))
+            assert len(result) == 150
+            assert result[0]["node_id"] == "n-0"
+            assert result[149]["node_id"] == "n-149"
+            assert mock_client.nodes.call_count == 2
+            # Verify offset was passed correctly
+            mock_client.nodes.assert_any_call(limit=100, offset=0)
+            mock_client.nodes.assert_any_call(limit=100, offset=100)
+

 # ---------------------------------------------------------------------------
 # _exec_on_node_sync
@@ -50,35 +113,111 @@ class TestListNodesImpl:
 class TestExecOnNodeSync:
    def test_success(self):
        turn_result = TurnResult(
+            ws_id="ws-123",
            tool_results=[("bash", "hello world")],
        )
-        with patch("mcp_cluster_ops.server.TurnstoneServer") as mock_cls:
-            mock_client = MagicMock()
-            mock_client.send_and_wait.return_value = turn_result
-            mock_cls.return_value.__enter__ = MagicMock(return_value=mock_client)
-            mock_cls.return_value.__exit__ = MagicMock(return_value=False)
+        with (
+            patch("mcp_cluster_ops.server.TurnstoneConsole") as mock_console_cls,
+            patch("mcp_cluster_ops.server.TurnstoneServer") as mock_server_cls,
+        ):
+            mock_console = MagicMock()
+            mock_console.route_create_workstream.return_value = {
+                "ws_id": "ws-123",
+                "node_url": "http://node-1:8080",
+                "node_id": "node-1",
+                "name": "ws-123",
+            }
+            _mock_console_ctx(mock_console_cls, mock_console)

-            node_id, result = _exec_on_node_sync(
-                {"host": "localhost"}, "node-1", "echo hello", 60.0
-            )
+            mock_server = MagicMock()
+            mock_server.send_and_wait.return_value = turn_result
+            _mock_server_ctx(mock_server_cls, mock_server)
+
+            node_id, result = _exec_on_node_sync(_CONSOLE_KW_AUTH, "node-1", "echo hello", 60.0)
            assert node_id == "node-1"
            assert result.ok
-            mock_client.send_and_wait.assert_called_once()
-            call_kwargs = mock_client.send_and_wait.call_args
-            assert call_kwargs.kwargs["target_node"] == "node-1"
-            assert call_kwargs.kwargs["auto_approve"] is True
+
+            # Verify console created ws on the right node
+            mock_console.route_create_workstream.assert_called_once_with(
+                target_node="node-1",
+                auto_approve=True,
+            )
+
+            # Verify server connected to the node URL with the token
+            mock_server_cls.assert_called_once_with(
+                base_url="http://node-1:8080",
+                token="tok_test",
+            )
+
+            # Verify send_and_wait got the right ws_id
+            call_kwargs = mock_server.send_and_wait.call_args
+            assert call_kwargs.args[1] == "ws-123"
+
+            # Verify workstream was closed
+            mock_console.route_close.assert_called_once_with("ws-123")

    def test_timeout(self):
-        turn_result = TurnResult(timed_out=True)
-        with patch("mcp_cluster_ops.server.TurnstoneServer") as mock_cls:
-            mock_client = MagicMock()
-            mock_client.send_and_wait.return_value = turn_result
-            mock_cls.return_value.__enter__ = MagicMock(return_value=mock_client)
-            mock_cls.return_value.__exit__ = MagicMock(return_value=False)
+        turn_result = TurnResult(ws_id="ws-456", timed_out=True)
+        with (
+            patch("mcp_cluster_ops.server.TurnstoneConsole") as mock_console_cls,
+            patch("mcp_cluster_ops.server.TurnstoneServer") as mock_server_cls,
+        ):
+            mock_console = MagicMock()
+            mock_console.route_create_workstream.return_value = {
+                "ws_id": "ws-456",
+                "node_url": "http://node-1:8080",
+                "node_id": "node-1",
+            }
+            _mock_console_ctx(mock_console_cls, mock_console)

-            _, result = _exec_on_node_sync({"host": "localhost"}, "node-1", "sleep 9999", 1.0)
+            mock_server = MagicMock()
+            mock_server.send_and_wait.return_value = turn_result
+            _mock_server_ctx(mock_server_cls, mock_server)
+
+            _, result = _exec_on_node_sync(_CONSOLE_KW, "node-1", "sleep 9999", 1.0)
            assert result.timed_out
            assert not result.ok
+            # Workstream still closed even on timeout
+            mock_console.route_close.assert_called_once_with("ws-456")
+
+    def test_send_failure_still_closes_workstream(self):
+        """Workstream must be closed even if send_and_wait raises."""
+        with (
+            patch("mcp_cluster_ops.server.TurnstoneConsole") as mock_console_cls,
+            patch("mcp_cluster_ops.server.TurnstoneServer") as mock_server_cls,
+        ):
+            mock_console = MagicMock()
+            mock_console.route_create_workstream.return_value = {
+                "ws_id": "ws-789",
+                "node_url": "http://node-1:8080",
+                "node_id": "node-1",
+            }
+            _mock_console_ctx(mock_console_cls, mock_console)
+
+            mock_server = MagicMock()
+            mock_server.send_and_wait.side_effect = ConnectionError("lost connection")
+            _mock_server_ctx(mock_server_cls, mock_server)
+
+            with contextlib.suppress(ConnectionError):
+                _exec_on_node_sync(_CONSOLE_KW, "node-1", "echo hi", 60.0)
+
+            mock_console.route_close.assert_called_once_with("ws-789")
+
+    def test_malformed_route_response_no_leak(self):
+        """If route response is missing ws_id, no route_close is attempted."""
+        with patch("mcp_cluster_ops.server.TurnstoneConsole") as mock_console_cls:
+            mock_console = MagicMock()
+            mock_console.route_create_workstream.return_value = {
+                # Missing "ws_id" and "node_url"
+                "node_id": "node-1",
+            }
+            _mock_console_ctx(mock_console_cls, mock_console)
+
+            with pytest.raises(KeyError):
+                _exec_on_node_sync(_CONSOLE_KW, "node-1", "echo hi", 60.0)
+
+            # route_close must NOT be called — ws_id was never assigned
+            mock_console.route_close.assert_not_called()


 # ---------------------------------------------------------------------------
@@ -88,40 +227,32 @@ class TestExecOnNodeSync:

 class TestDispatchParallel:
    def test_parallel_success(self):
-        def fake_exec(server_kw: Any, node_id: str, command: str, timeout: float) -> Any:
-            return (node_id, TurnResult(tool_results=[("bash", f"output-{node_id}")]))
+        def fake_exec(console_kw: Any, node_id: str, command: str, timeout: float) -> Any:
+            return (
+                node_id,
+                TurnResult(tool_results=[("bash", f"output-{node_id}")]),
+            )

        with patch("mcp_cluster_ops.server._exec_on_node_sync", side_effect=fake_exec):
            results = asyncio.run(
-                _dispatch_parallel(
-                    {"host": "localhost"},
-                    ["a", "b", "c"],
-                    "echo hi",
-                    60.0,
-                    8192,
-                )
+                _dispatch_parallel(_CONSOLE_KW, ["a", "b", "c"], "echo hi", 60.0, 8192)
            )
            assert len(results) == 3
            assert all(r["ok"] for r in results)
            outputs = {r["node"]: r["output"] for r in results}
            assert outputs["a"] == "output-a"
            assert outputs["b"] == "output-b"
+            assert outputs["c"] == "output-c"

    def test_partial_failure(self):
-        def fake_exec(server_kw: Any, node_id: str, command: str, timeout: float) -> Any:
+        def fake_exec(console_kw: Any, node_id: str, command: str, timeout: float) -> Any:
            if node_id == "bad":
                raise ConnectionError("connection refused")
            return (node_id, TurnResult(tool_results=[("bash", "ok")]))

        with patch("mcp_cluster_ops.server._exec_on_node_sync", side_effect=fake_exec):
            results = asyncio.run(
-                _dispatch_parallel(
-                    {"host": "localhost"},
-                    ["good", "bad"],
-                    "echo hi",
-                    60.0,
-                    8192,
-                )
+                _dispatch_parallel(_CONSOLE_KW, ["good", "bad"], "echo hi", 60.0, 8192)
            )
            assert len(results) == 2
            good = next(r for r in results if r["node"] == "good")
@@ -131,18 +262,12 @@ class TestDispatchParallel:
            assert "connection refused" in bad["error"]

    def test_all_fail(self):
-        def fake_exec(server_kw: Any, node_id: str, command: str, timeout: float) -> Any:
+        def fake_exec(console_kw: Any, node_id: str, command: str, timeout: float) -> Any:
            raise RuntimeError(f"fail-{node_id}")

        with patch("mcp_cluster_ops.server._exec_on_node_sync", side_effect=fake_exec):
            results = asyncio.run(
-                _dispatch_parallel(
-                    {"host": "localhost"},
-                    ["a", "b"],
-                    "echo hi",
-                    60.0,
-                    8192,
-                )
+                _dispatch_parallel(_CONSOLE_KW, ["a", "b"], "echo hi", 60.0, 8192)
            )
            assert all(not r["ok"] for r in results)
            assert "fail-a" in results[0]["error"]