From 9a518657a32814e258b433740056ef8ea75fdc34 Mon Sep 17 00:00:00 2001 From: Patrick Buckley Date: Tue, 31 Mar 2026 11:28:39 -0700 Subject: [PATCH] feat: replace console HTTP polling with persistent SSE streams (#266) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: replace console HTTP polling with persistent SSE streams Console collector now subscribes to each server node's /v1/api/events/global SSE stream for real-time state updates instead of polling /v1/api/dashboard and /health every 15 seconds. Server changes: - Emit ws_created/ws_closed events on global queue from create/close handlers - Add node_snapshot on SSE connect (workstreams, health, aggregate) - Add ?expected_node_id= identity verification (409 on mismatch) - Add health_changed callback to BackendHealthMonitor circuit breaker - Add periodic aggregate emitter thread (10s) Console collector changes: - Single asyncio event loop on one thread multiplexes all SSE connections (scales to 1000+ nodes vs thread-per-node) - Discovery loop spawns/cancels async SSE tasks per node - Snapshot reconciliation on connect, delta application for live events - Fix ws_state→cluster_state event type mismatch - Remove polling code (poll_interval, max_poll_workers, --poll-interval CLI) SDK changes: - Add NodeSnapshotEvent, HealthChangedEvent, AggregateEvent dataclasses - Add stream_node_events() method (async + sync) * fix: address review feedback on node event streams - Fix stop() to let SSE manager exit naturally instead of force-stopping the event loop (ensures finally cleanup runs) - Guard against empty/invalid SSE data from ping frames - Treat missing node_id as identity mismatch (409) when expected_node_id is provided - Fix stale docstring on _update_metrics --- docs/architecture.md | 40 +- docs/console.md | 15 +- docs/diagrams/11-console-data-flow.puml | 89 ++-- docs/diagrams/png/11-console-data-flow.png | 4 +- docs/docker.md | 1 - tests/test_console.py | 297 +++++++---- tests/test_tls_client.py | 4 +- turnstone/api/server_spec.py | 7 +- turnstone/console/collector.py | 554 +++++++++++++-------- turnstone/console/server.py | 7 - turnstone/core/config.py | 1 - turnstone/core/healthcheck.py | 24 +- turnstone/sdk/events.py | 33 ++ turnstone/sdk/server.py | 22 + turnstone/server.py | 207 +++++++- 15 files changed, 913 insertions(+), 392 deletions(-) diff --git a/docs/architecture.md b/docs/architecture.md index fec087b7..66990f65 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -74,7 +74,7 @@ turnstone/ _sync.py Background event loop for sync wrappers _types.py TurnResult + TurnstoneAPIError console/ - collector.py ClusterCollector — aggregates state from all nodes via HTTP + collector.py ClusterCollector — aggregates state from all nodes via SSE scheduler.py TaskScheduler — background cron/at scheduler, dispatches via HTTP server.py Cluster dashboard HTTP server + SSE + CLI entry point static/ Cluster dashboard web UI (page-specific HTML, CSS, JS) @@ -1201,31 +1201,31 @@ bell + status line to stderr to alert the user. ### Cluster Console ``` -Monitoring (3 daemon threads) Control + Proxy (async Starlette) +Monitoring (2 daemon threads) Control + Proxy (async Starlette) +------------------+ +----------------------------+ -| Event subscriber | | POST /v1/api/cluster/ | -| SSE on | | workstreams/new | -| /events/glob | | → POST to target server | +| Node discovery | | POST /v1/api/cluster/ | +| Service registry | | workstreams/new | +| every 60 seconds | | → POST to target server | +------------------+ +----------------------------+ -| Node discovery | | GET /node/{node_id}/ | -| Service registry | | → httpx.AsyncClient | -| every 15 seconds | | proxy to server_url | -+------------------+ | GET /node/{id}/v1/api/events | -| Poll loop | | → SSE stream proxy | -| GET /v1/api/dash | | POST /node/{id}/v1/api/send | -| GET /health | | → forwarded to server | -| ThreadPoolExec | +----------------------------+ -+------------------+ +| SSE manager | | GET /node/{node_id}/ | +| asyncio loop | | → httpx.AsyncClient | +| 1 task per node | | proxy to server_url | +| /events/global | | GET /node/{id}/v1/api/events | +| snapshot+deltas | | → SSE stream proxy | ++------------------+ | POST /node/{id}/v1/api/send | + | → forwarded to server | + +----------------------------+ ``` The console HTTP layer is a Starlette/ASGI app served by uvicorn. The SSE endpoint uses `EventSourceResponse` with the same listener queue pattern as -the main server. `ClusterCollector`'s background threads (event subscriber, -node discovery, poll loop) use `ThreadPoolExecutor` -for parallel HTTP polling. The poll loop diffs workstream IDs between poll -cycles and fans out synthetic `ws_created`/`ws_closed` SSE events for any -changes, ensuring browser clients stay in sync even when real-time cluster -events are missed. +the main server. `ClusterCollector` runs two daemon threads: a discovery loop +that queries the service registry every 60 seconds, and an SSE manager that +runs a single asyncio event loop multiplexing persistent SSE connections to +all nodes via `GET /v1/api/events/global`. Each node delivers a full snapshot +on connect followed by real-time delta events — state changes, health +transitions, and aggregate metrics arrive sub-second instead of on a 15-second +poll cycle. The console has two write-path capabilities: diff --git a/docs/console.md b/docs/console.md index 67677108..ea4207d7 100644 --- a/docs/console.md +++ b/docs/console.md @@ -1,6 +1,6 @@ # Cluster Dashboard (turnstone-console) -`turnstone-console` is a cluster management service that provides cluster-wide visibility and control across all turnstone nodes. It discovers nodes via the `services` database table, polls each node's HTTP API for workstream data, and receives real-time state changes via HTTP polling. +`turnstone-console` is a cluster management service that provides cluster-wide visibility and control across all turnstone nodes. It discovers nodes via the `services` database table and subscribes to each node's SSE event stream for real-time workstream, health, and metric updates. The console also supports **workstream creation** (dispatched via HTTP proxy to target nodes) and a **reverse proxy** that serves each node's server UI through the console port — so users only need network access to the console, not to individual server nodes. @@ -21,7 +21,7 @@ turnstone-console ──────┤ Data flows in two directions: -- **Inbound (monitoring):** The console discovers nodes via the `services` database table (nodes register on startup and send periodic heartbeats). It periodically polls each node's `GET /v1/api/dashboard` for full workstream snapshots and `GET /health` for node health. +- **Inbound (monitoring):** The console discovers nodes via the `services` database table (nodes register on startup and send periodic heartbeats). It opens a persistent SSE connection to each node's `GET /v1/api/events/global` endpoint, receiving a full snapshot on connect followed by real-time delta events (state changes, health transitions, aggregate metrics). - **Outbound (control):** The console proxies workstream creation requests to target nodes via HTTP. - **Proxy (pass-through):** The console reverse-proxies each node's server UI at `/node/{node_id}/`, forwarding HTTP and SSE traffic so the browser never contacts server nodes directly. @@ -30,8 +30,7 @@ Data flows in two directions: | Source | Method | Direction | Data | |--------|--------|-----------|------| | `services` table | Database query | Read | Node discovery (node_id, server_url, started) | -| Node HTTP API | `GET {server_url}/v1/api/dashboard` | Read | Full workstream list with tokens, context, activity | -| Node HTTP API | `GET {server_url}/health` | Read | Node health status | +| Node SSE | `GET {server_url}/v1/api/events/global` | Stream | Snapshot on connect, then real-time delta events (state, health, aggregate) | | Node HTTP API | `POST {server_url}/v1/api/workstreams/new` | Write | Workstream creation | | Node HTTP API | `GET/POST {server_url}/*` | Proxy | Server UI, API requests, SSE streams | @@ -41,9 +40,9 @@ Data flows in two directions: The collector (`turnstone/console/collector.py`) maintains an in-memory snapshot of all nodes and workstreams. Two daemon threads handle data acquisition: -1. **Node discovery** — queries the `services` database table every 15 seconds. Adds newly discovered nodes, removes expired ones (stale heartbeats), emits `node_joined` / `node_lost` events to SSE listeners. +1. **Node discovery** — queries the `services` database table every 60 seconds. Adds newly discovered nodes, removes expired ones (stale heartbeats), emits `node_joined` / `node_lost` events to SSE listeners, and spawns/cancels SSE tasks for new/lost nodes. -2. **Poll loop** — fetches `GET /v1/api/dashboard` and `GET /health` from each known node every 10 seconds. Uses `ThreadPoolExecutor(max_workers=50)` for parallelism. Each poll replaces the node's workstream list with the authoritative server data. +2. **SSE manager** — a single asyncio event loop on one thread multiplexes persistent SSE connections to all discovered nodes via `GET /v1/api/events/global`. Each connection receives a `node_snapshot` on connect (workstreams, health, aggregate) followed by real-time delta events (`ws_state`, `ws_created`, `ws_closed`, `ws_rename`, `health_changed`, `aggregate`). On disconnect, the node is marked unreachable and the connection is retried with exponential backoff (1s–30s). An `?expected_node_id=` query parameter provides identity verification against IP reuse (server returns 409 on mismatch). A `get_snapshot()` method builds the full cluster state under a single lock acquisition — overview aggregates and per-node workstream lists in one atomic read. This is served both as a REST endpoint and as the initial SSE event on client connect. @@ -54,7 +53,7 @@ All reads and writes to the node/workstream map are protected by a single `threa ### Scale Considerations - **50,000 workstreams** (1,000 nodes × 50 per node) at ~500 bytes each = ~25 MB in memory -- **1,000 nodes** polled in parallel — fan-out concurrency is configurable via `cluster.node_fan_out_limit` (default 200), yielding 5 batches at ~100ms each = ~0.5 second poll cycle +- **1,000 nodes** connected via persistent SSE — a single asyncio event loop multiplexes all connections with negligible overhead. Ensure `ulimit -n` >= 4096 for fd headroom - **Filtering and pagination** run in-memory on the full workstream list — sub-millisecond at this scale - **SSE fan-out** uses per-client queues (2,000 events) — backed-up clients get events dropped, not blocking - **Database** — for clusters sharing PostgreSQL, use [PgBouncer](pgbouncer.md) in transaction pooling mode @@ -629,7 +628,6 @@ CLI flags for `turnstone-console`: |------|---------|-------------| | `--host` | `0.0.0.0` | Bind host | | `--port` | `8090` | HTTP port | -| `--poll-interval` | `10` | Node polling interval (seconds) | | `--auth-token` | `$TURNSTONE_AUTH_TOKEN` | Bearer token for server node communication and proxy | | `--log-level` | `INFO` | Log level | @@ -640,7 +638,6 @@ Config file (`~/.config/turnstone/config.toml`): host = "0.0.0.0" port = 8090 url = "http://localhost:8090" # used by CLI /cluster commands -poll_interval = 10 ``` --- diff --git a/docs/diagrams/11-console-data-flow.puml b/docs/diagrams/11-console-data-flow.puml index 2c234158..f44fe918 100644 --- a/docs/diagrams/11-console-data-flow.puml +++ b/docs/diagrams/11-console-data-flow.puml @@ -10,55 +10,76 @@ participant "ClusterCollector" as CC participant "Node-A\n(server)" as NodeA participant "Node-B\n(server)" as NodeB -== Thread 1: HTTP Polling (every 10s) == +== Thread 1: Node Discovery (every 60s) == -CC -> CC : Iterate registered nodes +CC -> CC : list_services("server",\nmax_age_seconds=120) activate CC #C8E6C9 -CC -> NodeA : GET /v1/api/dashboard -activate NodeA -NodeA --> CC : {workstreams: [...],\naggregate: {total_tokens, ...}} -deactivate NodeA - -CC -> NodeA : GET /health -activate NodeA -NodeA --> CC : {status:"ok", version:"0.9.2",\nmodel:"...", workstreams:{...}} -deactivate NodeA - -CC -> NodeB : GET /v1/api/dashboard -activate NodeB -NodeB --> CC : {workstreams: [...],\naggregate: {total_tokens, ...}} -deactivate NodeB - -CC -> CC : Diff old vs new workstream IDs -CC -> CC : Replace NodeSnapshot["nodeA"]\n.workstreams, .health, .aggregate -CC -> CC : _fanout(ws_created) for\nnewly appeared workstreams -CC -> CC : _fanout(ws_closed) for\nremoved workstreams - -note right of CC - Poll-diff fanout ensures - browser SSE clients learn - about workstreams that - appeared between polls. -end note +CC -> CC : New node? → spawn SSE task\nLost node? → cancel SSE task +CC -> CC : _fanout(node_joined)\n_fanout(node_lost) deactivate CC +== Thread 2: SSE Manager (asyncio event loop) == + +note over CC + Single asyncio event loop multiplexes + one persistent SSE connection per node. + Scales to 1000+ nodes. +end note + +CC -> NodeA : GET /v1/api/events/global\n?expected_node_id=nodeA +activate NodeA +activate CC #BBDEFB + +NodeA --> CC : data: {"type":"node_snapshot",\n"node_id":"nodeA",\n"workstreams":[...],\n"health":{...},\n"aggregate":{...}} + +note right of CC + Snapshot populates NodeSnapshot + in-memory state. Reconciles + against stale data (emits + ws_created/ws_closed diffs). +end note + +loop real-time delta events + NodeA --> CC : data: {"type":"ws_state",\n"ws_id":"ws1","state":"running"} + CC -> CC : Update NodeSnapshot\n_fanout(cluster_state) +end + +alt health transition + NodeA --> CC : data: {"type":"health_changed",\n"circuit_state":"open"} + CC -> CC : Update node.health +end + +alt periodic aggregate (every 10s) + NodeA --> CC : data: {"type":"aggregate",\n"total_tokens":50000} + CC -> CC : Update node.aggregate +end + +deactivate CC +deactivate NodeA + +alt SSE disconnect + CC -> CC : Mark node unreachable\nReconnect with backoff\n(1s → 30s cap) +end + +alt identity mismatch (409 or snapshot node_id differs) + CC -> CC : Mark node unreachable\nStop reconnecting to this URL +end + == Browser SSE Stream == Browser -> Server : GET /v1/api/cluster/events activate Server -Server -> CC : get_snapshot() +Server -> CC : get_snapshot_and_register(queue) +note right : Atomic: snapshot + listener\nregistration under both locks\n→ no event gap CC --> Server : ClusterSnapshot\n(full current state) -Server -> CC : register_listener(queue) -note right : Per-client queue.Queue(maxsize=2000)\nSSE via EventSourceResponse + run_in_executor() - Server -> Browser : data: {"type":"snapshot",...}\n(full state as first SSE event) loop continuous (incremental updates) - CC -> Server : event via listener queue\n(from polling thread) + CC -> Server : event via listener queue\n(from SSE manager thread) Server -> Browser : data: {"type":"cluster_state",...}\n\n end @@ -81,7 +102,7 @@ Server --> Browser : JSON response Browser -> Server : GET /v1/api/cluster/overview Server -> CC : get_overview() -CC --> Server : {nodes: 2, workstreams: 12,\nstates: {running:5, ...},\naggregate: {total_tokens: 50000},\nversion_drift: false, versions: ["0.9.2"]} +CC --> Server : {nodes: 2, workstreams: 12,\nstates: {running:5, ...},\naggregate: {total_tokens: 50000},\nversion_drift: false, versions: ["0.9.7"]} Server --> Browser : JSON response Browser -> Server : GET /v1/api/cluster/nodes?sort=activity diff --git a/docs/diagrams/png/11-console-data-flow.png b/docs/diagrams/png/11-console-data-flow.png index 2098d889..f4e18c15 100644 --- a/docs/diagrams/png/11-console-data-flow.png +++ b/docs/diagrams/png/11-console-data-flow.png @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:14824b81fa87f9a29e9b182b54132d3e438dd83f880b20b111b4bf51bc4b39d1 -size 317947 +oid sha256:040f7d9ec7d676da40b9487e0825caf2c1574cbdd9f16d0998d90e0c2e4f8861 +size 360309 diff --git a/docs/docker.md b/docs/docker.md index abafb70e..5fab818b 100644 --- a/docs/docker.md +++ b/docs/docker.md @@ -69,7 +69,6 @@ All configuration is via environment variables in `.env` (copy from `.env.exampl | Variable | Default | Description | |----------|---------|-------------| | `CONSOLE_PORT` | `8090` | Host port mapping | -| `CONSOLE_POLL_INTERVAL` | `10` | Node polling interval (seconds) | ### Auth diff --git a/tests/test_console.py b/tests/test_console.py index d319666b..1d1effbc 100644 --- a/tests/test_console.py +++ b/tests/test_console.py @@ -3,7 +3,7 @@ import asyncio import json import queue -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock import pytest @@ -29,32 +29,15 @@ class MockStorage: # --------------------------------------------------------------------------- -def _make_collector(storage=None, poll_interval=0, discovery_interval=999): - """Create a collector with zero poll interval (no jitter delay in tests).""" +def _make_collector(storage=None, discovery_interval=999): + """Create a collector for tests (discovery disabled by default).""" s = storage or MockStorage() return ClusterCollector( storage=s, - poll_interval=poll_interval, discovery_interval=discovery_interval, ) -def _dashboard_response(workstreams=None, aggregate=None): - """Build a /v1/api/dashboard-style response dict.""" - return { - "workstreams": workstreams or [], - "aggregate": aggregate - or { - "total_tokens": 0, - "total_tool_calls": 0, - "active_count": 0, - "total_count": 0, - "uptime_seconds": 0, - "node": "local", - }, - } - - # --------------------------------------------------------------------------- # ClusterCollector — unit tests # --------------------------------------------------------------------------- @@ -157,30 +140,35 @@ class TestCollectorDiscovery: assert c._nodes["node-a"].started == 1234567890.0 -class TestCollectorPolling: - """Polling /v1/api/dashboard from nodes.""" +class TestCollectorSnapshot: + """Applying node_snapshot SSE events.""" - def test_apply_poll_populates_workstreams(self): + def test_apply_snapshot_populates_workstreams(self): c = _make_collector() c._nodes["node-a"] = NodeSnapshot(node_id="node-a", server_url="http://a:8080") - dashboard = _dashboard_response( - workstreams=[ - { - "id": "ws1", - "name": "test", - "state": "running", - "tokens": 1000, - "context_ratio": 0.15, - "activity": "bash: ls", - "activity_state": "tool", - "tool_calls": 3, - "title": "My task", - }, - ], - aggregate={"total_tokens": 1000, "total_tool_calls": 3}, + c._apply_snapshot( + "node-a", + { + "type": "node_snapshot", + "node_id": "node-a", + "workstreams": [ + { + "id": "ws1", + "name": "test", + "state": "running", + "tokens": 1000, + "context_ratio": 0.15, + "activity": "bash: ls", + "activity_state": "tool", + "tool_calls": 3, + "title": "My task", + }, + ], + "health": {"status": "ok"}, + "aggregate": {"total_tokens": 1000, "total_tool_calls": 3}, + }, ) - c._apply_poll("node-a", dashboard, {"status": "ok"}) detail = c.get_node_detail("node-a") assert len(detail["workstreams"]) == 1 @@ -189,7 +177,7 @@ class TestCollectorPolling: assert detail["workstreams"][0]["server_url"] == "http://a:8080" assert detail["health"]["status"] == "ok" - def test_apply_poll_replaces_stale_workstreams(self): + def test_apply_snapshot_replaces_stale_workstreams(self): c = _make_collector() c._nodes["node-a"] = NodeSnapshot( node_id="node-a", @@ -197,30 +185,44 @@ class TestCollectorPolling: workstreams={"old-ws": {"id": "old-ws", "name": "old", "state": "idle"}}, ) - dashboard = _dashboard_response( - workstreams=[{"id": "new-ws", "name": "new", "state": "running"}] + c._apply_snapshot( + "node-a", + { + "type": "node_snapshot", + "node_id": "node-a", + "workstreams": [{"id": "new-ws", "name": "new", "state": "running"}], + "health": {}, + "aggregate": {}, + }, ) - c._apply_poll("node-a", dashboard, {}) detail = c.get_node_detail("node-a") assert len(detail["workstreams"]) == 1 assert detail["workstreams"][0]["id"] == "new-ws" - def test_apply_poll_ignores_unknown_node(self): + def test_apply_snapshot_ignores_unknown_node(self): c = _make_collector() # Should not raise - c._apply_poll("unknown", _dashboard_response(), {}) + c._apply_snapshot( + "unknown", {"type": "node_snapshot", "workstreams": [], "health": {}, "aggregate": {}} + ) - def test_apply_poll_emits_ws_created_for_new_workstream(self): + def test_apply_snapshot_emits_ws_created_for_new_workstream(self): c = _make_collector() c._nodes["node-a"] = NodeSnapshot(node_id="node-a", server_url="http://a:8080") q: queue.Queue[dict] = queue.Queue() c.register_listener(q) - dashboard = _dashboard_response( - workstreams=[{"id": "ws1", "name": "new-task", "state": "idle"}] + c._apply_snapshot( + "node-a", + { + "type": "node_snapshot", + "node_id": "node-a", + "workstreams": [{"id": "ws1", "name": "new-task", "state": "idle"}], + "health": {}, + "aggregate": {}, + }, ) - c._apply_poll("node-a", dashboard, {}) event = q.get_nowait() assert event["type"] == "ws_created" @@ -228,7 +230,7 @@ class TestCollectorPolling: assert event["name"] == "new-task" assert event["node_id"] == "node-a" - def test_apply_poll_emits_ws_closed_for_removed_workstream(self): + def test_apply_snapshot_emits_ws_closed_for_removed_workstream(self): c = _make_collector() c._nodes["node-a"] = NodeSnapshot( node_id="node-a", @@ -238,13 +240,22 @@ class TestCollectorPolling: q: queue.Queue[dict] = queue.Queue() c.register_listener(q) - c._apply_poll("node-a", _dashboard_response(), {}) + c._apply_snapshot( + "node-a", + { + "type": "node_snapshot", + "node_id": "node-a", + "workstreams": [], + "health": {}, + "aggregate": {}, + }, + ) event = q.get_nowait() assert event["type"] == "ws_closed" assert event["ws_id"] == "ws1" - def test_apply_poll_no_events_when_unchanged(self): + def test_apply_snapshot_no_events_when_unchanged(self): c = _make_collector() c._nodes["node-a"] = NodeSnapshot( node_id="node-a", @@ -254,15 +265,21 @@ class TestCollectorPolling: q: queue.Queue[dict] = queue.Queue() c.register_listener(q) - # Same state — no events expected - dashboard = _dashboard_response( - workstreams=[{"id": "ws1", "name": "same", "state": "idle"}] + c._apply_snapshot( + "node-a", + { + "type": "node_snapshot", + "node_id": "node-a", + "workstreams": [{"id": "ws1", "name": "same", "state": "idle"}], + "health": {}, + "aggregate": {}, + }, ) - c._apply_poll("node-a", dashboard, {}) assert q.empty() - def test_apply_poll_emits_state_change(self): + def test_apply_snapshot_emits_state_change_as_cluster_state(self): + """State change events must use type 'cluster_state' for the frontend.""" c = _make_collector() c._nodes["node-a"] = NodeSnapshot( node_id="node-a", @@ -272,30 +289,141 @@ class TestCollectorPolling: q: queue.Queue[dict] = queue.Queue() c.register_listener(q) - dashboard = _dashboard_response( - workstreams=[{"id": "ws1", "name": "same", "state": "running"}] + c._apply_snapshot( + "node-a", + { + "type": "node_snapshot", + "node_id": "node-a", + "workstreams": [{"id": "ws1", "name": "same", "state": "running"}], + "health": {}, + "aggregate": {}, + }, ) - c._apply_poll("node-a", dashboard, {}) event = q.get_nowait() - assert event["type"] == "ws_state" + assert event["type"] == "cluster_state" assert event["ws_id"] == "ws1" assert event["state"] == "running" - def test_apply_poll_skips_empty_id_workstream(self): + def test_apply_snapshot_skips_empty_id_workstream(self): c = _make_collector() c._nodes["node-a"] = NodeSnapshot(node_id="node-a", server_url="http://a:8080") q: queue.Queue[dict] = queue.Queue() c.register_listener(q) - dashboard = _dashboard_response(workstreams=[{"name": "no-id", "state": "idle"}]) - c._apply_poll("node-a", dashboard, {}) + c._apply_snapshot( + "node-a", + { + "type": "node_snapshot", + "node_id": "node-a", + "workstreams": [{"name": "no-id", "state": "idle"}], + "health": {}, + "aggregate": {}, + }, + ) assert q.empty() assert len(c._nodes["node-a"].workstreams) == 0 - def test_poll_401_preserves_workstreams_and_marks_unreachable(self): - """A 401 from the server must NOT wipe workstream data.""" + +class TestCollectorDelta: + """Applying individual SSE delta events.""" + + def test_apply_delta_ws_state_fans_out_as_cluster_state(self): + """Server emits ws_state; collector must translate to cluster_state.""" + c = _make_collector() + c._nodes["node-a"] = NodeSnapshot( + node_id="node-a", + server_url="http://a:8080", + workstreams={"ws1": {"id": "ws1", "name": "test", "state": "idle"}}, + ) + q: queue.Queue[dict] = queue.Queue() + c.register_listener(q) + + c._apply_delta( + "node-a", {"type": "ws_state", "ws_id": "ws1", "state": "running", "tokens": 500} + ) + + event = q.get_nowait() + assert event["type"] == "cluster_state" + assert event["state"] == "running" + # Verify in-memory state was updated + assert c._nodes["node-a"].workstreams["ws1"]["state"] == "running" + + def test_apply_delta_ws_created(self): + c = _make_collector() + c._nodes["node-a"] = NodeSnapshot(node_id="node-a", server_url="http://a:8080") + q: queue.Queue[dict] = queue.Queue() + c.register_listener(q) + + c._apply_delta("node-a", {"type": "ws_created", "ws_id": "ws1", "name": "new"}) + + event = q.get_nowait() + assert event["type"] == "ws_created" + assert "ws1" in c._nodes["node-a"].workstreams + + def test_apply_delta_ws_closed(self): + c = _make_collector() + c._nodes["node-a"] = NodeSnapshot( + node_id="node-a", + server_url="http://a:8080", + workstreams={"ws1": {"id": "ws1", "name": "old", "state": "idle"}}, + ) + q: queue.Queue[dict] = queue.Queue() + c.register_listener(q) + + c._apply_delta("node-a", {"type": "ws_closed", "ws_id": "ws1"}) + + event = q.get_nowait() + assert event["type"] == "ws_closed" + assert "ws1" not in c._nodes["node-a"].workstreams + + def test_apply_delta_ws_rename(self): + c = _make_collector() + c._nodes["node-a"] = NodeSnapshot( + node_id="node-a", + server_url="http://a:8080", + workstreams={"ws1": {"id": "ws1", "name": "old-name", "state": "idle"}}, + ) + q: queue.Queue[dict] = queue.Queue() + c.register_listener(q) + + c._apply_delta("node-a", {"type": "ws_rename", "ws_id": "ws1", "name": "new-name"}) + + event = q.get_nowait() + assert event["type"] == "ws_rename" + assert event["name"] == "new-name" + assert c._nodes["node-a"].workstreams["ws1"]["name"] == "new-name" + + def test_apply_delta_health_changed(self): + c = _make_collector() + c._nodes["node-a"] = NodeSnapshot( + node_id="node-a", + server_url="http://a:8080", + health={"status": "ok", "backend": {"status": "up", "circuit_state": "closed"}}, + ) + + c._apply_delta("node-a", {"type": "health_changed", "circuit_state": "open"}) + + health = c._nodes["node-a"].health + assert health["backend"]["circuit_state"] == "open" + assert health["backend"]["status"] == "down" + assert health["status"] == "degraded" + + def test_apply_delta_aggregate(self): + c = _make_collector() + c._nodes["node-a"] = NodeSnapshot(node_id="node-a", server_url="http://a:8080") + + c._apply_delta( + "node-a", + {"type": "aggregate", "total_tokens": 5000, "total_tool_calls": 42, "active_count": 3}, + ) + + assert c._nodes["node-a"].aggregate["total_tokens"] == 5000 + assert c._nodes["node-a"].aggregate["total_tool_calls"] == 42 + + def test_mark_unreachable_preserves_workstreams(self): + """Disconnection marks unreachable but preserves workstream data.""" c = _make_collector() c._nodes["node-a"] = NodeSnapshot( node_id="node-a", @@ -304,47 +432,12 @@ class TestCollectorPolling: workstreams={"ws1": {"id": "ws1", "name": "existing", "state": "idle"}}, ) - # Mock httpx to return 401 - import httpx as _httpx + c._mark_unreachable("node-a") - mock_response = _httpx.Response( - 401, - json={"error": "Unauthorized"}, - request=_httpx.Request("GET", "http://a:8080/v1/api/dashboard"), - ) - - with patch.object(c._http_client, "get", return_value=mock_response): - c._poll_all_nodes() - - # Workstream data must be preserved, node marked unreachable assert c._nodes["node-a"].reachable is False assert "ws1" in c._nodes["node-a"].workstreams assert c._nodes["node-a"].workstreams["ws1"]["name"] == "existing" - def test_poll_403_preserves_workstreams(self): - """A 403 should also preserve state and mark unreachable.""" - c = _make_collector() - c._nodes["node-a"] = NodeSnapshot( - node_id="node-a", - server_url="http://a:8080", - reachable=True, - workstreams={"ws1": {"id": "ws1", "name": "keep-me", "state": "running"}}, - ) - - import httpx as _httpx - - mock_response = _httpx.Response( - 403, - json={"error": "Forbidden"}, - request=_httpx.Request("GET", "http://a:8080/v1/api/dashboard"), - ) - - with patch.object(c._http_client, "get", return_value=mock_response): - c._poll_all_nodes() - - assert c._nodes["node-a"].reachable is False - assert "ws1" in c._nodes["node-a"].workstreams - class TestCollectorFanout: """SSE fan-out to registered listeners.""" diff --git a/tests/test_tls_client.py b/tests/test_tls_client.py index 827abdaa..a570bf70 100644 --- a/tests/test_tls_client.py +++ b/tests/test_tls_client.py @@ -84,5 +84,5 @@ def test_collector_tls_defaults(): storage_mock = MagicMock() collector = ClusterCollector(storage=storage_mock) - # Should create httpx client without errors - assert collector._http_client is not None + # Should store TLS settings for async client creation + assert collector._tls_verify is True diff --git a/turnstone/api/server_spec.py b/turnstone/api/server_spec.py index 16980ef3..0279facb 100644 --- a/turnstone/api/server_spec.py +++ b/turnstone/api/server_spec.py @@ -137,8 +137,11 @@ SERVER_ENDPOINTS: list[EndpointSpec] = [ "/v1/api/events/global", "GET", "Global SSE event stream", - description="Global Server-Sent Events stream for state-change broadcasts " - "across all workstreams. Returns text/event-stream.", + description="Server-Sent Events stream for node-level state broadcasts. " + "Emits a node_snapshot event on connect (workstreams, health, aggregate), " + "followed by real-time delta events (ws_state, ws_activity, ws_created, " + "ws_closed, ws_rename, health_changed, aggregate). " + "Pass ?expected_node_id=X for identity verification (returns 409 on mismatch).", tags=["Streaming"], ), # --- Saved workstreams --- diff --git a/turnstone/console/collector.py b/turnstone/console/collector.py index 6cac6ba7..dd3cf928 100644 --- a/turnstone/console/collector.py +++ b/turnstone/console/collector.py @@ -1,22 +1,26 @@ """Cluster state collector — aggregates data from all turnstone nodes. -Discovers nodes via the service registry (StorageBackend), polls each -node's /v1/api/dashboard endpoint for workstream data. +Discovers nodes via the service registry (StorageBackend) and subscribes +to each node's ``/v1/api/events/global`` SSE stream for real-time state +updates. A single asyncio event loop on one dedicated thread multiplexes +all SSE connections, scaling to 1000+ nodes. """ from __future__ import annotations +import asyncio import contextlib import json import logging import queue +import random import threading import time -from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any import httpx +import httpx_sse if TYPE_CHECKING: from turnstone.console.metrics import ConsoleMetrics @@ -34,7 +38,7 @@ class NodeSnapshot: node_id: str = "" server_url: str = "" started: float = 0.0 - last_seen: float = 0.0 # monotonic time of last successful poll + last_seen: float = 0.0 # monotonic time of last successful data max_ws: int = 10 # max workstreams (capacity) workstreams: dict[str, dict[str, Any]] = field(default_factory=dict) health: dict[str, Any] = field(default_factory=dict) @@ -43,19 +47,17 @@ class NodeSnapshot: class ClusterCollector: - """Aggregates cluster state from the service registry and per-node HTTP APIs. + """Aggregates cluster state from the service registry and per-node SSE streams. Two daemon threads: 1. Node discovery — queries the service registry every ``discovery_interval`` seconds - 2. Poll loop — fetches /v1/api/dashboard from each node every ``poll_interval`` seconds + 2. SSE manager — single asyncio event loop multiplexing SSE connections to all nodes """ def __init__( self, storage: StorageBackend, - poll_interval: float = 15.0, - discovery_interval: float = 15.0, - max_poll_workers: int = 200, + discovery_interval: float = 60.0, http_timeout: float = 30.0, auth_token: str = "", token_manager: ServiceTokenManager | None = None, @@ -65,16 +67,14 @@ class ClusterCollector: console_metrics: ConsoleMetrics | None = None, ): self._storage = storage - self._poll_interval = poll_interval self._discovery_interval = discovery_interval - self._max_poll_workers = max_poll_workers self._http_timeout = http_timeout self._token_manager = token_manager self._router = router self._console_metrics = console_metrics + self._tls_verify = tls_verify + self._tls_cert = tls_cert # Static auth header — only used when no token_manager is present. - # When a token_manager exists, auth is injected per-request via - # extra_headers in _poll_all_nodes to avoid stale JWT expiry. self._static_auth: dict[str, str] | None = None if auth_token and token_manager is None: self._static_auth = {"Authorization": f"Bearer {auth_token}"} @@ -83,39 +83,41 @@ class ClusterCollector: self._nodes: dict[str, NodeSnapshot] = {} self._running = False self._threads: list[threading.Thread] = [] - self._poll_pool = ThreadPoolExecutor(max_workers=max_poll_workers) - self._http_client = httpx.Client( - timeout=httpx.Timeout(connect=10, read=http_timeout, write=5, pool=http_timeout), - limits=httpx.Limits( - max_connections=max_poll_workers + 10, - max_keepalive_connections=min(max_poll_workers, 200), - ), - verify=tls_verify, - cert=tls_cert, - ) # SSE fan-out to browser clients self._listeners: list[queue.Queue[dict[str, Any]]] = [] self._listeners_lock = threading.Lock() + # SSE manager state (managed by the asyncio event loop thread) + self._sse_loop: asyncio.AbstractEventLoop | None = None + self._sse_tasks: dict[str, asyncio.Task[None]] = {} + self._sse_stop_events: dict[str, asyncio.Event] = {} + self._sse_async_client: httpx.AsyncClient | None = None + def upgrade_tls(self, tls_verify: Any = True, tls_cert: tuple[str, str] | None = None) -> None: - """Replace the httpx client with one using mTLS context.""" - old = self._http_client - self._http_client = httpx.Client( - timeout=httpx.Timeout( - connect=10, read=self._http_timeout, write=5, pool=self._http_timeout - ), - limits=httpx.Limits( - max_connections=self._max_poll_workers + 10, - max_keepalive_connections=min(self._max_poll_workers, 200), - ), - verify=tls_verify, - cert=tls_cert, + """Update TLS settings for future SSE connections.""" + self._tls_verify = tls_verify + self._tls_cert = tls_cert + # If the async client is running, replace it on the event loop. + if self._sse_loop is not None and self._sse_loop.is_running(): + asyncio.run_coroutine_threadsafe(self._replace_async_client(), self._sse_loop) + + async def _replace_async_client(self) -> None: + """Replace the async httpx client (called on the SSE event loop). + + Closing the old client terminates its underlying connections, which + causes active ``_node_sse_task`` coroutines to raise and reconnect + using the new client with updated TLS settings. + """ + old = self._sse_async_client + self._sse_async_client = httpx.AsyncClient( + timeout=httpx.Timeout(connect=10, read=None, write=5, pool=None), + limits=httpx.Limits(max_connections=2000, max_keepalive_connections=1500), + verify=self._tls_verify, + cert=self._tls_cert, ) - # Don't close old client — concurrent _fetch_node() threads may still - # be using it. It will be GC'd once all references are released, and - # the current client is closed in stop(). - del old + if old is not None: + await old.aclose() # -- lifecycle ----------------------------------------------------------- @@ -124,7 +126,7 @@ class ClusterCollector: self._running = True for target, name in [ (self._discovery_loop, "console-discovery"), - (self._poll_loop, "console-poll"), + (self._sse_manager_thread, "console-sse"), ]: t = threading.Thread(target=target, name=name, daemon=True) t.start() @@ -132,10 +134,22 @@ class ClusterCollector: log.info("ClusterCollector started") def stop(self) -> None: - """Stop all threads and clean up resources.""" + """Stop all threads and clean up resources. + + Sets ``_running = False`` which causes the SSE manager coroutine to + exit naturally (its ``while self._running`` loop terminates), running + its ``finally`` cleanup (cancel tasks, close AsyncClient). + """ self._running = False - self._poll_pool.shutdown(wait=False) - self._http_client.close() + # Request cancellation of all SSE tasks so they don't block the + # manager's cleanup. The manager coroutine exits when _running is + # False and handles remaining task cancellation in its finally block. + if self._sse_loop is not None and self._sse_loop.is_running(): + for node_id in list(self._sse_tasks): + asyncio.run_coroutine_threadsafe(self._stop_node(node_id), self._sse_loop) + # Wait for background threads to finish their shutdown. + for t in self._threads: + t.join(timeout=5) log.info("ClusterCollector stopped") def _fanout(self, event: dict[str, Any]) -> None: @@ -145,6 +159,137 @@ class ClusterCollector: with contextlib.suppress(queue.Full): q.put_nowait(event) + # -- auth helpers -------------------------------------------------------- + + def _auth_headers(self) -> dict[str, str]: + """Build auth headers for the current SSE connection.""" + if self._token_manager is not None: + return {"Authorization": f"Bearer {self._token_manager.token}"} + return dict(self._static_auth) if self._static_auth else {} + + # -- SSE manager --------------------------------------------------------- + + def _sse_manager_thread(self) -> None: + """Run asyncio event loop that manages all node SSE connections.""" + self._sse_loop = asyncio.new_event_loop() + try: + self._sse_loop.run_until_complete(self._sse_manager()) + finally: + self._sse_loop.close() + self._sse_loop = None + + async def _sse_manager(self) -> None: + """Top-level coroutine — runs until collector stops.""" + self._sse_async_client = httpx.AsyncClient( + timeout=httpx.Timeout(connect=10, read=None, write=5, pool=None), + limits=httpx.Limits(max_connections=2000, max_keepalive_connections=1500), + verify=self._tls_verify, + cert=self._tls_cert, + ) + try: + while self._running: + await asyncio.sleep(1) + finally: + # Cancel all remaining tasks + for task in self._sse_tasks.values(): + task.cancel() + for task in self._sse_tasks.values(): + with contextlib.suppress(asyncio.CancelledError): + await task + self._sse_tasks.clear() + self._sse_stop_events.clear() + await self._sse_async_client.aclose() + self._sse_async_client = None + + async def _start_node(self, node_id: str) -> None: + """Start an SSE task for a node (called on the SSE event loop).""" + if node_id in self._sse_tasks: + return # already running + stop = asyncio.Event() + self._sse_stop_events[node_id] = stop + self._sse_tasks[node_id] = asyncio.create_task( + self._node_sse_task(node_id, stop), + name=f"sse-{node_id}", + ) + + async def _stop_node(self, node_id: str) -> None: + """Stop an SSE task for a node (called on the SSE event loop).""" + stop = self._sse_stop_events.pop(node_id, None) + if stop: + stop.set() + task = self._sse_tasks.pop(node_id, None) + if task: + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + + async def _node_sse_task(self, node_id: str, stop_event: asyncio.Event) -> None: + """Persistent SSE connection to a single server node.""" + backoff = 1.0 + while not stop_event.is_set() and self._running: + url = self._get_node_url(node_id) + if not url or self._sse_async_client is None: + break + base = url.rstrip("/") + try: + async with httpx_sse.aconnect_sse( + self._sse_async_client, + "GET", + f"{base}/v1/api/events/global", + params={"expected_node_id": node_id}, + headers=self._auth_headers(), + ) as source: + if source.response.status_code == 409: + log.warning("Node identity mismatch for %s at %s", node_id, url) + self._mark_unreachable(node_id) + break # stop reconnecting — wrong node at this URL + source.response.raise_for_status() + async for sse in source.aiter_sse(): + if stop_event.is_set(): + break + if not sse.data: + continue # ping/comment frame + try: + data = json.loads(sse.data) + except json.JSONDecodeError: + log.debug("Invalid SSE JSON from node %s", node_id) + continue + etype = data.get("type", "") + if etype == "node_snapshot": + # Client-side identity check (defense in depth) + if data.get("node_id") != node_id: + log.warning( + "Snapshot node_id mismatch: expected %s, got %s", + node_id, + data.get("node_id"), + ) + self._mark_unreachable(node_id) + break + self._apply_snapshot(node_id, data) + backoff = 1.0 + else: + self._apply_delta(node_id, data) + except asyncio.CancelledError: + raise + except Exception: + log.debug("SSE error for node %s", node_id, exc_info=True) + self._mark_unreachable(node_id) + await asyncio.sleep(min(backoff, 30) + random.random()) + backoff = min(backoff * 2, 30) + + def _get_node_url(self, node_id: str) -> str: + """Get the server URL for a node (thread-safe).""" + with self._lock: + node = self._nodes.get(node_id) + return node.server_url if node else "" + + def _mark_unreachable(self, node_id: str) -> None: + """Mark a node as unreachable (thread-safe).""" + with self._lock: + node = self._nodes.get(node_id) + if node: + node.reachable = False + # -- node discovery ------------------------------------------------------ def _discovery_loop(self) -> None: @@ -161,6 +306,8 @@ class ClusterCollector: raw_services = self._storage.list_services("server", max_age_seconds=120) active_ids = set() pending_events: list[dict[str, Any]] = [] + new_nodes: list[str] = [] + lost_nodes: list[str] = [] with self._lock: for svc in raw_services: @@ -183,6 +330,7 @@ class ClusterCollector: max_ws=meta.get("max_ws", 10), ) pending_events.append({"type": "node_joined", "node_id": nid}) + new_nodes.append(nid) log.info("Discovered node: %s", nid) else: self._nodes[nid].server_url = url or self._nodes[nid].server_url @@ -193,10 +341,18 @@ class ClusterCollector: for nid in lost: del self._nodes[nid] pending_events.append({"type": "node_lost", "node_id": nid}) + lost_nodes.append(nid) log.info("Lost node: %s", nid) for event in pending_events: self._fanout(event) + # Manage SSE tasks for new/lost nodes + if self._sse_loop is not None and self._sse_loop.is_running(): + for nid in new_nodes: + asyncio.run_coroutine_threadsafe(self._start_node(nid), self._sse_loop) + for nid in lost_nodes: + asyncio.run_coroutine_threadsafe(self._stop_node(nid), self._sse_loop) + # Notify the routing layer so it can refresh its hash-ring cache # when the rebalancer has published a new version. if self._router is not None: @@ -211,114 +367,67 @@ class ClusterCollector: self._router.version, ) - # -- polling ------------------------------------------------------------- + # -- SSE event handlers -------------------------------------------------- - def _poll_loop(self) -> None: - """Periodically fetch /v1/api/dashboard from each node.""" - while self._running: - try: - self._poll_all_nodes() - except Exception: - log.exception("Poll loop error") - time.sleep(self._poll_interval) + def _reconcile_node( + self, node_id: str, node: NodeSnapshot, new_ws_list: list[dict[str, Any]] + ) -> list[dict[str, Any]]: + """Diff new workstream data against the current snapshot. - @staticmethod - def _node_jitter(node_id: str, window: float) -> float: - """Deterministic per-node delay within a sliding window. - - Uses a Mersenne prime (2^31 - 1) to hash the node_id into a - stable offset so each node is polled at a different point in - the cycle. The offset is consistent across restarts for the - same node_id, giving an even spread without randomness. + Returns a list of pending events. Caller must hold ``_lock``. + Updates ``node.workstreams`` in place. """ - h = hash(node_id) & 0x7FFFFFFF # positive 31-bit - return (h % 2147483647) / 2147483647 * window # M31 = 2^31 - 1 + pending: list[dict[str, Any]] = [] + old_ids = {k for k in node.workstreams if k} + new_ws: dict[str, dict[str, Any]] = {} + for ws in new_ws_list: + ws_id = ws.get("id", "") + if not ws_id: + continue + ws["node"] = node_id + ws["server_url"] = node.server_url + new_ws[ws_id] = ws + new_ids = set(new_ws.keys()) + # Additions + for ws_id in sorted(new_ids - old_ids): + ws = new_ws[ws_id] + pending.append( + { + "type": "ws_created", + "ws_id": ws_id, + "name": ws.get("name", ""), + "node_id": node_id, + } + ) + # Removals + for ws_id in sorted(old_ids - new_ids): + pending.append({"type": "ws_closed", "ws_id": ws_id}) + # State and name changes on existing workstreams + for ws_id in sorted(new_ids & old_ids): + old_ws = node.workstreams.get(ws_id, {}) + new_w = new_ws[ws_id] + old_state = old_ws.get("state", "") + new_state = new_w.get("state", "") + if old_state != new_state: + pending.append( + { + "type": "cluster_state", + "ws_id": ws_id, + "state": new_state, + "node_id": node_id, + "tokens": new_w.get("tokens", 0), + "content": new_w.get("content", ""), + } + ) + old_name = old_ws.get("name", "") + new_name = new_w.get("name", "") + if old_name != new_name and new_name: + pending.append({"type": "ws_rename", "ws_id": ws_id, "name": new_name}) + node.workstreams = new_ws + return pending - def _poll_all_nodes(self) -> None: - """Fetch dashboard data from all known nodes in parallel. - - Submissions are throttled by the thread pool size to avoid a - thundering herd — at most ``max_poll_workers`` concurrent HTTP - requests are in flight at any time. Each worker sleeps a - deterministic per-node jitter (derived from its node_id) to - spread requests across the first half of the poll interval. - """ - # Snapshot current auth header for this poll cycle. Per-request - # headers avoid mutating shared client state (thread-safe). - if self._token_manager is not None: - poll_headers: dict[str, str] | None = { - "Authorization": f"Bearer {self._token_manager.token}" - } - else: - poll_headers = self._static_auth - with self._lock: - targets = [ - (n.node_id, n.server_url) - for n in self._nodes.values() - if n.server_url and n.server_url.startswith("http") - ] - - if not targets: - return - - jitter_window = self._poll_interval / 2 - - def _jittered_fetch( - nid: str, url: str, headers: dict[str, str] | None - ) -> tuple[dict[str, Any], dict[str, Any]]: - delay = self._node_jitter(nid, jitter_window) - if delay > 0.1: - time.sleep(delay) - return self._fetch_node(nid, url, headers) - - futures = { - self._poll_pool.submit(_jittered_fetch, nid, url, poll_headers): nid - for nid, url in targets - } - for future in as_completed(futures): - nid = futures[future] - try: - dashboard, health = future.result() - self._apply_poll(nid, dashboard, health) - except httpx.HTTPStatusError as exc: - if exc.response.status_code in (401, 403): - log.warning( - "Auth failure polling node %s: HTTP %d", nid, exc.response.status_code - ) - else: - log.debug("Failed to poll node %s: HTTP %d", nid, exc.response.status_code) - with self._lock: - if nid in self._nodes: - self._nodes[nid].reachable = False - except Exception: - log.warning("Failed to poll node %s", nid, exc_info=True) - with self._lock: - if nid in self._nodes: - self._nodes[nid].reachable = False - - def _fetch_node( - self, - node_id: str, - server_url: str, - extra_headers: dict[str, str] | None = None, - ) -> tuple[dict[str, Any], dict[str, Any]]: - """Fetch /v1/api/dashboard and /health from a single node.""" - base = server_url.rstrip("/") - dash_resp = self._http_client.get(f"{base}/v1/api/dashboard", headers=extra_headers) - dash_resp.raise_for_status() - dash_data: dict[str, Any] = dash_resp.json() - try: - health_resp = self._http_client.get(f"{base}/health", headers=extra_headers) - health_data: dict[str, Any] = health_resp.json() - except Exception: - log.debug("Failed to fetch health from %s", node_id, exc_info=True) - health_data = {} - return dash_data, health_data - - def _apply_poll(self, node_id: str, dashboard: dict[str, Any], health: dict[str, Any]) -> None: - """Apply polled data to the in-memory node snapshot.""" - ws_list = dashboard.get("workstreams", []) - aggregate = dashboard.get("aggregate", {}) + def _apply_snapshot(self, node_id: str, data: dict[str, Any]) -> None: + """Apply a ``node_snapshot`` SSE event to the in-memory state.""" pending_events: list[dict[str, Any]] = [] with self._lock: node = self._nodes.get(node_id) @@ -326,63 +435,116 @@ class ClusterCollector: return node.last_seen = time.monotonic() node.reachable = True - node.health = health - node.aggregate = aggregate - # Build new workstream map - old_ids = {k for k in node.workstreams if k} - new_ws: dict[str, dict[str, Any]] = {} - for ws in ws_list: - ws_id = ws.get("id", "") - if not ws_id: - continue - ws["node"] = node_id - ws["server_url"] = node.server_url - new_ws[ws_id] = ws - new_ids = set(new_ws.keys()) - # Detect additions not yet known to SSE clients - for ws_id in sorted(new_ids - old_ids): - ws = new_ws[ws_id] + node.health = data.get("health", {}) + node.aggregate = data.get("aggregate", {}) + pending_events = self._reconcile_node(node_id, node, data.get("workstreams", [])) + for event in pending_events: + self._fanout(event) + + def _apply_delta(self, node_id: str, data: dict[str, Any]) -> None: + """Apply a single delta SSE event to the in-memory state.""" + etype = data.get("type", "") + if not etype: + return + pending_events: list[dict[str, Any]] = [] + with self._lock: + node = self._nodes.get(node_id) + if not node: + return + node.last_seen = time.monotonic() + + if etype == "ws_state": + # Server emits ws_state; translate to cluster_state for browser. + # Only fan out if the workstream is known — a ws_state arriving + # before ws_created (race) is silently absorbed on reconnect. + ws_id = data.get("ws_id", "") + ws = node.workstreams.get(ws_id) + if ws: + ws["state"] = data.get("state", ws.get("state", "")) + ws["tokens"] = data.get("tokens", ws.get("tokens", 0)) + ws["context_ratio"] = data.get("context_ratio", ws.get("context_ratio", 0)) + ws["activity"] = data.get("activity", ws.get("activity", "")) + ws["activity_state"] = data.get("activity_state", ws.get("activity_state", "")) + pending_events.append( + { + "type": "cluster_state", + "ws_id": ws_id, + "state": data.get("state", ""), + "node_id": node_id, + "tokens": data.get("tokens", 0), + "content": data.get("content", ""), + } + ) + + elif etype == "ws_activity": + ws_id = data.get("ws_id", "") + ws = node.workstreams.get(ws_id) + if ws: + ws["activity"] = data.get("activity", "") + ws["activity_state"] = data.get("activity_state", "") + # Activity events are not forwarded to cluster SSE — only state changes + + elif etype == "ws_created": + ws_id = data.get("ws_id", "") + if ws_id and ws_id not in node.workstreams: + node.workstreams[ws_id] = { + "id": ws_id, + "name": data.get("name", ""), + "state": "idle", + "node": node_id, + "server_url": node.server_url, + "model": data.get("model", ""), + "model_alias": data.get("model_alias", ""), + "tokens": 0, + "context_ratio": 0.0, + "activity": "", + "activity_state": "", + "tool_calls": 0, + "title": "", + } pending_events.append( { "type": "ws_created", "ws_id": ws_id, - "name": ws.get("name", ""), + "name": data.get("name", ""), "node_id": node_id, } ) - # Detect removals - for ws_id in sorted(old_ids - new_ids): + + elif etype == "ws_closed": + ws_id = data.get("ws_id", "") + node.workstreams.pop(ws_id, None) pending_events.append({"type": "ws_closed", "ws_id": ws_id}) - # Detect state changes on existing workstreams - for ws_id in sorted(new_ids & old_ids): - old_ws = node.workstreams.get(ws_id, {}) - new_w = new_ws[ws_id] - old_state = old_ws.get("state", "") - new_state = new_w.get("state", "") - if old_state != new_state: - pending_events.append( - { - "type": "ws_state", - "ws_id": ws_id, - "state": new_state, - "node_id": node_id, - "tokens": new_w.get("tokens", 0), - "content": new_w.get("content", ""), - } - ) - # Detect name/title changes - old_name = old_ws.get("name", "") - new_name = new_w.get("name", "") - if old_name != new_name and new_name: - pending_events.append( - { - "type": "ws_rename", - "ws_id": ws_id, - "name": new_name, - } - ) - node.workstreams = new_ws - # Fan out diffs to SSE listeners outside the lock + + elif etype == "ws_rename": + ws_id = data.get("ws_id", "") + name = data.get("name", "") + ws = node.workstreams.get(ws_id) + if ws and name: + ws["name"] = name + pending_events.append({"type": "ws_rename", "ws_id": ws_id, "name": name}) + + elif etype == "health_changed": + # Update the health dict's circuit state in-place + circuit = data.get("circuit_state", "") + if circuit: + if not node.health: + node.health = {} + backend = node.health.setdefault("backend", {}) + backend["circuit_state"] = circuit + backend["status"] = "up" if circuit == "closed" else "down" + node.health["status"] = "ok" if circuit == "closed" else "degraded" + # Not forwarded to cluster SSE — next snapshot refreshes UI + + elif etype == "aggregate": + node.aggregate = { + "total_tokens": data.get("total_tokens", 0), + "total_tool_calls": data.get("total_tool_calls", 0), + "active_count": data.get("active_count", 0), + "total_count": data.get("total_count", 0), + } + # Not forwarded to cluster SSE — overview queries read from snapshot + for event in pending_events: self._fanout(event) diff --git a/turnstone/console/server.py b/turnstone/console/server.py index dc4b0e35..101edc49 100644 --- a/turnstone/console/server.py +++ b/turnstone/console/server.py @@ -6140,12 +6140,6 @@ def main() -> None: default=8090, help="Port to listen on (default: 8090)", ) - parser.add_argument( - "--poll-interval", - type=float, - default=10.0, - help="Node polling interval in seconds (default: 10)", - ) from turnstone.core.log import add_log_args add_log_args(parser) @@ -6228,7 +6222,6 @@ def main() -> None: collector = ClusterCollector( storage=auth_storage, - poll_interval=args.poll_interval, auth_token=collector_token if collector_token_mgr is None else "", token_manager=collector_token_mgr, router=router, diff --git a/turnstone/core/config.py b/turnstone/core/config.py index 71671829..40f9b872 100644 --- a/turnstone/core/config.py +++ b/turnstone/core/config.py @@ -117,7 +117,6 @@ _CONFIG_MAP: dict[str, dict[str, str]] = { "host": "host", "port": "port", "url": "console_url", - "poll_interval": "poll_interval", "log_level": "log_level", }, "auth": { diff --git a/turnstone/core/healthcheck.py b/turnstone/core/healthcheck.py index 87aa6bfb..64f194b5 100644 --- a/turnstone/core/healthcheck.py +++ b/turnstone/core/healthcheck.py @@ -43,6 +43,7 @@ class BackendHealthMonitor: provider: str = "openai", initial_model: str = "", on_model_changed: Callable[[str, int | None], None] | None = None, + on_state_changed: Callable[[str], None] | None = None, ) -> None: self._client = client self._probe_interval = probe_interval @@ -54,6 +55,7 @@ class BackendHealthMonitor: self._provider = provider self._last_detected_model = initial_model self._on_model_changed = on_model_changed + self._on_state_changed = on_state_changed self._lock = threading.Lock() self._state = CircuitState.CLOSED @@ -82,8 +84,17 @@ class BackendHealthMonitor: # Passive tracking (called by request path) # ------------------------------------------------------------------ + def _fire_state_callback(self, state_val: str | None) -> None: + """Fire on_state_changed callback outside the lock.""" + if state_val is not None and self._on_state_changed is not None: + try: + self._on_state_changed(state_val) + except Exception: + log.debug("on_state_changed callback error", exc_info=True) + def record_success(self) -> None: """Called on successful LLM call. Resets failure count, closes circuit.""" + state_to_dispatch: str | None = None with self._lock: self._consecutive_failures = 0 if self._state != CircuitState.CLOSED: @@ -93,9 +104,12 @@ class BackendHealthMonitor: self._last_state_change = time.monotonic() log.info("Circuit breaker CLOSED (was %s): backend recovered", prev.value) self._update_metrics() + state_to_dispatch = self._state.value + self._fire_state_callback(state_to_dispatch) def record_failure(self) -> None: """Called on LLM call failure. May open circuit.""" + state_to_dispatch: str | None = None with self._lock: self._consecutive_failures += 1 if self._state == CircuitState.HALF_OPEN: @@ -105,6 +119,7 @@ class BackendHealthMonitor: self._last_state_change = time.monotonic() log.warning("Circuit breaker OPEN: probe failed in HALF_OPEN") self._update_metrics() + state_to_dispatch = self._state.value elif ( self._state == CircuitState.CLOSED and self._consecutive_failures >= self._failure_threshold @@ -116,6 +131,8 @@ class BackendHealthMonitor: self._consecutive_failures, ) self._update_metrics() + state_to_dispatch = self._state.value + self._fire_state_callback(state_to_dispatch) # ------------------------------------------------------------------ # Query helpers @@ -246,7 +263,12 @@ class BackendHealthMonitor: # ------------------------------------------------------------------ def _update_metrics(self) -> None: - """Push state to metrics collector. Called with *self._lock* held.""" + """Push circuit-breaker state to metrics collector. + + Called with *self._lock* held. State-change callbacks are dispatched + by the callers (``record_success`` / ``record_failure``) after the + lock is released, not by this method. + """ from turnstone.core.metrics import metrics metrics.set_backend_status(self._state == CircuitState.CLOSED) diff --git a/turnstone/sdk/events.py b/turnstone/sdk/events.py index ccbcb65c..0dbffccd 100644 --- a/turnstone/sdk/events.py +++ b/turnstone/sdk/events.py @@ -300,6 +300,36 @@ class ClusterSnapshotEvent(ClusterEvent): timestamp: float = 0.0 +@dataclass +class NodeSnapshotEvent(ClusterEvent): + """Full node state delivered on SSE connect to ``/v1/api/events/global``.""" + + type: str = "node_snapshot" + node_id: str = "" + workstreams: list[dict[str, Any]] = field(default_factory=list) + health: dict[str, Any] = field(default_factory=dict) + aggregate: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class HealthChangedEvent(ClusterEvent): + """Circuit breaker state transition on a server node.""" + + type: str = "health_changed" + circuit_state: str = "" + + +@dataclass +class AggregateEvent(ClusterEvent): + """Periodic aggregate metrics from a server node.""" + + type: str = "aggregate" + total_tokens: int = 0 + total_tool_calls: int = 0 + active_count: int = 0 + total_count: int = 0 + + # --------------------------------------------------------------------------- # Type registries (built after all classes are defined) # --------------------------------------------------------------------------- @@ -354,5 +384,8 @@ _CLUSTER_REGISTRY: dict[str, type[ClusterEvent]] = { ClusterWsClosedEvent, ClusterWsRenameEvent, ClusterSnapshotEvent, + NodeSnapshotEvent, + HealthChangedEvent, + AggregateEvent, ] } diff --git a/turnstone/sdk/server.py b/turnstone/sdk/server.py index d04bb6c6..c6f5ac9f 100644 --- a/turnstone/sdk/server.py +++ b/turnstone/sdk/server.py @@ -37,6 +37,7 @@ from turnstone.sdk._base import _BaseClient from turnstone.sdk._sync import _SyncRunner from turnstone.sdk._types import TurnResult from turnstone.sdk.events import ( + ClusterEvent, ContentEvent, ErrorEvent, ReasoningEvent, @@ -199,6 +200,22 @@ class AsyncTurnstoneServer(_BaseClient): async for data in self._stream_sse("/v1/api/events/global"): yield ServerEvent.from_dict(data) + async def stream_node_events( + self, *, expected_node_id: str = "" + ) -> AsyncIterator[ClusterEvent]: + """Iterate over node-level SSE events (snapshot + deltas). + + Connects to ``/v1/api/events/global`` with the optional + ``expected_node_id`` param for identity verification. Yields + ``ClusterEvent`` instances (``NodeSnapshotEvent``, ``HealthChangedEvent``, + etc.) suitable for console collector consumption. + """ + params: dict[str, str] = {} + if expected_node_id: + params["expected_node_id"] = expected_node_id + async for data in self._stream_sse("/v1/api/events/global", params=params): + yield ClusterEvent.from_dict(data) + # -- high-level convenience ---------------------------------------------- async def send_and_wait( @@ -509,6 +526,11 @@ class TurnstoneServer: def stream_global_events(self) -> Iterator[ServerEvent]: return self._runner.run_iter(self._async.stream_global_events()) + def stream_node_events(self, *, expected_node_id: str = "") -> Iterator[ClusterEvent]: + return self._runner.run_iter( + self._async.stream_node_events(expected_node_id=expected_node_id) + ) + # -- high-level convenience ---------------------------------------------- def send_and_wait( diff --git a/turnstone/server.py b/turnstone/server.py index cfa90bfa..fd8a636f 100644 --- a/turnstone/server.py +++ b/turnstone/server.py @@ -919,17 +919,111 @@ async def events_sse(request: Request) -> Response: return EventSourceResponse(event_generator(), ping=5) +def _build_node_snapshot(app_state: Any) -> dict[str, Any]: + """Build a complete node state snapshot for SSE consumers. + + Includes workstream list, health, and aggregate — everything the console + collector needs to populate a ``NodeSnapshot`` without polling. + """ + from turnstone.core.memory import get_workstream_display_name + + mgr: WorkstreamManager = app_state.workstreams + wss = mgr.list_all() + total_tokens = 0 + total_tool_calls = 0 + active_count = 0 + ws_list = [] + for ws in wss: + ui = ws.ui + if hasattr(ui, "_ws_lock"): + with ui._ws_lock: # type: ignore[union-attr] + tok = ui._ws_prompt_tokens + ui._ws_completion_tokens # type: ignore[union-attr] + tc = sum(ui._ws_tool_calls.values()) # type: ignore[union-attr] + ctx = ui._ws_context_ratio # type: ignore[union-attr] + activity = ui._ws_current_activity # type: ignore[union-attr] + activity_state = ui._ws_activity_state # type: ignore[union-attr] + else: + tok = tc = 0 + ctx = 0.0 + activity = activity_state = "" + total_tokens += tok + total_tool_calls += tc + if ws.state.value != "idle": + active_count += 1 + title = "" + if ws.session: + title = get_workstream_display_name(ws.session.ws_id) or "" + ws_list.append( + { + "id": ws.id, + "name": ws.name, + "state": ws.state.value, + "title": title, + "tokens": tok, + "context_ratio": round(ctx, 3), + "activity": activity, + "activity_state": activity_state, + "tool_calls": tc, + "model": ws.session.model if ws.session else "", + "model_alias": ws.session.model_alias if ws.session else "", + } + ) + return { + "type": "node_snapshot", + "node_id": getattr(app_state, "node_id", ""), + "workstreams": ws_list, + "health": _build_health_dict(app_state), + "aggregate": { + "total_tokens": total_tokens, + "total_tool_calls": total_tool_calls, + "active_count": active_count, + "total_count": len(ws_list), + }, + } + + async def global_events_sse(request: Request) -> Response: - """GET /v1/api/events/global — global SSE event stream.""" + """GET /v1/api/events/global — global SSE event stream. + + Supports optional ``?expected_node_id=X`` query parameter for node identity + verification. If present and the server's node_id does not match, returns + 409 Conflict immediately. + + On connect, emits a ``node_snapshot`` event with the full node state + (workstreams, health, aggregate) followed by real-time delta events. + The snapshot and listener registration are atomic — no events are lost. + """ + # -- Node identity check -------------------------------------------------- + expected = request.query_params.get("expected_node_id") + actual_node_id = getattr(request.app.state, "node_id", "") + if expected and expected != actual_node_id: + return JSONResponse( + { + "error": "node_id mismatch" if actual_node_id else "node_id unavailable", + "expected": expected, + "actual": actual_node_id, + }, + status_code=409, + ) + + # -- Atomic snapshot + listener registration ------------------------------ client_queue: queue.Queue[dict[str, Any]] = queue.Queue(maxsize=1000) listeners = request.app.state.global_listeners listeners_lock = request.app.state.global_listeners_lock + + # Hold the listeners lock while building the snapshot AND registering. + # The fanout thread also acquires this lock when snapshotting the listener + # list, so events that land on global_queue during snapshot build will be + # distributed to our queue after we release — gap-free. with listeners_lock: + snapshot = _build_node_snapshot(request.app.state) listeners.append(client_queue) async def event_generator() -> AsyncGenerator[dict[str, str], None]: _metrics.record_sse_connect() try: + # Emit snapshot as first event + yield {"data": json.dumps(snapshot)} loop = asyncio.get_running_loop() executor = request.app.state.sse_executor while True: @@ -1100,17 +1194,20 @@ def _count_ws_states(wss: list[Workstream]) -> dict[str, int]: return counts -async def health(request: Request) -> JSONResponse: - """GET /health — server health status.""" - mgr: WorkstreamManager = request.app.state.workstreams +def _build_health_dict(app_state: Any) -> dict[str, Any]: + """Assemble health status dict from app state. + + Shared by the ``/health`` endpoint and the global SSE snapshot. + """ + mgr: WorkstreamManager = app_state.workstreams wss = mgr.list_all() states = _count_ws_states(wss) - monitor = getattr(request.app.state, "health_monitor", None) + monitor = getattr(app_state, "health_monitor", None) backend_ok = monitor.is_healthy if monitor else True data: dict[str, Any] = { "status": "ok" if backend_ok else "degraded", "version": __version__, - "node_id": getattr(request.app.state, "node_id", ""), + "node_id": getattr(app_state, "node_id", ""), "uptime_seconds": round(time.monotonic() - _metrics.start_time, 2), "model": _metrics.model, "max_ws": mgr.max_workstreams, @@ -1120,14 +1217,19 @@ async def health(request: Request) -> JSONResponse: "circuit_state": monitor.circuit_state.value if monitor else "closed", }, } - mc = getattr(request.app.state, "mcp_client", None) + mc = getattr(app_state, "mcp_client", None) if mc: data["mcp"] = { "servers": mc.server_count, "resources": mc.resource_count, "prompts": mc.prompt_count, } - return JSONResponse(data) + return data + + +async def health(request: Request) -> JSONResponse: + """GET /health — server health status.""" + return JSONResponse(_build_health_dict(request.app.state)) async def metrics_endpoint(request: Request) -> Response: @@ -1546,10 +1648,21 @@ async def create_workstream(request: Request) -> JSONResponse: ws.session.set_watch_runner( runner, dispatch_fn=_make_watch_dispatch(ws, ws.session, ws.ui) ) + # Emit creation event on global queue for SSE consumers (console) + gq: queue.Queue[dict[str, Any]] = request.app.state.global_queue + with contextlib.suppress(queue.Full): + gq.put_nowait( + { + "type": "ws_created", + "ws_id": ws.id, + "name": ws.name, + "model": ws.session.model if ws.session else "", + "model_alias": ws.session.model_alias if ws.session else "", + } + ) # Emit eviction event if a workstream was evicted to make room evicted = mgr.last_evicted if evicted is not None: - gq: queue.Queue[dict[str, Any]] = request.app.state.global_queue with contextlib.suppress(queue.Full): gq.put_nowait( { @@ -1669,6 +1782,9 @@ async def close_workstream(request: Request) -> JSONResponse: ws_id = str(body.get("ws_id", "")) mgr = request.app.state.workstreams if mgr.close(ws_id): + gq: queue.Queue[dict[str, Any]] = request.app.state.global_queue + with contextlib.suppress(queue.Full): + gq.put_nowait({"type": "ws_closed", "ws_id": ws_id, "reason": "closed"}) return JSONResponse({"status": "ok"}) return JSONResponse({"error": "Cannot close last workstream"}, status_code=400) @@ -2076,6 +2192,58 @@ async def internal_migrate(request: Request) -> JSONResponse: # --------------------------------------------------------------------------- +def _emit_health_changed(circuit_state: str, gq: queue.Queue[dict[str, Any]]) -> None: + """Push a health_changed event onto the global SSE queue. + + Called from the BackendHealthMonitor callback on circuit breaker transitions. + """ + with contextlib.suppress(queue.Full): + gq.put_nowait({"type": "health_changed", "circuit_state": circuit_state}) + + +def _aggregate_emitter_thread( + mgr: WorkstreamManager, + global_queue: queue.Queue[dict[str, Any]], + interval: float = 10.0, +) -> None: + """Periodically emit aggregate token/tool_call totals on the global SSE queue. + + Runs as a daemon thread so the console receives periodic updates without + having to poll ``/v1/api/dashboard``. + """ + while True: + time.sleep(interval) + total_tokens = 0 + total_tool_calls = 0 + active_count = 0 + try: + for ws in mgr.list_all(): + ui = ws.ui + if hasattr(ui, "_ws_lock"): + with ui._ws_lock: # type: ignore[union-attr] + tok = ui._ws_prompt_tokens + ui._ws_completion_tokens # type: ignore[union-attr] + tc = sum(ui._ws_tool_calls.values()) # type: ignore[union-attr] + else: + tok = 0 + tc = 0 + total_tokens += tok + total_tool_calls += tc + if ws.state.value != "idle": + active_count += 1 + with contextlib.suppress(queue.Full): + global_queue.put_nowait( + { + "type": "aggregate", + "total_tokens": total_tokens, + "total_tool_calls": total_tool_calls, + "active_count": active_count, + "total_count": len(mgr.list_all()), + } + ) + except Exception: + log.debug("Aggregate emitter error", exc_info=True) + + def _idle_cleanup_thread( mgr: WorkstreamManager, timeout_sec: float, @@ -2134,6 +2302,13 @@ async def _lifespan(app: Starlette) -> AsyncGenerator[None, None]: daemon=True, ) fanout.start() + # Start aggregate emitter thread for SSE consumers + agg_emitter = threading.Thread( + target=_aggregate_emitter_thread, + args=(app.state.workstreams, app.state.global_queue), + daemon=True, + ) + agg_emitter.start() # Start idle cleanup thread if configured if app.state.idle_timeout > 0: cleanup = threading.Thread( @@ -2606,6 +2781,13 @@ def main() -> None: if new_reg is not None: new_reg.shutdown() + # Set up global event queue for state-change broadcasts (created early so + # the health monitor callback can reference it). + global_queue: queue.Queue[dict[str, Any]] = queue.Queue(maxsize=10000) + global_listeners: list[queue.Queue[dict[str, Any]]] = [] + global_listeners_lock = threading.Lock() + WebUI._global_queue = global_queue + health_monitor = BackendHealthMonitor( client=client, probe_interval=config_store.get("health.backend_probe_interval"), @@ -2615,6 +2797,7 @@ def main() -> None: provider=provider_name, initial_model=model, on_model_changed=_handle_model_change, + on_state_changed=lambda state: _emit_health_changed(state, global_queue), ) health_monitor.start() @@ -2628,12 +2811,6 @@ def main() -> None: trusted_proxies=config_store.get("ratelimit.trusted_proxies"), ) - # Set up global event queue for state-change broadcasts - global_queue: queue.Queue[dict[str, Any]] = queue.Queue(maxsize=10000) - global_listeners: list[queue.Queue[dict[str, Any]]] = [] - global_listeners_lock = threading.Lock() - WebUI._global_queue = global_queue - # Config builders — shared between startup logging and session factory. # Re-read from ConfigStore each call so hot-reload works. from turnstone.core.judge import JudgeConfig