From 5d0ec99449b4ac68bd02c9be125a27afb763aa99 Mon Sep 17 00:00:00 2001 From: Patrick Buckley Date: Mon, 11 May 2026 00:41:10 -0700 Subject: [PATCH] feat(console): event-driven wait_for_workstream + idle cleanup via ChildEventBus MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Retire two polling patterns in coord that have clean event sources. PR 2 of 3 in the coord-completion stack; sits on top of PR #505 (reactive node discovery via PG LISTEN/NOTIFY). `wait_for_workstream` (coord's block-wait tool) polled storage every 0.5 s in a worker thread regardless of whether anything had changed — a 600 s wait incurred ~2400 round-trips. Now subscribes to a new in-process `ChildEventBus` (`turnstone/core/child_event_bus.py`) and blocks on `threading.Event.wait(min(remaining, WAIT_HEARTBEAT_INTERVAL))`: - `CoordinatorAdapter` owns the bus; `_dispatch_child_event` calls `bus.notify(child_ws_id)` after each `_enqueue_on_ui` for the state-class branch (cluster_state, ws_closed, ws_rename, intent_verdict, approval_resolved, approve_request). - Wait loop clears the Event BEFORE the storage snapshot to close the subscribe/check race; a notify between clear and the next `wait()` leaves the Event set so the loop re-reads without losing the wake-up. - 2 s heartbeat cap preserves the existing `wait_progress` SSE cadence for the sidebar UI while cutting SSE traffic ~4x vs the pre-bus 500 ms cadence in the quiescent case. - Worst-case completion latency is 2 s (vs pre-bus 0.5 s) because `set_state` buffers non-ERROR writes through `StateWriter` (async-flushed) while `emit_state` fans out immediately — a bus-driven wake can beat the flusher and read pre-transition state, then re-block until heartbeat. Deliberate trade-off; the SSE-traffic reduction outweighs the regression on the most common terminal transition. - Defense-in-depth: ownership-filter `cleaned` to own-subtree before `register_waiter` so a foreign ws_id passed by an untrusted coord LLM (prompt injection) can't observe wake-up timing as a side channel. Predicate (`_row_in_own_subtree`) requires both `parent_ws_id == coord_ws_id` AND `user_id == coord_user_id` parity — same gate strength as the existing `_is_own_subtree` mutating-op guard, so a corrupted / cross-tenant `parent_ws_id` alone can't satisfy it. Shared with `_snapshot_all` so the snapshot's `denied` shape stays in lockstep with the bus filter (Copilot review on #506). Coord idle-cleanup thread polled the storage scan every `check_every` seconds (~30 s on default 2 h timeout) even when no coord was anywhere near idle. Now subscribes to `SessionManager._state_subscribers` with a `tick_now` event and blocks on `tick_now.wait(check_every)` — any state change wakes the sweeper without waiting a full interval, AND the timeout still fires the periodic sweep for the DB-orphan-only case. A `min_sweep_interval=5 s` floor bounds DB-call traffic at ~0.2/s under sustained activity so the loop can't tight-spin `close_idle` at the rate of its own DB latency (6x improvement over the pre-refactor fixed 30 s cadence under any activity, and prompt state-change-driven wakes when below the floor). `CoordinatorClient` constructor takes `child_event_bus` as a required kwarg — there's no external SDK shape to preserve and keeping it optional would silently mask a wiring bug in any future caller. Tests construct their own `ChildEventBus()` per fixture. Tests: 16 unit tests for `ChildEventBus` (register / unregister symmetry, multi-waiter fan-out, multi-child waiter, subscribe/check race, concurrent register / notify smoke); 7 new adapter tests (bus notify fires for all 6 state-class events, drops for unknown child / wrong ws_id); 7 new coord-client wait tests (subscribe- after-terminal, notify wakes, unrelated notify doesn't wake, heartbeat fires without notify, unregister on exit, multi-waiter independence, cross-tenant denial via the user_id-parity filter); 8 idle-cleanup tests (initial sweep, heartbeat cadence, exception swallowing, stop_event clean exit, state-change wake, subscriber cleanup, mid-sweep wake, `min_sweep_interval` floor). All pass; ruff + mypy clean. Full non-live suite: 6227 passed (+2 vs prior baseline). --- tests/test_child_event_bus.py | 215 +++++++++++++++++++ tests/test_console_idle_cleanup.py | 258 ++++++++++++++++++++--- tests/test_coordinator_adapter.py | 79 +++++++ tests/test_coordinator_client.py | 256 +++++++++++++++++++++- tests/test_coordinator_end_to_end.py | 3 + turnstone/console/coordinator_adapter.py | 27 +++ turnstone/console/coordinator_client.py | 222 +++++++++++++------ turnstone/console/server.py | 99 +++++++-- turnstone/core/child_event_bus.py | 121 +++++++++++ turnstone/core/session_manager.py | 14 ++ 10 files changed, 1185 insertions(+), 109 deletions(-) create mode 100644 tests/test_child_event_bus.py create mode 100644 turnstone/core/child_event_bus.py diff --git a/tests/test_child_event_bus.py b/tests/test_child_event_bus.py new file mode 100644 index 00000000..2c6763f8 --- /dev/null +++ b/tests/test_child_event_bus.py @@ -0,0 +1,215 @@ +"""Unit tests for :class:`turnstone.core.child_event_bus.ChildEventBus`. + +The bus is the in-process wakeup primitive for ``wait_for_workstream`` +(see :mod:`turnstone.console.coordinator_client`). It's a small dict +of ws_id → set[threading.Event] under a lock — focused tests for +register/notify symmetry, no-subscriber notify, multi-waiter fan-out, +multi-child waiter, and concurrent register/notify (smoke). End-to-end +integration with the dispatch sink lives in +``test_coordinator_adapter.py`` and ``test_coordinator_client.py``. +""" + +from __future__ import annotations + +import threading +import time + +import pytest + +from turnstone.core.child_event_bus import ChildEventBus + + +def test_register_returns_event_that_starts_unset() -> None: + """A waiter must not see leftover state from before it registered — + a fresh wait should always block until the first notify.""" + bus = ChildEventBus() + event = bus.register_waiter(["ws-1"]) + assert isinstance(event, threading.Event) + assert not event.is_set() + + +def test_notify_wakes_waiter_on_matching_ws_id() -> None: + bus = ChildEventBus() + event = bus.register_waiter(["ws-1"]) + bus.notify("ws-1") + assert event.is_set() + + +def test_notify_does_not_wake_waiter_on_unrelated_ws_id() -> None: + """Different ws_ids must keep independent waiter sets — a notify on + a stranger ws can't wake the wait or the bus stops being keyed.""" + bus = ChildEventBus() + event = bus.register_waiter(["ws-1"]) + bus.notify("ws-other") + assert not event.is_set() + + +def test_notify_with_no_subscribers_is_noop() -> None: + """The dispatch sink calls notify on every translated event; the + steady state has no wait tool active. Must not raise.""" + bus = ChildEventBus() + bus.notify("ws-nobody-cares") # no exception + + +def test_multi_waiter_each_gets_independent_event() -> None: + """Two waits on the same ws_id must wake independently — clearing + one Event must not silence the other.""" + bus = ChildEventBus() + e1 = bus.register_waiter(["ws-1"]) + e2 = bus.register_waiter(["ws-1"]) + assert e1 is not e2 + bus.notify("ws-1") + assert e1.is_set() + assert e2.is_set() + + +def test_multi_child_waiter_fires_on_any_listed_ws_id() -> None: + """A wait on [A, B, C] returns a single Event registered against + all three. Notify on ANY of A/B/C must wake the wait — the + caller's snapshot re-read disambiguates which one changed.""" + bus = ChildEventBus() + event = bus.register_waiter(["ws-a", "ws-b", "ws-c"]) + bus.notify("ws-b") + assert event.is_set() + + +def test_unregister_removes_event_from_all_listed_ws_ids() -> None: + """After unregister, notify on any of the previously-watched ws_ids + must NOT wake the Event — leaks would mean every future notify on + that ws_id wakes a long-dead wait.""" + bus = ChildEventBus() + event = bus.register_waiter(["ws-a", "ws-b"]) + bus.unregister_waiter(["ws-a", "ws-b"], event) + bus.notify("ws-a") + bus.notify("ws-b") + assert not event.is_set() + + +def test_unregister_is_idempotent() -> None: + """A double-unregister must silently no-op — finally blocks may + run twice in odd shutdown paths, the bus must not raise.""" + bus = ChildEventBus() + event = bus.register_waiter(["ws-1"]) + bus.unregister_waiter(["ws-1"], event) + bus.unregister_waiter(["ws-1"], event) # no exception + + +def test_unregister_pops_empty_buckets() -> None: + """Empty per-ws_id buckets must be popped so a long-lived bus + doesn't accumulate dead keys after many waits have churned through. + Reaches into the private state — the property is structural, not + behavioral, so the assertion is also.""" + bus = ChildEventBus() + event = bus.register_waiter(["ws-1"]) + assert "ws-1" in bus._waiters + bus.unregister_waiter(["ws-1"], event) + assert "ws-1" not in bus._waiters + + +def test_unregister_keeps_bucket_with_remaining_waiters() -> None: + """Removing one waiter from a multi-waiter bucket must not drop + the others — popping the bucket would silently disable notifies + for every concurrent wait on the same ws_id.""" + bus = ChildEventBus() + e1 = bus.register_waiter(["ws-1"]) + e2 = bus.register_waiter(["ws-1"]) + bus.unregister_waiter(["ws-1"], e1) + bus.notify("ws-1") + assert not e1.is_set() + assert e2.is_set() + + +def test_empty_and_falsy_ws_ids_are_skipped_on_register() -> None: + """Defensive: ``wait_for_workstream`` cleans its inputs but the bus + is reachable from other callers in future use; falsy ids should be + silently dropped, not registered against an empty-string key.""" + bus = ChildEventBus() + event = bus.register_waiter(["", "ws-1", ""]) + # Only the real ws_id should bucket the waiter. + assert list(bus._waiters.keys()) == ["ws-1"] + bus.notify("") # no crash, no spurious wake + assert not event.is_set() + bus.notify("ws-1") + assert event.is_set() + + +def test_notify_wakes_waiter_blocking_on_event_wait() -> None: + """End-to-end wake-up latency: a wait blocked on ``Event.wait`` + must return promptly after a notify on a watched ws_id. This is + the property that retires the 0.5s polling cadence.""" + bus = ChildEventBus() + event = bus.register_waiter(["ws-1"]) + woken_at = [0.0] + + def _waiter() -> None: + event.wait(timeout=2.0) + woken_at[0] = time.monotonic() + + t = threading.Thread(target=_waiter, daemon=True) + t.start() + # Give the waiter a beat to enter Event.wait, then notify. + time.sleep(0.05) + notified_at = time.monotonic() + bus.notify("ws-1") + t.join(timeout=1.0) + assert not t.is_alive(), "waiter did not wake within 1s of notify" + # Latency budget is generous; the contract is "well under the legacy + # 0.5s poll cadence", not microsecond timing. + assert woken_at[0] - notified_at < 0.2 + + +def test_clear_before_check_race_does_not_lose_wake() -> None: + """The wait-loop pattern is ``clear(); snapshot(); ...; wait()``. + A notify between clear and wait must leave the Event set, so the + next wait returns immediately and the loop re-snapshots. Same + standard subscribe/check race the wait loop guards against.""" + bus = ChildEventBus() + event = bus.register_waiter(["ws-1"]) + # Simulate wait-loop ordering: clear, then notify "between" clear + # and the next wait. + event.clear() + bus.notify("ws-1") + # The next wait must return True immediately (set is sticky until + # the next clear). + assert event.wait(timeout=0.1) is True + + +def test_concurrent_register_and_notify_is_safe() -> None: + """Smoke test: many threads registering / notifying / unregistering + in parallel must not raise or deadlock. Doesn't assert specific + interleavings — only structural safety of the lock discipline.""" + bus = ChildEventBus() + stop = threading.Event() + errors: list[BaseException] = [] + + def _worker(ws_id: str) -> None: + try: + for _ in range(200): + if stop.is_set(): + return + ev = bus.register_waiter([ws_id]) + bus.notify(ws_id) + bus.unregister_waiter([ws_id], ev) + except BaseException as e: # noqa: BLE001 + errors.append(e) + + threads = [threading.Thread(target=_worker, args=(f"ws-{i}",), daemon=True) for i in range(8)] + for t in threads: + t.start() + for t in threads: + t.join(timeout=5.0) + stop.set() + assert not errors, f"worker threads raised: {errors!r}" + # All buckets should have been popped (every register paired with + # unregister). + assert bus._waiters == {} + + +@pytest.mark.parametrize("ws_id", ["", None]) +def test_notify_silently_ignores_falsy_ws_id(ws_id: object) -> None: + """Defensive: the dispatch sink already guards against empty + ws_ids, but a falsy slip-through must not raise.""" + bus = ChildEventBus() + event = bus.register_waiter(["ws-1"]) + bus.notify(ws_id) # type: ignore[arg-type] + assert not event.is_set() diff --git a/tests/test_console_idle_cleanup.py b/tests/test_console_idle_cleanup.py index e21a9432..a3e11c8a 100644 --- a/tests/test_console_idle_cleanup.py +++ b/tests/test_console_idle_cleanup.py @@ -5,11 +5,15 @@ lifting is in ``SessionManager.close_idle`` (covered in ``test_session_manager.py``) and ``bulk_close_stale_orphans`` (covered in ``test_storage_sqlite.py``). These tests verify the glue: -- the helper runs an initial sweep BEFORE its first sleep (cold-start +- the helper runs an initial sweep BEFORE its first wait (cold-start cleanup without blocking the lifespan), - the helper swallows exceptions so a transient DB blip can't kill the daemon thread, -- the helper exits cleanly when ``stop_event`` is set. +- the helper exits cleanly when ``stop_event`` is set, +- the helper subscribes to ``mgr.subscribe_to_state`` and a state-change + event wakes the next sweep early (event-driven, not polling), +- the helper unsubscribes when the thread exits so the subscriber + doesn't leak past one cleanup-thread lifetime. The ``stop_event`` parameter is exclusively for tests — production callers pass ``None`` and the daemon runs for process lifetime. @@ -17,28 +21,36 @@ callers pass ``None`` and the daemon runs for process lifetime. from __future__ import annotations +import contextlib import threading -from unittest.mock import patch +import time +from typing import TYPE_CHECKING from turnstone.console.server import _coord_idle_cleanup_thread +if TYPE_CHECKING: + from collections.abc import Callable + class _StubMgr: + """Minimal SessionManager substitute exposing only what the cleanup + thread touches: ``close_idle``, ``subscribe_to_state``, + ``unsubscribe_from_state``. Records call ordering for assertions + and lets the test fire state-change events manually via + :meth:`fire_state_change`. + """ + def __init__( self, *, stop_event: threading.Event, expected_calls: int, raise_after: int = -1 ) -> None: self.calls: list[float] = [] - self.sleep_calls_at_each_close: list[int] = [] self._stop_event = stop_event self._expected = expected_calls self._raise_after = raise_after - self._sleep_count = 0 + self._subscribers: list[Callable[[str, object], None]] = [] + self._sub_lock = threading.Lock() def close_idle(self, timeout_sec: float) -> list[str]: - # Snapshot how many sleeps preceded this close — lets the - # "initial sweep" test verify the first close_idle ran with - # zero preceding sleeps. - self.sleep_calls_at_each_close.append(self._sleep_count) self.calls.append(timeout_sec) try: if 0 <= self._raise_after < len(self.calls): @@ -50,39 +62,78 @@ class _StubMgr: self._stop_event.set() return [] - def record_sleep(self, _seconds: float) -> None: - self._sleep_count += 1 + def subscribe_to_state(self, callback: Callable[[str, object], None]) -> None: + with self._sub_lock: + self._subscribers.append(callback) + + def unsubscribe_from_state(self, callback: Callable[[str, object], None]) -> None: + with self._sub_lock, contextlib.suppress(ValueError): + self._subscribers.remove(callback) + + @property + def subscribers_count(self) -> int: + with self._sub_lock: + return len(self._subscribers) + + def fire_state_change(self, ws_id: str = "ws-x", state: object = "idle") -> None: + with self._sub_lock: + snapshot = list(self._subscribers) + for cb in snapshot: + cb(ws_id, state) def _run_until_done(mgr: _StubMgr, stop_event: threading.Event, timeout_sec: float) -> None: - with patch("turnstone.console.server.time.sleep", mgr.record_sleep): - thread = threading.Thread( - target=_coord_idle_cleanup_thread, - args=(mgr, timeout_sec, stop_event), - daemon=True, - ) - thread.start() - thread.join(timeout=2.0) - assert not thread.is_alive(), "helper failed to exit on stop_event" + # ``min_sweep_interval=0.0`` disables the production cadence floor + # (default 5 s) so tests can fire many close_idle calls back-to-back + # without waiting real time between them. The floor is exercised + # in its own dedicated test below. + thread = threading.Thread( + target=_coord_idle_cleanup_thread, + args=(mgr, timeout_sec, stop_event), + kwargs={"min_sweep_interval": 0.0}, + daemon=True, + ) + thread.start() + thread.join(timeout=2.0) + assert not thread.is_alive(), "helper failed to exit on stop_event" -def test_coord_idle_cleanup_runs_initial_sweep_before_sleep() -> None: - """The first close_idle call must happen BEFORE the first time.sleep — +def test_coord_idle_cleanup_runs_initial_sweep_before_wait() -> None: + """The first close_idle call must happen BEFORE the first wait — otherwise cold-start orphans wait one ``check_every`` interval (~30 min on default 2h timeout) for the first reap. Crucial because the - lifespan no longer does a synchronous initial sweep.""" + lifespan no longer does a synchronous initial sweep. + + Verified structurally: a single ``expected_calls=1`` run completes + in well under one ``check_every`` (here 0.04 s timeout → 0.01 s + check_every), so the initial sweep must have happened before any + real wait could have blocked it. + """ stop_event = threading.Event() mgr = _StubMgr(stop_event=stop_event, expected_calls=1) - _run_until_done(mgr, stop_event, timeout_sec=120.0) - assert mgr.sleep_calls_at_each_close == [0], "first close_idle should run before any sleep" + started = time.monotonic() + _run_until_done(mgr, stop_event, timeout_sec=0.04) + elapsed = time.monotonic() - started + assert len(mgr.calls) == 1 + # check_every = min(300.0, 0.04/4) = 0.01 s. An initial sweep + # gated behind one full wait would have taken ~0.01+ s anyway, so + # the upper bound here is "much less than one check_every plus + # process noise" — the explicit 1.0 s gives generous CI headroom + # while still asserting the test is testing the right thing. + assert elapsed < 1.0 def test_coord_idle_cleanup_calls_close_idle_each_tick() -> None: + """Heartbeat path: with no state-change events, close_idle fires + each ``check_every`` interval. Test uses a tiny timeout so the + test runs fast — the contract under test is "the loop iterates", + not the production cadence. + """ stop_event = threading.Event() mgr = _StubMgr(stop_event=stop_event, expected_calls=3) - _run_until_done(mgr, stop_event, timeout_sec=120.0) + _run_until_done(mgr, stop_event, timeout_sec=0.04) assert len(mgr.calls) == 3 - assert all(t == 120.0 for t in mgr.calls) + assert all(t == 0.04 for t in mgr.calls) def test_coord_idle_cleanup_survives_close_idle_exceptions() -> None: @@ -91,7 +142,7 @@ def test_coord_idle_cleanup_survives_close_idle_exceptions() -> None: blip would silently leak orphans forever.""" stop_event = threading.Event() mgr = _StubMgr(stop_event=stop_event, expected_calls=4, raise_after=1) - _run_until_done(mgr, stop_event, timeout_sec=120.0) + _run_until_done(mgr, stop_event, timeout_sec=0.04) # All four calls must have fired despite calls 2-4 raising. assert len(mgr.calls) == 4 @@ -102,5 +153,154 @@ def test_coord_idle_cleanup_exits_cleanly_on_stop_event() -> None: daemon-process termination.""" stop_event = threading.Event() mgr = _StubMgr(stop_event=stop_event, expected_calls=2) - _run_until_done(mgr, stop_event, timeout_sec=120.0) + _run_until_done(mgr, stop_event, timeout_sec=0.04) assert stop_event.is_set() + + +def test_state_change_wakes_close_idle_before_heartbeat() -> None: + """The event-driven path is the whole point of the refactor: a + workstream state-change must wake the cleanup sweep without + waiting one ``check_every`` interval. Tested with a long + timeout_sec so the heartbeat would NOT have fired in the test + window — the close_idle call past the initial sweep must come + from a state-change wake. + """ + stop_event = threading.Event() + mgr = _StubMgr(stop_event=stop_event, expected_calls=2) + # check_every = min(300.0, 120.0/4) = 30 s — well outside the test + # window. Any close_idle call past the initial sweep must come + # from a fire_state_change-driven wake-up. + thread = threading.Thread( + target=_coord_idle_cleanup_thread, + args=(mgr, 120.0, stop_event), + kwargs={"min_sweep_interval": 0.0}, + daemon=True, + ) + thread.start() + # Wait for the initial sweep to complete AND the thread to enter + # its first ``tick_now.wait`` (signalled here by the subscriber + # being registered + calls advancing to 1). + deadline = time.monotonic() + 1.0 + while time.monotonic() < deadline: + if mgr.subscribers_count == 1 and len(mgr.calls) >= 1: + break + time.sleep(0.01) + assert mgr.subscribers_count == 1, "thread didn't subscribe to state" + assert len(mgr.calls) == 1, "initial sweep didn't fire" + # One state-change fire wakes the first ``wait`` → close_idle runs + # again → stop_event is set (expected_calls=2) → thread exits. + mgr.fire_state_change() + thread.join(timeout=2.0) + assert not thread.is_alive(), "thread didn't exit after state-change-driven sweep" + # 2 = initial + state-change-driven. If the state change weren't + # being honoured, close_idle would have stalled on the 30 s wait + # and the thread.join would have timed out. + assert len(mgr.calls) == 2 + + +def test_subscriber_unregisters_when_thread_exits() -> None: + """The cleanup thread's state-change subscriber must be removed + when the thread exits — otherwise long-running processes that + restart their cleanup threads (admin model-CRUD path, tests) leak + subscribers and every state change fires N stale callbacks. + """ + stop_event = threading.Event() + mgr = _StubMgr(stop_event=stop_event, expected_calls=1) + _run_until_done(mgr, stop_event, timeout_sec=0.04) + assert mgr.subscribers_count == 0, "subscriber leaked past thread exit" + + +def test_state_change_during_close_idle_triggers_followup_sweep() -> None: + """A state-change fired during the initial sweep (e.g. close_idle's + own ``close()`` calls firing subscribers) must wake the next + ``tick_now.wait`` rather than being lost to the clear-before-sweep + ordering. The clear runs INSIDE the loop just before close_idle, + so a fire during the initial sweep — which precedes the loop — + arrives at an already-set event that the first wait sees set and + returns on immediately. + """ + stop_event = threading.Event() + mgr = _StubMgr(stop_event=stop_event, expected_calls=2) + + real_close_idle = mgr.close_idle + + # One-shot fire during the initial sweep, mirroring what + # close_idle's own close() calls do in production (set_state → + # state-change subscribers). + fired = [False] + + def _instrumented_close_idle(timeout_sec: float) -> list[str]: + result = real_close_idle(timeout_sec) + if not fired[0]: + fired[0] = True + mgr.fire_state_change() + return result + + mgr.close_idle = _instrumented_close_idle # type: ignore[method-assign] + + thread = threading.Thread( + target=_coord_idle_cleanup_thread, + args=(mgr, 120.0, stop_event), + kwargs={"min_sweep_interval": 0.0}, + daemon=True, + ) + thread.start() + thread.join(timeout=2.0) + assert not thread.is_alive(), "thread blocked on the next wait — mid-sweep wake was lost" + # 2 = initial sweep + state-change-driven follow-up. Without the + # event surviving the clear-before-sweep ordering, the thread + # would have blocked on the 30 s ``wait`` and the test would have + # timed out at thread.join. + assert len(mgr.calls) == 2 + + +def test_min_sweep_interval_floors_close_idle_cadence_under_sustained_wakes() -> None: + """Cadence floor: even when state-change events keep firing + ``tick_now.set()``, ``close_idle`` must not run more often than + ``min_sweep_interval`` — otherwise the loop tight-spins close_idle + at the rate of its own DB latency, doing 600-1500x more DB work + than the pre-refactor fixed-30 s cadence. + + Wires a state-change subscriber that fires another state change + from inside close_idle, so the bus would tick forever if not + floored. Asserts the elapsed-between-sweeps is at least + ``min_sweep_interval`` modulo small wall-clock noise. + """ + stop_event = threading.Event() + mgr = _StubMgr(stop_event=stop_event, expected_calls=3) + + real_close_idle = mgr.close_idle + sweep_times: list[float] = [] + + def _instrumented_close_idle(timeout_sec: float) -> list[str]: + sweep_times.append(time.monotonic()) + result = real_close_idle(timeout_sec) + # Always fire another state-change to simulate sustained + # activity (each turn fires thinking/running/attention/idle). + # If the floor were absent, the next wake would race the next + # close_idle immediately and ``sweep_times`` deltas would be + # bounded by close_idle latency (microseconds), not the floor. + mgr.fire_state_change() + return result + + mgr.close_idle = _instrumented_close_idle # type: ignore[method-assign] + + # 0.15 s floor keeps the test fast (~0.3 s total) while still + # representing a meaningful gap relative to close_idle's + # near-zero stub latency. + thread = threading.Thread( + target=_coord_idle_cleanup_thread, + args=(mgr, 120.0, stop_event), + kwargs={"min_sweep_interval": 0.15}, + daemon=True, + ) + thread.start() + thread.join(timeout=3.0) + assert not thread.is_alive(), "thread didn't exit" + assert len(sweep_times) >= 2, "fewer than two sweeps fired" + # Gap between sweep 1 (post-initial) and sweep 2 must respect + # the floor. Initial sweep at sweep_times[0] is unfloored + # (no prior sweep to compare against), so the meaningful + # assertion is on sweep_times[1] - sweep_times[0]. + gap = sweep_times[1] - sweep_times[0] + assert gap >= 0.12, f"floor breached: gap {gap:.3f}s < min_sweep_interval 0.15s" diff --git a/tests/test_coordinator_adapter.py b/tests/test_coordinator_adapter.py index 47328ce4..057c7e74 100644 --- a/tests/test_coordinator_adapter.py +++ b/tests/test_coordinator_adapter.py @@ -716,3 +716,82 @@ class TestCoordinatorAdapterDispatchChildEvent: }, ) assert recorder.enqueued == [] + + def test_dispatch_notifies_child_event_bus_on_state_event(self) -> None: + """Every translated state-class event must call + ``ChildEventBus.notify(ws_id)`` so a registered + ``wait_for_workstream`` waiter wakes promptly. Notify fires + AFTER the UI enqueue so the SSE fan-out keeps priority — the + order assertion here is structural (one notify call, matching + ws_id) since the bus side-effect lookup is what guards against + regressions, not the relative event ordering. + """ + adapter, _, _ = self._setup() + adapter._registry.merge_children("coord-a", ["child-a1"]) + bus = adapter.child_event_bus + event = bus.register_waiter(["child-a1"]) + adapter._dispatch_child_event( + { + "type": "cluster_state", + "ws_id": "child-a1", + "state": "idle", + } + ) + assert event.is_set(), "bus notify did not fire on cluster_state dispatch" + + def test_dispatch_notifies_for_all_state_class_event_types(self) -> None: + """The dispatch sink translates six event types into the + ``child_ws_*`` SSE shape; all six must also fire the bus so + a wait on any of them wakes. ``ws_created`` is intentionally + NOT in this set — waiters register against ws_ids they already + know exist (the wait tool takes a pre-known list).""" + for etype, extra in [ + ("cluster_state", {"state": "running"}), + ("ws_closed", {"reason": "evicted"}), + ("ws_rename", {"name": "renamed"}), + ("intent_verdict", {"verdict": {"call_id": "c1"}}), + ("approval_resolved", {"approved": True}), + ("approve_request", {"detail": {}}), + ]: + adapter, _, _ = self._setup() + adapter._registry.merge_children("coord-a", ["child-a1"]) + bus = adapter.child_event_bus + event = bus.register_waiter(["child-a1"]) + adapter._dispatch_child_event( + {"type": etype, "ws_id": "child-a1", **extra}, + ) + assert event.is_set(), f"bus notify did not fire on {etype} dispatch" + + def test_dispatch_does_not_notify_for_unrelated_ws_id(self) -> None: + """Bus is keyed by ws_id — a dispatch for ws X must not wake a + waiter registered against ws Y, or every state change anywhere + in the system would shake every concurrent wait.""" + adapter, _, _ = self._setup() + adapter._registry.merge_children("coord-a", ["child-a1"]) + bus = adapter.child_event_bus + event = bus.register_waiter(["child-other"]) + adapter._dispatch_child_event( + { + "type": "cluster_state", + "ws_id": "child-a1", + "state": "idle", + } + ) + assert not event.is_set(), "bus notify spuriously fired on unrelated ws_id" + + def test_dispatch_does_not_notify_for_unknown_child(self) -> None: + """Events whose ws_id isn't in any coord's registry are dropped + BEFORE the bus notify (early return at ``coord_id is None``). + Notify only fires for events the dispatch sink fully translated, + keeping the bus side-effect aligned with the UI enqueue.""" + adapter, _, _ = self._setup() + bus = adapter.child_event_bus + event = bus.register_waiter(["ws-orphan"]) + adapter._dispatch_child_event( + { + "type": "cluster_state", + "ws_id": "ws-orphan", + "state": "idle", + } + ) + assert not event.is_set(), "bus notify fired for ws_id the dispatch dropped" diff --git a/tests/test_coordinator_client.py b/tests/test_coordinator_client.py index fd57e5bb..7c96ace2 100644 --- a/tests/test_coordinator_client.py +++ b/tests/test_coordinator_client.py @@ -9,6 +9,7 @@ storage-call path. from __future__ import annotations import json +import time from typing import TYPE_CHECKING, Any import httpx @@ -20,6 +21,7 @@ from turnstone.console.coordinator_client import ( CoordinatorTokenManager, ) from turnstone.core.auth import JWT_AUD_CONSOLE, validate_jwt +from turnstone.core.child_event_bus import ChildEventBus from turnstone.core.storage._sqlite import SQLiteBackend if TYPE_CHECKING: @@ -145,6 +147,7 @@ def _mock_client( coord_ws_id="coord-1", user_id="user-1", http_client=http, + child_event_bus=ChildEventBus(), ) return client, captured @@ -470,6 +473,7 @@ def _make_read_client(storage: SQLiteBackend) -> CoordinatorClient: coord_ws_id="coord-1", user_id="user-1", http_client=http, + child_event_bus=ChildEventBus(), ) @@ -663,6 +667,7 @@ def _make_client_with_cluster_response( coord_ws_id="coord-1", user_id="user-1", http_client=http, + child_event_bus=ChildEventBus(), ) @@ -1443,6 +1448,23 @@ def test_wait_for_workstream_denies_foreign_ws_id(populated_storage): assert result["elapsed"] < 1.0 +def test_wait_for_workstream_denies_cross_tenant_child(populated_storage): + """Defense-in-depth (Copilot #506): a row whose ``parent_ws_id`` + matches the coordinator but whose ``user_id`` belongs to a + different tenant must collapse to ``denied`` — otherwise a + forged / migration-era / pre-tenant-gate row would let a + coordinator's LLM observe foreign-tenant state through + ``wait_for_workstream``. The ``populated_storage`` fixture's + ``cross-tenant-child`` row has exactly this shape + (parent_ws_id="coord-1", user_id="user-2"). + """ + client = _make_read_client(populated_storage) + result = client.wait_for_workstream(["cross-tenant-child"], timeout=5, mode="any") + assert result["results"]["cross-tenant-child"]["state"] == "denied" + assert result["complete"] is False + assert result["elapsed"] < 1.0 + + def test_wait_for_workstream_missing_ws_id_indistinguishable_from_denied(populated_storage): """A ws_id that doesn't exist collapses into the same 'denied' shape as a foreign ws_id so wait can't be used as an existence @@ -1531,10 +1553,22 @@ def test_wait_for_workstream_dedupes_ws_ids(populated_storage): assert list(result["results"].keys()) == ["child-a"] -def test_wait_for_workstream_uses_batched_storage_calls(populated_storage, monkeypatch): - """Per-tick polling must issue batched storage calls — at the - documented cap (32 ws_ids over a 600s wait) the naive per-id - shape produced ~38k row reads. Guard against regression.""" +def test_wait_for_workstream_never_falls_back_to_per_id_storage_calls( + populated_storage, monkeypatch +): + """All storage reads issued by ``wait_for_workstream`` must go + through the batched paths. At the documented cap (32 ws_ids over + a 600 s wait) the naive per-id shape produced ~38k row reads, so + a regression to per-id is the meaningful failure mode this test + guards against. + + The primary safety net is the ``pytest.fail`` mock on the per-id + ``get_workstream`` / ``sum_workstream_tokens`` paths — any call + there blows up loudly with the regression message. The + additional ``batch_calls`` / ``sum_calls`` assertions cover the + subtler regression where the call IS batched but only covers a + subset of ws_ids (e.g. one ws_id per call in a loop). + """ client = _make_read_client(populated_storage) batch_calls: list[list[str]] = [] sum_calls: list[list[str]] = [] @@ -1565,11 +1599,16 @@ def test_wait_for_workstream_uses_batched_storage_calls(populated_storage, monke result = client.wait_for_workstream(["child-a", "child-b"], timeout=5, mode="any") assert result["complete"] is True - # One tick is enough since child-a is already idle (terminal). - assert len(batch_calls) == 1 - assert len(sum_calls) == 1 - assert set(batch_calls[0]) == {"child-a", "child-b"} - assert set(sum_calls[0]) == {"child-a", "child-b"} + # Every batched call carried the full ws_id set. The exact count + # (currently 2: one pre-loop ownership filter + one snapshot tick) + # is incidental; if either gains another batched read it stays + # batched, which is the property under test. + assert batch_calls, "no batched get_workstreams_batch call observed" + assert sum_calls, "no batched sum_workstream_tokens_batch call observed" + first_batch = set(batch_calls[0]) + first_sum = set(sum_calls[0]) + assert first_batch == {"child-a", "child-b"} + assert first_sum == {"child-a", "child-b"} def test_wait_for_workstream_handles_non_string_mode(populated_storage): @@ -1581,6 +1620,205 @@ def test_wait_for_workstream_handles_non_string_mode(populated_storage): assert "invalid mode" in result["error"] +# --------------------------------------------------------------------------- +# wait_for_workstream — event-driven (ChildEventBus wired in) +# --------------------------------------------------------------------------- +# +# When the coord adapter wires its ``child_event_bus`` into the client, +# the wait loop blocks on a per-call ``threading.Event`` keyed by ws_id +# and only re-snapshots storage on state-change wakes or the heartbeat +# cap. The legacy ``time.sleep`` poll path remains intact for tests +# that don't wire the bus (above), so this section adds focused +# coverage of the bus-driven behaviour without re-running the full +# matrix of mode / since / cross-tenant cases. + + +def _make_read_client_with_bus(storage, bus) -> CoordinatorClient: + """Like ``_make_read_client`` but wires a real ``ChildEventBus``. + + Caller owns the bus so the test can call ``bus.notify(ws_id)`` to + simulate the dispatch-sink wake-up. + """ + transport = httpx.MockTransport(lambda r: httpx.Response(200)) + http = httpx.Client(transport=transport) + return CoordinatorClient( + console_base_url="http://x", + storage=storage, + token_factory=lambda: "t", + coord_ws_id="coord-1", + user_id="user-1", + http_client=http, + child_event_bus=bus, + ) + + +def test_wait_with_bus_returns_immediately_when_already_terminal(populated_storage): + """Subscribe-after-terminal race: the wait registers its waiter + BEFORE the first snapshot, then re-snapshots — an already-terminal + child must return at once without spinning the heartbeat cap. + """ + from turnstone.core.child_event_bus import ChildEventBus + + bus = ChildEventBus() + client = _make_read_client_with_bus(populated_storage, bus) + result = client.wait_for_workstream(["child-a"], timeout=5, mode="any") + assert result["complete"] is True + assert result["results"]["child-a"]["state"] == "idle" + assert result["elapsed"] < 1.0 + # Waiter must be unregistered on exit so a long-lived bus doesn't + # accumulate dead keys across many waits. + assert "child-a" not in bus._waiters + + +def test_wait_with_bus_wakes_on_notify(populated_storage): + """The core property of the refactor: a state-change ``notify`` + must wake the wait promptly — well under the legacy 0.5 s poll + cadence AND the 2 s heartbeat cap. Test fires a state update + + notify after a short delay and asserts the wait returns quickly. + """ + import threading as _t + + from turnstone.core.child_event_bus import ChildEventBus + + bus = ChildEventBus() + client = _make_read_client_with_bus(populated_storage, bus) + # child-b starts running; flip to idle + notify after the wait + # blocks. 100 ms is enough that the wait is parked in event.wait() + # but short enough that the test runs fast. + timer = _t.Timer( + 0.1, + lambda: ( + populated_storage.update_workstream_state("child-b", "idle"), + bus.notify("child-b"), + ), + ) + timer.start() + start = time.monotonic() + result = client.wait_for_workstream(["child-b"], timeout=5.0, mode="any") + elapsed = time.monotonic() - start + assert result["complete"] is True + assert result["results"]["child-b"]["state"] == "idle" + # Bus-driven wake should fire well under 1 s; legacy poll would + # take ~0.5 s but bus-driven should be ~0.1 s (the timer delay) + # plus a few ms. Generous 0.6 s budget for CI noise. + assert elapsed < 0.6, f"wake-up too slow: {elapsed}s" + + +def test_wait_with_bus_unrelated_notify_does_not_wake(populated_storage): + """A notify on a ws_id the wait isn't watching must NOT wake it — + otherwise every state change anywhere on the system would shake + every concurrent wait into a redundant storage snapshot. + """ + from turnstone.core.child_event_bus import ChildEventBus + + bus = ChildEventBus() + client = _make_read_client_with_bus(populated_storage, bus) + # child-b is running indefinitely; mode='all' will time out unless + # a relevant notify fires. Fire only unrelated notifies — wait + # should still hit the full timeout. + import threading as _t + + def _fire_unrelated() -> None: + for _ in range(5): + bus.notify("ws-unrelated-1") + bus.notify("ws-unrelated-2") + time.sleep(0.05) + + t = _t.Thread(target=_fire_unrelated, daemon=True) + t.start() + start = time.monotonic() + result = client.wait_for_workstream(["child-b"], timeout=0.5, mode="all") + elapsed = time.monotonic() - start + assert result["complete"] is False, "unrelated notify falsely satisfied wait" + # Wait should burn its full timeout (give or take heartbeat + # granularity). The bus path doesn't have a 0.5 s poll, so the + # bound is "approximately timeout". + assert elapsed >= 0.5 + t.join(timeout=1.0) + + +def test_wait_with_bus_heartbeat_still_progresses_without_notify(populated_storage): + """Without any notify, the wait must still progress through ticks + via the heartbeat cap so ``progress_callback`` keeps firing for + the sidebar UI. Verified by counting callback firings over an + interval longer than the heartbeat. + """ + from turnstone.core.child_event_bus import ChildEventBus + + bus = ChildEventBus() + client = _make_read_client_with_bus(populated_storage, bus) + # Shrink the heartbeat for test speed via the ClassVar seam — + # instance attribute shadows the class-level default. Production + # stays at 2.0 s; the test exercises the heartbeat-fires-without- + # notify property in well under 1 s. + client._WAIT_HEARTBEAT_INTERVAL = 0.1 # type: ignore[misc] + snapshots: list[dict[str, dict[str, object]]] = [] + + def _cb(snap: dict[str, dict[str, object]], _elapsed: float) -> None: + snapshots.append(snap) + + # child-b is running indefinitely; wait will time out at 0.4 s. + # With heartbeat = 0.1 s, we expect ~3-5 callback firings + # (initial tick + ~3-4 heartbeats). Loose lower bound to avoid + # CI flakiness. + start = time.monotonic() + result = client.wait_for_workstream(["child-b"], timeout=0.4, mode="all", progress_callback=_cb) + elapsed = time.monotonic() - start + assert result["complete"] is False + assert elapsed >= 0.4 + # At least 2 callback firings: the initial snapshot plus at least + # one heartbeat-driven re-tick. Tight upper bound would be + # ~ceil(0.4/0.1) + 1 = 5 firings. + assert len(snapshots) >= 2, f"heartbeat didn't fire: {len(snapshots)} snapshots" + + +def test_wait_with_bus_unregisters_waiter_on_exit(populated_storage): + """Both the success path and the timeout path must unregister the + waiter — otherwise a long-lived bus accumulates dead + ``threading.Event`` instances forever. + """ + from turnstone.core.child_event_bus import ChildEventBus + + bus = ChildEventBus() + client = _make_read_client_with_bus(populated_storage, bus) + # Success path (already-terminal child). + client.wait_for_workstream(["child-a"], timeout=5, mode="any") + assert bus._waiters == {}, "success path leaked waiter" + # Timeout path (running child, mode='all' that times out). + client.wait_for_workstream(["child-a", "child-b"], timeout=0.3, mode="all") + assert bus._waiters == {}, "timeout path leaked waiter" + + +def test_wait_with_bus_multi_waiter_independence(populated_storage): + """Two concurrent waits on the same ws_id must be independent — + one wait completing must not affect the other's wake-up state. + Smoke-tests the multi-Event-per-bucket bus behaviour against the + real wait-loop. + """ + import threading as _t + + from turnstone.core.child_event_bus import ChildEventBus + + bus = ChildEventBus() + client = _make_read_client_with_bus(populated_storage, bus) + + results: dict[str, dict[str, object]] = {} + + def _do_wait(label: str) -> None: + results[label] = client.wait_for_workstream(["child-a"], timeout=5, mode="any") + + threads = [_t.Thread(target=_do_wait, args=(f"t{i}",), daemon=True) for i in range(3)] + for t in threads: + t.start() + for t in threads: + t.join(timeout=5.0) + for label in ("t0", "t1", "t2"): + assert results[label]["complete"] is True + assert results[label]["results"]["child-a"]["state"] == "idle" + # All waiters must be unregistered after exit. + assert bus._waiters == {} + + # --------------------------------------------------------------------------- # wait_for_workstream — last-message bundling # --------------------------------------------------------------------------- diff --git a/tests/test_coordinator_end_to_end.py b/tests/test_coordinator_end_to_end.py index ba696f45..0115a433 100644 --- a/tests/test_coordinator_end_to_end.py +++ b/tests/test_coordinator_end_to_end.py @@ -40,6 +40,7 @@ from turnstone.console.server import ( _require_coord_mgr, ) from turnstone.core.auth import AuthResult +from turnstone.core.child_event_bus import ChildEventBus from turnstone.core.session_manager import SessionManager from turnstone.core.session_routes import ( SessionEndpointConfig, @@ -286,6 +287,7 @@ def test_coordinator_client_spawn_close_delete(tmp_path): coord_ws_id="coord-42", user_id="user-1", http_client=http, + child_event_bus=ChildEventBus(), ) # spawn --------------------------------------------------------------- @@ -387,6 +389,7 @@ def _read_client(storage: SQLiteBackend) -> CoordinatorClient: coord_ws_id="coord-root", user_id="user-1", http_client=http, + child_event_bus=ChildEventBus(), ) diff --git a/turnstone/console/coordinator_adapter.py b/turnstone/console/coordinator_adapter.py index 2185cbfe..d5933101 100644 --- a/turnstone/console/coordinator_adapter.py +++ b/turnstone/console/coordinator_adapter.py @@ -19,6 +19,7 @@ from typing import TYPE_CHECKING, Any from turnstone.core import session_worker from turnstone.core.adapters._ui_cleanup import cleanup_session_ui +from turnstone.core.child_event_bus import ChildEventBus from turnstone.core.child_source import ClusterChildSource from turnstone.core.children_registry import ChildrenRegistry from turnstone.core.log import get_logger @@ -67,6 +68,14 @@ class CoordinatorAdapter: # method names which still exist as thin shims for the # cluster-routing + cleanup callers. self._registry = ChildrenRegistry() + # In-process wakeup primitive for ``wait_for_workstream``. The + # dispatch sink (:meth:`_dispatch_child_event`) calls + # ``notify(child_ws_id)`` after each translated child event; + # waiters block on per-call ``threading.Event``s instead of + # polling storage. Owned by the adapter so the manager-level + # exposure can simply delegate; ``CoordinatorClient`` picks it + # up via the coord client factory closure. + self._child_event_bus = ChildEventBus() # Cross-node child events arrive via ``ClusterChildSource`` # (Stage 3 Step 2): a strategy that subscribes to the # collector's listener channel and runs a daemon thread that @@ -77,6 +86,17 @@ class CoordinatorAdapter: # the collector reference is available. self._child_source: ClusterChildSource | None = None + @property + def child_event_bus(self) -> ChildEventBus: + """In-process wakeup bus consumed by ``wait_for_workstream``. + + Exposed so the coord client factory in the console bootstrap + can pass it to :class:`CoordinatorClient` without reaching + into a private attr, and so :class:`SessionManager` can + delegate its own ``child_event_bus`` property here. + """ + return self._child_event_bus + def attach(self, manager: SessionManager) -> None: """Late-bind the owning :class:`SessionManager`. @@ -639,6 +659,13 @@ class CoordinatorAdapter: "detail": event.get("detail") or {}, } _enqueue_on_ui(owning_ws.ui, coord_id, child_event) + # Wake any in-process ``wait_for_workstream`` subscriber on + # this child. Notify runs AFTER the UI enqueue so the SSE + # fan-out keeps priority (a wait that wakes early sees the + # state already enqueued for its owning dashboard). The bus + # is a no-op when no waiter is registered — the steady + # state for the hot dispatch path. + self._child_event_bus.notify(ws_id) def _enqueue_on_ui(ui: Any, coord_ws_id: str, payload: dict[str, Any]) -> None: diff --git a/turnstone/console/coordinator_client.py b/turnstone/console/coordinator_client.py index 3cb038fb..288dc933 100644 --- a/turnstone/console/coordinator_client.py +++ b/turnstone/console/coordinator_client.py @@ -73,12 +73,27 @@ WAIT_MAX_WS_IDS: int = 32 # wait_for_workstream again with the same ws_ids — each call re-arms freshly. WAIT_MAX_TIMEOUT: float = 600.0 -# Storage-poll cadence. 500ms is short enough that the wait terminates -# promptly after a child finishes (well under the human-perceptible-latency -# floor), and long enough that a 60s wait incurs at most 120 cheap row -# reads — still cheaper than the 20+ inspect_workstream model turns the -# tool replaces. -WAIT_POLL_INTERVAL: float = 0.5 +# Maximum ``event.wait`` interval in the bus-driven wait loop. +# A long-running stuck child would otherwise look dead in the sidebar UI +# because the ``wait_progress`` SSE emission piggybacks on the wait loop +# — capping at 2 s keeps the heartbeat visible without flooding storage. +# Today's polling effectively snapshots every 500 ms; 2 s preserves a +# similar liveness feel while cutting per-listener SSE traffic ~4x in the +# steady-state-quiescent case. Tunable post-merge if profiling shows +# storage-read pressure on state-change wakes. +# +# **Worst-case completion latency**: 2 s. ``SessionManager.set_state`` +# buffers non-ERROR storage writes through ``StateWriter`` (async-flushed +# at ~1 s cadence) while ``emit_state`` fans the event out immediately — +# a bus-driven wake can therefore beat the flusher and read pre-transition +# state on a terminal transition, then re-block on ``event.wait`` until +# the heartbeat cap fires. Pre-bus the 0.5 s poll bounded this at 0.5 s. +# Going to 2 s is intentional: the 4x SSE-traffic reduction in the +# steady-state-quiescent case outweighs the worst-case latency +# regression on the most common terminal transition, and a model issuing +# a follow-up ``inspect_workstream`` (the pre-bus pattern this tool +# replaces) was already paying multi-second model-turn latency per probe. +WAIT_HEARTBEAT_INTERVAL: float = 2.0 # Per-ws cap on the inline ``message`` field bundled into wait_for_workstream # results. Sized so a fan-out of 32 children at the cap is ~320 KiB of @@ -184,6 +199,7 @@ def load_task_envelope(storage: Any, ws_id: str) -> tuple[dict[str, Any], bool]: if TYPE_CHECKING: from collections.abc import Callable + from turnstone.core.child_event_bus import ChildEventBus from turnstone.core.storage._protocol import StorageBackend log = get_logger(__name__) @@ -313,6 +329,7 @@ class CoordinatorClient: user_id: str, timeout: float = 30.0, http_client: httpx.Client | None = None, + child_event_bus: ChildEventBus, ) -> None: self._base_url = console_base_url.rstrip("/") self._storage = storage @@ -325,6 +342,12 @@ class CoordinatorClient: # with the coordinator session. self._http = http_client or httpx.Client(timeout=timeout) self._owns_http = http_client is None + # In-process wakeup bus for ``wait_for_workstream``. The wait + # loop blocks on a ``threading.Event`` keyed by ws_id and only + # re-snapshots storage on state-change wakes or the heartbeat + # cap. Owned by ``CoordinatorAdapter`` in production; tests + # pass their own instance. + self._child_event_bus = child_event_bus # tasks per-ws lock cache — populated lazily by _task_lock(). # Single-session so a plain dict behind a coarse lock is fine; # WeakValueDictionary isn't needed (entries live as long as the @@ -447,6 +470,27 @@ class CoordinatorClient: return False return bool(row.get("user_id")) and row.get("user_id") == self._user_id + def _row_in_own_subtree(self, ws_id: str, row: dict[str, Any] | None) -> bool: + """Row-level subtree predicate sharing one home for read paths. + + Both :meth:`wait_for_workstream`'s pre-loop ownership filter and + its inner ``_snapshot_all`` already have the workstream row in + hand (from ``get_workstreams_batch``). Funneling them through + the same 4-line check keeps the predicate in lockstep with + :meth:`_is_own_subtree` (used by mutating ops) — both require + ``parent_ws_id`` AND ``user_id`` parity so a corrupted or + forged ``parent_ws_id`` alone can't satisfy the gate on either + path. Returns False on a missing / None row so callers can + safely pass ``rows.get(wid)``. + """ + if ws_id == self._coord_ws_id: + return True + if row is None: + return False + if row.get("parent_ws_id") != self._coord_ws_id: + return False + return bool(row.get("user_id")) and row.get("user_id") == self._user_id + # -- model-invoked mutating ops (HTTP) --------------------------------- def spawn( @@ -562,7 +606,7 @@ class CoordinatorClient: _WAIT_TERMINAL_STATES: ClassVar[frozenset[str]] = WAIT_TERMINAL_STATES _WAIT_MAX_WS_IDS: ClassVar[int] = WAIT_MAX_WS_IDS _WAIT_MAX_TIMEOUT: ClassVar[float] = WAIT_MAX_TIMEOUT - _WAIT_POLL_INTERVAL: ClassVar[float] = WAIT_POLL_INTERVAL + _WAIT_HEARTBEAT_INTERVAL: ClassVar[float] = WAIT_HEARTBEAT_INTERVAL def wait_for_workstream( self, @@ -721,12 +765,7 @@ class CoordinatorClient: snaps: dict[str, dict[str, Any]] = {} for wid in cleaned: row = rows.get(wid) - if row is None: - snaps[wid] = {"state": "denied", "tokens": 0} - continue - is_self = wid == self._coord_ws_id - is_own_child = row.get("parent_ws_id") == self._coord_ws_id - if not (is_self or is_own_child): + if row is None or not self._row_in_own_subtree(wid, row): snaps[wid] = {"state": "denied", "tokens": 0} continue snaps[wid] = { @@ -760,54 +799,119 @@ class CoordinatorClient: last_results: dict[str, dict[str, Any]] = {} complete = False - while True: - results = _snapshot_all() - last_results = results - if progress_callback is not None: - try: - progress_callback(results, time.monotonic() - start) - except Exception: - log.debug("coord_client.wait.progress_cb_failed", exc_info=True) - real_terminal = [_is_real_terminal(snap) for snap in results.values()] - settled = [_is_settled(snap) for snap in results.values()] - # ``since`` — orthogonal to mode. If the caller supplied a - # prior snapshot, any diff on a ws_id that IS in ``since_map`` - # exits the wait so a follow-up call doesn't re-count - # already-terminal children. ws_ids absent from ``since_map`` - # are ignored for the diff-exit check — they fall through to - # the normal mode='any' / mode='all' conditions below. This - # prevents a disjoint since-dict from exiting on tick one - # with complete=True (previous shape did, silently). - if since_map and any( - _diff_since(snap, since_map[wid]) - for wid, snap in results.items() - if wid in since_map - ): - complete = True - break - if mode == "any": - if any(real_terminal): + # Subscribe to in-process state-change events for the watched + # ws_ids when the bus is wired. ``register_waiter`` returns a + # single ``threading.Event`` registered against every id so a + # wait on [A, B, C] wakes on any of A/B/C changing. Bus is + # optional so test fixtures that don't wire it fall back to the + # legacy ``time.sleep`` cadence with no behaviour change. + # + # **Defense-in-depth ownership filter**: ``_dispatch_child_event`` + # fires ``bus.notify(ws_id)`` for every ws_id in *any* coord's + # registry on this console process, so a foreign ws_id passed by + # an untrusted coord LLM (prompt injection) would otherwise leak + # wake-up timing as a side channel — _snapshot_all returns + # ``denied`` for the content, but the *time* at which the wait + # un-blocked would correlate with the foreign ws_id's next + # state-class event. Filter ``cleaned`` to own-subtree ids + # before registering; foreign / missing ws_ids stay in the + # snapshot list so they still surface as ``denied`` in + # ``_snapshot_all`` and exit via the pure-denied short-circuit + # below. Predicate shared with ``_snapshot_all`` via + # :meth:`_row_in_own_subtree`. + try: + pre_rows = self._storage.get_workstreams_batch(cleaned) + except Exception: + log.debug("coord_client.wait.ownership_filter_failed", exc_info=True) + pre_rows = {wid: None for wid in cleaned} + own_subtree = [wid for wid in cleaned if self._row_in_own_subtree(wid, pre_rows.get(wid))] + bus = self._child_event_bus + wake_event = bus.register_waiter(own_subtree) if own_subtree else None + try: + while True: + # Clear BEFORE the storage snapshot to close the + # subscribe/check race: any ``notify`` between clear + # and the next ``wake_event.wait`` leaves the Event + # set, so the wait returns immediately and the loop + # re-snapshots without losing the wake-up. + if wake_event is not None: + wake_event.clear() + results = _snapshot_all() + last_results = results + if progress_callback is not None: + try: + progress_callback(results, time.monotonic() - start) + except Exception: + log.debug("coord_client.wait.progress_cb_failed", exc_info=True) + real_terminal = [_is_real_terminal(snap) for snap in results.values()] + settled = [_is_settled(snap) for snap in results.values()] + # ``since`` — orthogonal to mode. If the caller supplied a + # prior snapshot, any diff on a ws_id that IS in ``since_map`` + # exits the wait so a follow-up call doesn't re-count + # already-terminal children. ws_ids absent from ``since_map`` + # are ignored for the diff-exit check — they fall through to + # the normal mode='any' / mode='all' conditions below. This + # prevents a disjoint since-dict from exiting on tick one + # with complete=True (previous shape did, silently). + if since_map and any( + _diff_since(snap, since_map[wid]) + for wid, snap in results.items() + if wid in since_map + ): complete = True break - # Pure-denied list: every snap is settled but none is a - # real terminal — no work to wait for. Short-circuit so - # the model sees the denied results immediately rather - # than spinning the timeout (``complete=False`` because - # the wait condition never had a real chance to fire). - if all(settled): + if mode == "any": + if any(real_terminal): + complete = True + break + # Pure-denied list: every snap is settled but none is a + # real terminal — no work to wait for. Short-circuit so + # the model sees the denied results immediately rather + # than spinning the timeout (``complete=False`` because + # the wait condition never had a real chance to fire). + if all(settled): + break + else: # mode == "all" + if all(settled): + # Every ws_id is settled (real-terminal or denied). + # The wait condition is met — the model gets the + # full results dict and decides what each terminal + # state means. + complete = True + break + remaining = deadline - time.monotonic() + if remaining <= 0: break - else: # mode == "all" - if all(settled): - # Every ws_id is settled (real-terminal or denied). - # The wait condition is met — the model gets the - # full results dict and decides what each terminal - # state means. - complete = True - break - remaining = deadline - time.monotonic() - if remaining <= 0: - break - time.sleep(min(self._WAIT_POLL_INTERVAL, remaining)) + if wake_event is not None: + # Block until a child state-change notify fires OR + # the heartbeat cap expires (so a stuck child still + # emits a periodic ``wait_progress`` for the + # sidebar UX). Heartbeat cap is the only timer — + # the bus is the wake source. See + # ``WAIT_HEARTBEAT_INTERVAL`` (module top) for the + # worst-case completion-latency rationale: 2 s is + # a deliberate 4x trade vs the pre-bus 0.5 s poll. + wake_event.wait(min(remaining, self._WAIT_HEARTBEAT_INTERVAL)) + else: + # Pure-foreign / pure-denied list: every cleaned + # ws_id was filtered out of ``own_subtree`` so the + # bus has nothing to wake on. The pure-denied + # short-circuit above exits ``mode='any'`` on the + # first tick; ``mode='all'`` falls through to here + # and must burn the timeout. Use the heartbeat + # cadence for the deadline carve-up so + # ``progress_callback`` keeps firing. + time.sleep(min(self._WAIT_HEARTBEAT_INTERVAL, remaining)) + finally: + # Always unregister so a crash mid-wait can't leak the + # registration past one wait's lifetime. Bus discards + # empty buckets so long-lived buses don't accumulate dead + # keys after many waits. Unregister against the same + # ``own_subtree`` list the register call used — passing + # ``cleaned`` here would silently no-op for foreign ids + # but pass an unknown bucket to ``unregister_waiter``. + if wake_event is not None: + bus.unregister_waiter(own_subtree, wake_event) # Bundle each terminal child's last assistant message inline so the # coordinator LLM doesn't have to follow up with one # ``inspect_workstream`` per ws. Only ``idle`` / ``error`` ws_ids @@ -816,7 +920,7 @@ class CoordinatorClient: # subset across a small thread pool — at the WAIT_MAX_WS_IDS=32 # cap, 8 workers cuts a worst-case all-idle fan-out from 32 # sequential storage round-trips down to 4 batches, which lands - # inside the WAIT_POLL_INTERVAL the model already tolerates + # inside the WAIT_HEARTBEAT_INTERVAL the model already tolerates # between ticks. Storage backends use SQLAlchemy with # ``check_same_thread=False`` (SQLite) / a connection pool # (Postgres), so concurrent reads from the worker pool are safe. diff --git a/turnstone/console/server.py b/turnstone/console/server.py index 2df534b1..9cd1f77f 100644 --- a/turnstone/console/server.py +++ b/turnstone/console/server.py @@ -4140,6 +4140,7 @@ def _coord_idle_cleanup_thread( mgr: SessionManager, timeout_sec: float, stop_event: threading.Event | None = None, + min_sweep_interval: float = 5.0, ) -> None: """Periodically reap idle + DB-orphan coordinator workstreams. @@ -4150,7 +4151,7 @@ def _coord_idle_cleanup_thread( and which aren't currently loaded. The latter pass catches coords left behind by prior console process incarnations. - Runs an initial sweep BEFORE the first sleep so cold-start orphans are + Runs an initial sweep BEFORE the first wait so cold-start orphans are reaped immediately rather than waiting one ``check_every`` interval (~30 min on default 2h timeout). This intentionally diverges from the regular server pattern, which has no initial sweep — the regular server runs @@ -4158,26 +4159,90 @@ def _coord_idle_cleanup_thread( is a small fixed-size cache where orphans dominate the row count after a cold boot. + Wait shape: subscribes a callback to ``mgr._state_subscribers`` that + sets a ``tick_now`` event; the loop blocks on ``tick_now.wait(check_every)`` + so any workstream state-change wakes the sweeper without waiting a + full check interval, AND the timeout still fires the periodic sweep + even when no activity happens (catching the DB-orphan-only case). + Net: blocked most of the time instead of repeating storage scans. + + ``min_sweep_interval`` is the hard floor between successive + ``close_idle`` calls (default 5 s) — without it, sustained + state-change activity (each turn typically fires + thinking/running/attention/idle on the coord SessionManager) would + cause every ``tick_now.set`` mid-sweep to leave the next ``wait`` + returning immediately, and the loop would tight-spin ``close_idle`` + at the rate of its own DB latency (~20-50 calls/sec). The floor + bounds DB-call traffic at ``1 / min_sweep_interval`` per second + under any external activity while still letting a quiet system + fire on every state-change wake-up. Tests inject 0.0 to keep the + suite fast. + + Default 5 s is a 6x improvement on the pre-refactor fixed 30 s + cadence while bounding DB-call traffic at ~0.2 calls/sec under + sustained activity — an order of magnitude below ``close_idle``'s + DB-latency budget, but tight enough that idle-row reaping still + feels prompt to a human watching the sidebar. Tunable post-merge + if profiling shows close_idle latency dominates the cadence. + ``stop_event`` is for tests — when set, the thread exits cleanly after the next loop check. Production callers pass ``None`` (the daemon is process-lifetime). """ check_every = min(300.0, timeout_sec / 4) - # Initial sweep — runs once before entering the sleep loop. + tick_now = threading.Event() + + def _on_state_change(_ws_id: str, _state: Any) -> None: + # Any workstream state-change resets the idle clock for that + # ws AND may make a different ws newly-eligible (close-idle + # pass 2 evaluates DB rows by timestamp). Cheap signal, full + # re-evaluation deferred to the next loop iteration. + tick_now.set() + + mgr.subscribe_to_state(_on_state_change) try: - mgr.close_idle(timeout_sec) - except Exception: - log.debug("console.coord_idle_cleanup_initial_failed", exc_info=True) - while True: - if stop_event is not None and stop_event.is_set(): - return - time.sleep(check_every) - if stop_event is not None and stop_event.is_set(): - return + # Initial sweep — runs once before entering the wait loop. + # ``tick_now`` is intentionally not cleared here: any + # state-change event that arrives between subscribe and the + # first ``wait`` should fire close_idle immediately, not be + # discarded. try: mgr.close_idle(timeout_sec) except Exception: - log.debug("console.coord_idle_cleanup_failed", exc_info=True) + log.debug("console.coord_idle_cleanup_initial_failed", exc_info=True) + last_sweep_at = time.monotonic() + while True: + if stop_event is not None and stop_event.is_set(): + return + tick_now.wait(check_every) + if stop_event is not None and stop_event.is_set(): + return + # Clear BEFORE the cadence floor so any state-change event + # arriving during the cooldown (or during the close_idle + # below) leaves ``tick_now`` set — the next loop iteration + # then re-enters ``wait`` already-set and re-evaluates + # promptly. close_idle is idempotent so a spurious extra + # tick is just one redundant scan. + tick_now.clear() + # Cadence floor — see docstring for the tight-spin + # hazard rationale. Cooldown uses ``stop_event.wait`` + # (not ``time.sleep``) so the test stop hook still + # terminates promptly during the cooldown window. + since_last = time.monotonic() - last_sweep_at + if since_last < min_sweep_interval: + gap = min_sweep_interval - since_last + if stop_event is not None: + if stop_event.wait(gap): + return + else: + time.sleep(gap) + try: + mgr.close_idle(timeout_sec) + except Exception: + log.debug("console.coord_idle_cleanup_failed", exc_info=True) + last_sweep_at = time.monotonic() + finally: + mgr.unsubscribe_from_state(_on_state_change) # Guards concurrent attempts to bootstrap the coord subsystem from the @@ -4239,12 +4304,22 @@ def _bootstrap_coord_subsystem( def _token_factory() -> str: return tm.token + # ``coord_adapter`` is bound later in this same + # ``_bootstrap_coord_subsystem`` call, after the adapter and + # manager are constructed but before any session is created + # — so this factory is *defined* before the adapter exists but + # only ever *called* after it does. The free-variable lookup + # at call time resolves to the adapter built in this same + # bootstrap pass, giving the client a handle to the in-process + # wakeup bus the dispatch sink notifies on every child + # state-change event. return CoordinatorClient( console_base_url=console_bind_url, storage=storage, token_factory=_token_factory, coord_ws_id=ws_id, user_id=user_id, + child_event_bus=coord_adapter.child_event_bus, ) # Pre-compute config-derived integers BEFORE any thread starts so a diff --git a/turnstone/core/child_event_bus.py b/turnstone/core/child_event_bus.py new file mode 100644 index 00000000..a19d7c7b --- /dev/null +++ b/turnstone/core/child_event_bus.py @@ -0,0 +1,121 @@ +"""Per-workstream wakeup primitive for in-process child state-change subscribers. + +Retires the polling pattern in ``CoordinatorClient.wait_for_workstream``, +where the coord LLM's wait tool issued a storage snapshot every 0.5 s +regardless of whether anything had changed. The dispatch path +(:meth:`turnstone.console.coordinator_adapter.CoordinatorAdapter._dispatch_child_event`) +now calls :meth:`ChildEventBus.notify` after each translated child event; +waiters block on a per-call :class:`threading.Event` returned by +:meth:`register_waiter` and re-read storage only when an event fires or +the heartbeat cap expires. + +Bus is in-process only. Cross-process / cross-node child events are +already merged into ``_dispatch_child_event`` via the cluster collector's +SSE multiplex before the bus sees them — there is no locality branching +in the bus itself. + +Design constraints: + +- Waiter primitive is :class:`threading.Event` because the wait tool runs + on the coordinator's sync worker thread, not an asyncio loop. +- Concurrent ``register`` / ``unregister`` / ``notify`` is safe — a + single ``threading.Lock`` guards the dict. ``Event.set`` itself is + thread-safe and is called outside the lock so a slow waker can't block + registration. +- ``notify`` with no subscribers is a no-op (the steady state — most + state-change events fire while no wait tool is active). +- A waiter watching multiple ws_ids fires once on any of them; the + caller's ``_snapshot_all`` re-read resolves which one changed. +- Empty per-ws_id buckets are popped on unregister so long-lived buses + don't accumulate dead keys after wait churn. +""" + +from __future__ import annotations + +import threading +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Iterable + + +class ChildEventBus: + """Fan ``notify(child_ws_id)`` to every :class:`threading.Event` + registered against that ws_id. + + Use :meth:`register_waiter` once per wait call to obtain an Event, + then call :meth:`unregister_waiter` in a ``finally`` so a crash mid- + wait doesn't leak the registration. The dispatch side calls + :meth:`notify` on every translated child state-change event; an + empty bucket is a cheap dict lookup + immediate return. + """ + + def __init__(self) -> None: + self._waiters: dict[str, set[threading.Event]] = {} + self._lock = threading.Lock() + + def register_waiter(self, child_ws_ids: Iterable[str]) -> threading.Event: + """Return a fresh Event registered against every listed ws_id. + + A wait on ``[A, B, C]`` returns a single Event that fires when + *any* of A/B/C changes. The caller's snapshot re-read resolves + which one. Empty / falsy ids are silently skipped — callers that + clean their input upstream (e.g. ``wait_for_workstream``'s + dedup + cap) don't need to filter again here. + """ + event = threading.Event() + with self._lock: + for wid in child_ws_ids: + if not wid: + continue + self._waiters.setdefault(wid, set()).add(event) + return event + + def unregister_waiter( + self, + child_ws_ids: Iterable[str], + event: threading.Event, + ) -> None: + """Remove ``event`` from each listed ws_id's waiter set. + + Idempotent — already-removed Events silently no-op. Pops empty + sets so a long-lived bus doesn't accumulate dead keys after + many waits have come and gone. Must be called from the same + ``finally`` that paired with :meth:`register_waiter` so a + crash mid-wait doesn't leak the registration past one wait's + lifetime. + """ + with self._lock: + for wid in child_ws_ids: + if not wid: + continue + bucket = self._waiters.get(wid) + if bucket is None: + continue + bucket.discard(event) + if not bucket: + self._waiters.pop(wid, None) + + def notify(self, child_ws_id: str) -> None: + """Wake every Event registered for ``child_ws_id``. + + Called from the coord dispatch sink after each translated child + event. Snapshot the bucket under the lock, then call + ``Event.set`` outside the lock so a slow waker doesn't block + ``register`` / ``unregister`` / further ``notify``. ``Event.set`` + is thread-safe and idempotent — re-firing a still-set Event is + a no-op. + + Empty / falsy ws_ids are silently dropped; the same hot path + runs for every dispatched event regardless of whether anyone's + waiting, so the empty-bucket case must stay cheap. + """ + if not child_ws_id: + return + with self._lock: + bucket = self._waiters.get(child_ws_id) + if not bucket: + return + events = list(bucket) + for event in events: + event.set() diff --git a/turnstone/core/session_manager.py b/turnstone/core/session_manager.py index 5ad0b1e9..042c047a 100644 --- a/turnstone/core/session_manager.py +++ b/turnstone/core/session_manager.py @@ -22,6 +22,7 @@ from turnstone.core.workstream import Workstream, WorkstreamKind, WorkstreamStat if TYPE_CHECKING: from collections.abc import Callable + from turnstone.core.child_event_bus import ChildEventBus from turnstone.core.session import ChatSession, SessionUI from turnstone.core.state_writer import StateWriter from turnstone.core.storage._protocol import StorageBackend @@ -252,6 +253,19 @@ class SessionManager: def kind(self) -> WorkstreamKind: return self._adapter.kind + @property + def child_event_bus(self) -> ChildEventBus | None: + """Delegate to the adapter's per-workstream wakeup bus. + + Returns ``None`` for adapters that don't host one (today only the + coord adapter does; interactive's child surface is degenerate + and has nothing to wait on yet). Manager-level property gives + adapter-agnostic callers (tests, future cross-kind tools) a + stable lookup that doesn't depend on knowing which adapter is + attached. + """ + return getattr(self._adapter, "child_event_bus", None) + @property def _service_type(self) -> str | None: """``services.service_type`` this manager's hosting process registers