From 4d230cd013702a995b17e200afc26cdc0e1f866b Mon Sep 17 00:00:00 2001 From: Patrick Buckley Date: Tue, 21 Jul 2026 18:19:39 -0700 Subject: [PATCH] feat(session): coalesce concurrent /history reconstructions per workstream (#884) After a node restart every open pane resyncs via REST /history inside the same jitter window; client jitter spreads the peak but not the total. Concurrent requests for the same (ws_id, limit) now share ONE reconstruction (load_messages -> decoration -> projection) via a single-flight task map in the handler closure. Deliberately single-flight only, no TTL cache: the payload depends on live-mutable inputs with no total cheap invalidation signal (the surface_persisted_reasoning registry toggle emits no per-ws event; cold workstreams have no event counter), so a cache could serve stale reasoning/approval/cursor state for its whole TTL, while a joiner's worst-case staleness equals the flight duration -- the window a lone slow request already exposes. All auth/tenant/kind/existence gates stay per-request ahead of the join; only the caller-independent reconstruction is shared. A shared draw that hit a transient load_messages failure is not fanned out: joiners retry once, independently, so one storage blip cannot wipe every coalesced pane (the 200-empty payload renders as an authoritative empty pane in both clients, and the seedless clear_ui path has no SSE redelivery to repair it). The flight is a detached task (awaiters shield it) so an owner disconnect cannot strand joiners, and each task pops its own key in a finally, so the map only ever holds in-flight work. ws.history.load_failed rises to warning: it now names the draw that triggers joiner retries and renders as a pane wipe. --- tests/test_workstream_endpoints.py | 308 ++++++++++++++++++++++++++++- turnstone/core/session_routes.py | 152 +++++++++++++- 2 files changed, 450 insertions(+), 10 deletions(-) diff --git a/tests/test_workstream_endpoints.py b/tests/test_workstream_endpoints.py index 1b795ee43..41d66fecf 100644 --- a/tests/test_workstream_endpoints.py +++ b/tests/test_workstream_endpoints.py @@ -2,12 +2,16 @@ from __future__ import annotations +import asyncio import json +import logging import queue +import threading from types import SimpleNamespace from typing import TYPE_CHECKING, Any from unittest.mock import MagicMock, patch +import httpx import pytest from starlette.applications import Starlette from starlette.middleware import Middleware @@ -17,6 +21,8 @@ from starlette.testclient import TestClient if TYPE_CHECKING: + from collections.abc import Callable + from starlette.requests import Request from starlette.responses import Response @@ -954,11 +960,18 @@ def _interactive_endpoint_cfg( ) -def _build_history_app( +def _build_history_asgi_app( mock_mgr: Any, storage: Any, tenant_check: Any = None, -) -> TestClient: +) -> Starlette: + """Raw ASGI app for the lifted history factory. + + Returned un-wrapped (no ``TestClient``) so the coalescing tests can + drive CONCURRENT requests through ``httpx.ASGITransport`` on a + private event loop — ``TestClient`` is synchronous and can only + hold one request in flight at a time. + """ cfg = _interactive_endpoint_cfg(mock_mgr, tenant_check=tenant_check) handler = make_history_handler(cfg) app = Starlette( @@ -974,7 +987,15 @@ def _build_history_app( ) app.state.workstreams = mock_mgr app.state.auth_storage = storage - return TestClient(app) + return app + + +def _build_history_app( + mock_mgr: Any, + storage: Any, + tenant_check: Any = None, +) -> TestClient: + return TestClient(_build_history_asgi_app(mock_mgr, storage, tenant_check=tenant_check)) def _build_detail_app( @@ -1485,6 +1506,287 @@ def test_history_does_not_synthesize_orphan_results(self, _inject_storage): assert all(m.get("role") != "tool" for m in r.json()["messages"]) +class _GatedStorage: + """Delegates to the real backend; ``load_messages`` counts entries, + optionally raises, and blocks on ``gate`` until the test releases it + (a pre-set gate is a pass-through). The count-then-block ordering + is load-bearing for the coalescing tests: a request that did NOT + join an existing flight bumps ``load_calls`` at entry — before the + gate can block it — so the counter distinguishes join from + second-flight without any timing assumptions.""" + + def __init__(self, inner: Any) -> None: + self._inner = inner + self.gate = threading.Event() + self.load_calls = 0 + self.fail_next = 0 + self._lock = threading.Lock() + + def __getattr__(self, name: str) -> Any: + return getattr(self._inner, name) + + def load_messages(self, ws_id: str, **kwargs: Any) -> Any: + with self._lock: + self.load_calls += 1 + should_fail = self.fail_next > 0 + if should_fail: + self.fail_next -= 1 + if not self.gate.wait(timeout=10): + raise TimeoutError("test gate never released") + if should_fail: + raise RuntimeError("transient load failure (test)") + return self._inner.load_messages(ws_id, **kwargs) + + +async def _until(cond: Callable[[], bool], timeout: float = 5.0) -> None: + async with asyncio.timeout(timeout): + while not cond(): + await asyncio.sleep(0.01) + + +class TestHistoryCoalescing: + """Single-flight coalescing of concurrent ``/history`` requests (#884). + + The matrix these tests assert: join-vs-miss × gates-per-request × + failed-vs-clean shared draw × key isolation × no cross-flight + caching × owner cancellation. Driven through ``httpx.ASGITransport`` + on a private loop because ``TestClient`` cannot hold two requests in + flight at once. + + Determinism scheme: every synchronization point is a positive + observable edge — the owner's arrival on ``load_calls == 1`` (it + entered ``load_messages`` and parked on the gate), and the JOIN on + the handler's ``ws.history.coalesced`` debug record, which is + emitted synchronously at the map-hit branch before the joiner + awaits the flight (capture through stdlib handlers is verified: + turnstone's structlog config routes to stdlib logging, so pytest's + ``caplog`` sees it). The join-proof stays counter-based on top: + a request that did NOT join bumps ``load_calls`` before the gate + can block it (see ``_GatedStorage``), so ``load_calls == 1`` after + the join record proves sharing.""" + + def _register_ws(self, storage: Any, ws_id: str) -> None: + storage.register_workstream(ws_id, kind="interactive", user_id="test-user") + storage.save_message(ws_id, "user", "hello") + storage.save_message(ws_id, "assistant", "hi there") + + def _live_mgr(self, ws_id: str) -> MagicMock: + mock_ws = MagicMock() + mock_ws.id = ws_id + mock_ws.ui._pending_approval = None + mock_mgr = MagicMock() + mock_mgr.get.return_value = mock_ws + return mock_mgr + + def _counting_tenant(self) -> tuple[Any, dict[str, int]]: + calls = {"n": 0} + + def check(request: Any, ws_id: str, mgr: Any) -> Any: + calls["n"] += 1 + if request.headers.get("x-test-deny"): + return JSONResponse({"error": "Workstream not found"}, status_code=404) + return None + + return check, calls + + def _scaffold( + self, storage: Any, ws_id: str + ) -> tuple[_GatedStorage, dict[str, int], Starlette, str]: + self._register_ws(storage, ws_id) + gated = _GatedStorage(storage) + tenant, tenant_calls = self._counting_tenant() + app = _build_history_asgi_app(self._live_mgr(ws_id), gated, tenant_check=tenant) + return gated, tenant_calls, app, f"/v1/api/workstreams/{ws_id}/history" + + async def _drive_two_joiners( + self, + app: Starlette, + url: str, + gated: _GatedStorage, + caplog: Any, + *, + cancel_owner: bool = False, + ) -> tuple[Any, httpx.Response]: + """Drive two concurrent requests through one shared flight. + + Every wait is a positive observable edge. Owner arrival: + ``load_calls == 1`` (it entered ``load_messages`` and parked on + the gate). JOIN: the handler logs ``ws.history.coalesced`` + synchronously at the map-hit branch, before the joiner awaits + the flight — waiting on that record (via ``caplog``; capture + verified, structlog routes to stdlib logging) proves t2 joined + while the flight is still gated, with no wall-clock beat. The + trailing ``"ws="`` in the match keeps it from also matching + ``ws.history.coalesced_retry``. The counter join-proof stays + on top: a request that did NOT join bumps ``load_calls`` at + ``load_messages`` entry, BEFORE the gate can block it (see + ``_GatedStorage``). + + ``cancel_owner=True`` swaps the clean tail for the + owner-disconnect tail (cancel t1 before releasing the gate, + await the joiner, collect t1's outcome) — on that path the + first tuple member is the owner's ``CancelledError``, not a + ``Response``, hence the loose first-slot type. + """ + caplog.set_level(logging.DEBUG, logger="turnstone.core.session_routes") + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://t") as c: + t1 = asyncio.create_task(c.get(url)) + await _until(lambda: gated.load_calls == 1) + t2 = asyncio.create_task(c.get(url)) + await _until( + lambda: any("ws.history.coalesced ws=" in r.getMessage() for r in caplog.records) + ) + assert gated.load_calls == 1 + if cancel_owner: + t1.cancel() + gated.gate.set() + r2 = await t2 + (r1,) = await asyncio.gather(t1, return_exceptions=True) + return r1, r2 + gated.gate.set() + r1, r2 = await asyncio.gather(t1, t2) + return r1, r2 + + def test_concurrent_requests_share_one_flight(self, _inject_storage: Any, caplog: Any) -> None: + gated, _tenant_calls, app, url = self._scaffold(_inject_storage, "ws-flight-share") + + r1, r2 = asyncio.run(self._drive_two_joiners(app, url, gated, caplog)) + assert r1.status_code == 200 + assert r2.status_code == 200 + assert r1.json() == r2.json() + contents = [m.get("content") for m in r1.json()["messages"]] + assert contents == ["hello", "hi there"] + assert gated.load_calls == 1 + + def test_failed_shared_draw_not_fanned_out_to_joiner( + self, _inject_storage: Any, caplog: Any + ) -> None: + """A transient ``load_messages`` failure in the shared draw + yields the owner's 200-empty (same as a lone request today) but + the JOINER retries independently and gets the real rows — one + storage blip must not wipe every coalesced pane.""" + gated, _tenant_calls, app, url = self._scaffold(_inject_storage, "ws-flight-fail") + gated.fail_next = 1 + + r1, r2 = asyncio.run(self._drive_two_joiners(app, url, gated, caplog)) + assert r1.status_code == 200 + assert r1.json()["messages"] == [] + assert r2.status_code == 200 + assert [m.get("content") for m in r2.json()["messages"]] == ["hello", "hi there"] + # Owner draw + joiner retry — never a third. + assert gated.load_calls == 2 + + def test_gates_run_per_request_before_join(self, _inject_storage: Any) -> None: + """A caller failing its own tenant gate gets its 404 while the + flight is still in the air — it never joins and never triggers + a reconstruction of its own.""" + gated, _tenant_calls, app, url = self._scaffold(_inject_storage, "ws-flight-deny") + + async def run() -> tuple[httpx.Response, httpx.Response]: + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://t") as c: + t1 = asyncio.create_task(c.get(url)) + await _until(lambda: gated.load_calls == 1) + r2 = await c.get(url, headers={"x-test-deny": "1"}) + assert gated.load_calls == 1 + gated.gate.set() + r1 = await t1 + return r1, r2 + + r1, r2 = asyncio.run(run()) + assert r1.status_code == 200 + assert r2.status_code == 404 + assert gated.load_calls == 1 + + def test_flight_key_includes_limit(self, _inject_storage: Any) -> None: + """A limit=1 caller must not receive the limit-100 payload: + different limits are different flights.""" + gated, _tenant_calls, app, url = self._scaffold(_inject_storage, "ws-flight-limit") + + async def run() -> tuple[httpx.Response, httpx.Response]: + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://t") as c: + t1 = asyncio.create_task(c.get(url)) + await _until(lambda: gated.load_calls == 1) + t2 = asyncio.create_task(c.get(url, params={"limit": 1})) + # Positively waits for the SECOND reconstruction to + # enter load_messages while the first is still gated — + # the limit=1 request did not join. + await _until(lambda: gated.load_calls == 2) + gated.gate.set() + r1, r2 = await asyncio.gather(t1, t2) + return r1, r2 + + r1, r2 = asyncio.run(run()) + assert len(r1.json()["messages"]) == 2 + assert len(r2.json()["messages"]) == 1 + assert gated.load_calls == 2 + + def test_no_result_reuse_across_completed_flights(self, _inject_storage: Any) -> None: + """The flights map is a single-flight, not a cache: a request + arriving after a flight completed reconstructs afresh.""" + gated, _tenant_calls, app, url = self._scaffold(_inject_storage, "ws-flight-seq") + gated.gate.set() # pass-through + + async def run() -> tuple[httpx.Response, httpx.Response]: + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://t") as c: + r1 = await c.get(url) + r2 = await c.get(url) + return r1, r2 + + r1, r2 = asyncio.run(run()) + assert r1.status_code == 200 + assert r2.status_code == 200 + assert r1.json() == r2.json() + assert gated.load_calls == 2 + + def test_decoration_failure_is_shared_not_retried( + self, _inject_storage: Any, caplog: Any + ) -> None: + """The inverse of the load-failure test: a decoration/projection + failure AFTER a successful load is degraded-but-non-empty, so it + is shared to joiners as-is — ``load_failed`` stays False and no + joiner retry fires (``load_calls`` stays 1, vs the load-failure + test's 2). Pins the two failure modes apart: conflating them + (e.g. setting ``load_failed`` in the decoration except) would + turn every shared degraded draw into a double reconstruction.""" + gated, _tenant_calls, app, url = self._scaffold(_inject_storage, "ws-flight-decor") + + with patch( + "turnstone.core.history_decoration.project_history_messages", + side_effect=RuntimeError("projection failure (test)"), + ) as fake_project: + r1, r2 = asyncio.run(self._drive_two_joiners(app, url, gated, caplog)) + # The degraded path must actually have run — without this, a + # fixture whose cursor is None either way would let an + # ineffective patch pass the assertions below vacuously. + assert fake_project.called + assert r1.status_code == 200 + assert r2.status_code == 200 + # Degraded (un-projected, cursor=None) but NON-empty and identical — + # the joiner shared the draw instead of re-reconstructing. + assert r1.json() == r2.json() + assert r1.json()["messages"] + assert r1.json()["cursor"] is None + assert gated.load_calls == 1 + + def test_owner_disconnect_leaves_joiner_completing( + self, _inject_storage: Any, caplog: Any + ) -> None: + """The flight is a detached task: cancelling the request that + CREATED it (client disconnect) must not cancel the shared + reconstruction a joiner is awaiting.""" + gated, _tenant_calls, app, url = self._scaffold(_inject_storage, "ws-flight-cancel") + + r1, r2 = asyncio.run(self._drive_two_joiners(app, url, gated, caplog, cancel_owner=True)) + assert isinstance(r1, asyncio.CancelledError) + assert r2.status_code == 200 + assert [m.get("content") for m in r2.json()["messages"]] == ["hello", "hi there"] + assert gated.load_calls == 1 + + class TestBuildHistorySystemTurnPropagation: """``project_history_messages`` surfaces first-class operator-context ``system`` rows (``_source`` → ``source``) so a tab reconnecting via diff --git a/turnstone/core/session_routes.py b/turnstone/core/session_routes.py index 283dd40e1..aba89d75a 100644 --- a/turnstone/core/session_routes.py +++ b/turnstone/core/session_routes.py @@ -31,9 +31,10 @@ from __future__ import annotations +import asyncio from collections.abc import Awaitable, Callable, Iterable from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Protocol, cast +from typing import TYPE_CHECKING, Any, Protocol, TypeAlias, cast from starlette.responses import JSONResponse from starlette.routing import Route @@ -3395,6 +3396,13 @@ def _resume_cursor_and_trim( return messages[:orphan_idx], cursor +# One /history flight's result: (messages, resume cursor, load_failed). +# ``load_failed`` is True only when ``load_messages`` RAISED — never for +# a legitimately empty workstream; see ``_reconstruct`` inside +# :func:`make_history_handler`. +_HistoryFlightResult: TypeAlias = tuple[list[dict[str, Any]], int | None, bool] + + def make_history_handler(cfg: SessionEndpointConfig) -> Handler: """Lifted body for ``GET {prefix}/{ws_id}/history`` — message history. @@ -3429,13 +3437,38 @@ def make_history_handler(cfg: SessionEndpointConfig) -> Handler: lifted verbs' storage offload pattern; pre-lift coord ran them inline on the event loop). + Restart-herd coalescing (issue #884): concurrent requests for the + same ``(ws_id, limit)`` share ONE reconstruction (single-flight) — + after a node restart every open pane resyncs via REST ``/history`` + inside the same jitter window, and each un-coalesced request repeats + the full ``load_messages`` → decoration → projection pipeline + against a cold process. Deliberately single-flight only, NO + result/TTL cache: the payload depends on live-mutable inputs with no + total cheap invalidation signal (the ``surface_persisted_reasoning`` + registry toggle emits no per-ws event; cold workstreams have no + event counter at all), so a cached payload could serve stale + reasoning/approval/cursor state for the whole TTL, while a joiner's + worst-case staleness equals the flight duration — the same window a + lone slow request already exposes. All auth/existence gates run + per-request BEFORE a request may join a flight; only the + caller-independent reconstruction is shared. + Args: cfg: per-kind policy bundle. """ - async def history(request: Request) -> Response: - import asyncio + # In-flight reconstructions, keyed ``(ws_id, limit)``. Holds ONLY + # live flights — each task pops its own key in a ``finally`` before + # completing, so a later request can never read a completed (stale) + # result: this map is a single-flight, not a cache. Scoped to this + # cfg's closure (one map per mount) so interactive and coordinator + # ws_ids can never collide across kinds, mirroring the cfg-level + # isolation of every other lifted verb. Touched only from the event + # loop thread — no lock needed; the ``limit`` component is required + # (a limit=10 caller must not receive a limit=500 payload). + flights: dict[tuple[str, int], asyncio.Task[_HistoryFlightResult]] = {} + async def history(request: Request) -> Response: if cfg.permission_gate is not None: err = cfg.permission_gate(request) if err is not None: @@ -3509,6 +3542,105 @@ async def history(request: Request) -> Response: limit = 100 limit = max(1, min(limit, 500)) + # ---- single-flight join (issue #884) -------------------------- + # Everything ABOVE this line is per-request (auth, tenant, kind, + # existence, limit) and must stay above it: a request may only + # join a flight after its own gates passed. Everything below is + # the caller-independent reconstruction, shared via ``flights``. + key = (ws_id, limit) + task = flights.get(key) + joined = task is not None + if task is None: + task = asyncio.create_task(_run_flight(key, mgr, storage, request.app.state)) + flights[key] = task + else: + # The coalescing tests synchronize on this record — keep the + # ``ws.history.coalesced ws=`` prefix stable (a rename turns + # their join-wait into a confusing 5s timeout, not a clean + # failure). + log.debug("ws.history.coalesced ws=%s", ws_id[:8]) + # ``shield``: this request disconnecting must not cancel the + # shared reconstruction other joiners are awaiting — the flight + # is a detached task that completes (and pops its key) on its + # own. A cancelled REQUEST still propagates CancelledError out + # of the shield to Starlette as usual. Accepted cost: a LONE + # reader's mid-flight disconnect no longer aborts the + # reconstruction (pre-coalescing, the inline pipeline was + # cancellable at its next await) — the detached flight runs its + # bounded pipeline to completion once, unobserved. Reviewed + # and kept as-is: refcounting awaiters to cancel-at-zero would + # add a join-vs-cancel race (a late joiner grabbing a task the + # last leaver is cancelling) to a seam that is race-free + # precisely because the flight is never cancelled. + messages, cursor, load_failed = await asyncio.shield(task) + if joined and load_failed: + # The shared draw hit a transient ``load_messages`` failure + # and produced the 200-empty payload. Both clients render a + # 200-empty as an authoritative empty pane (the seedless + # clear_ui path has no SSE redelivery to repair it), so + # sharing the failed draw would fan ONE storage blip out as + # a pane wipe across every joiner. Joiners therefore retry + # once, independently and unshared — exactly the blast + # radius the un-coalesced endpoint had. The flight OWNER + # keeps its failed draw (same as a lone request today). + log.debug("ws.history.coalesced_retry ws=%s", ws_id[:8]) + messages, cursor, _ = await _reconstruct(mgr, storage, request.app.state, ws_id, limit) + # Each awaiter serializes its own JSONResponse from the shared + # payload — cost parity with the un-coalesced endpoint (one + # render per request). Sharing pre-rendered bytes across + # awaiters was reviewed and declined: it reshapes the flight + # result to (bytes, flag), splitting the _HistoryFlightResult + # contract the retry path still needs, for a herd-only + # serialization micro-win. + return JSONResponse({"ws_id": ws_id, "messages": messages, "cursor": cursor}) + + async def _run_flight( + key: tuple[str, int], + mgr: SessionManager, + storage: Any, + app_state: Any, + ) -> _HistoryFlightResult: + # The key MUST be popped when (and only when) this flight + # settles — success, internal exception, or task cancellation — + # and from inside the task itself, so ``flights`` never holds a + # completed task: a finished (stale) or poisoned flight can + # never be joined late and replay an old result or exception. + try: + return await _reconstruct(mgr, storage, app_state, key[0], key[1]) + finally: + flights.pop(key, None) + + async def _reconstruct( + mgr: SessionManager, + storage: Any, + app_state: Any, + ws_id: str, + limit: int, + ) -> _HistoryFlightResult: + """The shareable reconstruction: rows → decorated, projected payload. + + Runs once per flight and is awaited by every coalesced request, + so it must not read anything request- or caller-specific — the + auth/tenant/kind gates all ran per-request before the flight + was joined. Resolves its own ``mgr.get(ws_id)`` handle (rather + than inheriting the flight owner's) so the live reads below + (``_pending_approval``, the replay ring, agent trajectories) + anchor to flight start; a joiner can observe live state up to + one flight-duration old — the same window a single slow request + already exposes. A stale resume ``cursor`` handed to a joiner + degrades gracefully: if the ring evicts past it before that + client connects, the stream answers ``replay_truncated`` and + the client resyncs. + + Returns ``(messages, cursor, load_failed)`` — ``load_failed`` + is True only when ``load_messages`` RAISED (transient storage + failure), never for a legitimately empty workstream. The + caller uses it to keep an exception-empty payload from fanning + a pane wipe out to coalesced joiners (see the joiner retry in + ``history``). + """ + live_session = mgr.get(ws_id) + load_failed = False messages: list[dict[str, Any]] = [] # Fresh-connect resume cursor (the ``Last-Event-ID`` the client # opens its initial SSE with). Non-None only when the trailing @@ -3532,7 +3664,13 @@ async def history(request: Request) -> Response: include_compaction=True, ) except Exception: - log.debug("ws.history.load_failed ws=%s", ws_id[:8], exc_info=True) + # Warning, not debug: this 200-empty renders as an + # authoritative pane wipe in both clients, and under + # coalescing it is also what triggers the joiner retry — + # a storage blip here is operationally interesting for + # the same reason ``decoration_failed`` below is. + load_failed = True + log.warning("ws.history.load_failed ws=%s", ws_id[:8], exc_info=True) # Audit-trail decoration — attach persisted intent_verdict and # output_assessment data to each assistant.tool_calls entry so # the dashboard's history replay paints the same verdict pills @@ -3597,8 +3735,8 @@ async def history(request: Request) -> Response: # ``app.state.registry``; console stores its coord # registry as ``app.state.coord_registry``. The # lifted handler is shared, so we try both. - resolved_registry = getattr(request.app.state, "registry", None) or getattr( - request.app.state, "coord_registry", None + resolved_registry = getattr(app_state, "registry", None) or getattr( + app_state, "coord_registry", None ) if resolved_registry is not None and resolved_alias: try: @@ -3688,7 +3826,7 @@ async def history(request: Request) -> Response: ws_id[:8], exc_info=True, ) - return JSONResponse({"ws_id": ws_id, "messages": messages, "cursor": cursor}) + return messages, cursor, load_failed return history