diff --git a/examples/mcp-cluster-ops/README.md b/examples/mcp-cluster-ops/README.md index 61aa5efd..27d7a542 100644 --- a/examples/mcp-cluster-ops/README.md +++ b/examples/mcp-cluster-ops/README.md @@ -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" } } } diff --git a/examples/mcp-cluster-ops/mcp_cluster_ops/server.py b/examples/mcp-cluster-ops/mcp_cluster_ops/server.py index f971cc13..811357ee 100644 --- a/examples/mcp-cluster-ops/mcp_cluster_ops/server.py +++ b/examples/mcp-cluster-ops/mcp_cluster_ops/server.py @@ -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) diff --git a/examples/mcp-cluster-ops/tests/test_helpers.py b/examples/mcp-cluster-ops/tests/test_helpers.py index 9fee7624..0dd2cfbd 100644 --- a/examples/mcp-cluster-ops/tests/test_helpers.py +++ b/examples/mcp-cluster-ops/tests/test_helpers.py @@ -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) == [] diff --git a/examples/mcp-cluster-ops/tests/test_tools.py b/examples/mcp-cluster-ops/tests/test_tools.py index 8e52bd15..502b2446 100644 --- a/examples/mcp-cluster-ops/tests/test_tools.py +++ b/examples/mcp-cluster-ops/tests/test_tools.py @@ -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"]