From 4e791cfb1597ef9cc02219ebe84ce16c919c6fc1 Mon Sep 17 00:00:00 2001 From: Patrick Buckley Date: Fri, 24 Apr 2026 17:29:29 -0700 Subject: [PATCH] refactor(server): swap interactive workers to session_worker.send MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Five spawn sites in turnstone/server.py now share the worker dispatch: * ``send_message`` (``POST /v1/api/send``) — the main path. Now uses ``session_worker.send`` with separate ``_enqueue`` / ``_run`` closures; the queue-vs-spawn outcome is conveyed via a captured ``queue_outcome`` dict so the existing response shapes (``status: queued`` vs ``status: ok``) survive. * ``_make_watch_dispatch`` — watch results dispatch. * ``run_retry`` (post-rewind) and ``_run_initial`` (initial-message on workstream creation) — set ``_worker_running`` directly under ws._lock instead of going through session_worker (their structural shape doesn't fit a queue-vs-spawn decision) but stay consistent with the shared gate so they can't race with /send into parallel workers. * ``cancel_generation``'s ``was_running`` snapshot now reads ``_worker_running`` for parity with the dispatcher. Pre-dispatch cancel-await also gates on ``_worker_running`` for consistency. The ``busy_error`` / ``status: busy`` legacy branch (reached only when worker is alive but ws.session is None) is gone — the new path checks ws.session up front and returns the same 500 shape. Test fixtures in test_server_attachments_endpoints.py and test_watch_dispatch.py updated to set ws._worker_running explicitly (MagicMock auto-truthifies the field, which would otherwise mis-route all idle paths into queue mode). --- tests/test_server_attachments_endpoints.py | 7 +- tests/test_watch_dispatch.py | 35 +-- turnstone/server.py | 285 ++++++++++++--------- 3 files changed, 180 insertions(+), 147 deletions(-) diff --git a/tests/test_server_attachments_endpoints.py b/tests/test_server_attachments_endpoints.py index 956b8266..13280f43 100644 --- a/tests/test_server_attachments_endpoints.py +++ b/tests/test_server_attachments_endpoints.py @@ -452,6 +452,7 @@ class TestSendMessageAttachments: ws.ui = ui ws.session = session ws.worker_thread = None + ws._worker_running = False ws._lock = threading.RLock() mgr.get.return_value = ws return captured, session @@ -628,7 +629,7 @@ class TestQueuedSendWithAttachments: ui._ws_messages = 0 ui._ws_turn_tool_calls = 0 - # worker_thread needs .is_alive() → True to hit the queue branch + # _worker_running=True forces session_worker.send onto the queue path worker = MagicMock() worker.is_alive = MagicMock(return_value=True) @@ -638,6 +639,7 @@ class TestQueuedSendWithAttachments: ws.ui = ui ws.session = session ws.worker_thread = worker + ws._worker_running = True ws._lock = threading.RLock() mgr.get.return_value = ws return captured @@ -760,6 +762,7 @@ class TestQueuedAttachmentReservation: session.send = fake_send # type: ignore[method-assign] ws = mgr.get.return_value ws.worker_thread = None # idle → non-queue path + ws._worker_running = False # Auto-consume on a follow-up send: reserved attachment must not # be picked up (another turn isn't entitled to it). @@ -791,6 +794,7 @@ class TestQueuedAttachmentReservation: session.send = fake_send # type: ignore[method-assign] ws = mgr.get.return_value ws.worker_thread = None + ws._worker_running = False # A second send explicitly naming the reserved id: scope check # rejects it, so the attachment list is empty. @@ -885,6 +889,7 @@ class TestReserveThenDispatchRace: ws.ui = ui ws.session = session ws.worker_thread = None + ws._worker_running = False ws._lock = threading.RLock() mgr.get.return_value = ws diff --git a/tests/test_watch_dispatch.py b/tests/test_watch_dispatch.py index ef6d294b..c4c5120c 100644 --- a/tests/test_watch_dispatch.py +++ b/tests/test_watch_dispatch.py @@ -128,20 +128,15 @@ def test_busy_workstream_enqueues_message(): ws = Workstream() ui = _RecordingUI() - # Simulate a live worker thread. - blocker = threading.Event() - ws.worker_thread = threading.Thread(target=blocker.wait, args=(5,), daemon=True) - ws.worker_thread.start() + # Simulate a live worker — session_worker.send gates on + # ``_worker_running``, not ``Thread.is_alive``. + ws._worker_running = True - try: - dispatch = _make_watch_dispatch(ws, session, ui) - dispatch("queued msg") + dispatch = _make_watch_dispatch(ws, session, ui) + dispatch("queued msg") - item = session._watch_pending.get_nowait() - assert item == {"message": "queued msg"} - finally: - blocker.set() - ws.worker_thread.join(2) + item = session._watch_pending.get_nowait() + assert item == {"message": "queued msg"} def test_busy_workstream_drops_on_full_queue(): @@ -154,19 +149,13 @@ def test_busy_workstream_drops_on_full_queue(): ws = Workstream() ui = _RecordingUI() - blocker = threading.Event() - ws.worker_thread = threading.Thread(target=blocker.wait, args=(5,), daemon=True) - ws.worker_thread.start() + ws._worker_running = True # simulate a live worker - try: - dispatch = _make_watch_dispatch(ws, session, ui) - # Should not block or raise — just log a warning and drop. - dispatch("overflow msg") + dispatch = _make_watch_dispatch(ws, session, ui) + # Should not block or raise — just log a warning and drop. + dispatch("overflow msg") - assert session._watch_pending.full() - finally: - blocker.set() - ws.worker_thread.join(2) + assert session._watch_pending.full() # ── Lock guard ─────────────────────────────────────────────────────────────── diff --git a/turnstone/server.py b/turnstone/server.py index 1ac033d1..6e7efcb4 100644 --- a/turnstone/server.py +++ b/turnstone/server.py @@ -42,6 +42,7 @@ from starlette.staticfiles import StaticFiles from turnstone import __version__ from turnstone.api.docs import make_docs_handler, make_openapi_handler from turnstone.api.server_spec import build_server_spec +from turnstone.core import session_worker from turnstone.core.adapters.interactive_adapter import InteractiveAdapter from turnstone.core.auth import ( DENY_EMPTY_SUB, @@ -1435,43 +1436,48 @@ async def metrics_endpoint(request: Request) -> Response: def _make_watch_dispatch(ws: Workstream, session: ChatSession, ui: Any) -> Any: """Create a dispatch function for watch results on a workstream. - Handles both idle (start worker thread) and busy (enqueue for IDLE drain) - cases. Mirrors the ``send_message`` worker-thread pattern. + Handles both busy (enqueue into ``session._watch_pending`` for drain at IDLE) + and idle (spawn a worker thread) cases via the shared + :func:`turnstone.core.session_worker.send` dispatcher. The shared + dispatcher's ``_worker_running`` gate keeps watches and ``/v1/api/send`` + from racing into parallel workers on the same ChatSession. """ pending = session._watch_pending def dispatch(msg: str) -> None: - with ws._lock: - if ws.worker_thread and ws.worker_thread.is_alive(): - # Workstream is busy — queue for drain at IDLE (Path A) - try: - pending.put_nowait({"message": msg}) - except queue.Full: - log.warning( - "Watch pending queue full, dropping result for ws %s", - ws.id, - ) - return + def _enq() -> None: + # Watches don't pump through ChatSession.queue_message; instead + # they drop into the IDLE-drain pending queue. Swallow Full + # locally — drop-on-full is the long-standing watch behavior + # (no 429 surface). + try: + pending.put_nowait({"message": msg}) + except queue.Full: + log.warning( + "Watch pending queue full, dropping result for ws %s", + ws.id, + ) - # Workstream is idle — start a worker thread (Path B) - # Mirrors the send_message() run() pattern for proper cleanup. - def run() -> None: - me = threading.current_thread() - try: - session.send(msg) - except GenerationCancelled: - if ws.worker_thread is me and ui: - ui.on_stream_end() - ui.on_state_change("idle") - except Exception as exc: - if ws.worker_thread is me and ui: - ui.on_error(f"Watch error: {exc}") - ui.on_stream_end() - ui.on_state_change("error") + def _run() -> None: + me = threading.current_thread() + try: + session.send(msg) + except GenerationCancelled: + if ws.worker_thread is me and ui: + ui.on_stream_end() + ui.on_state_change("idle") + except Exception as exc: + if ws.worker_thread is me and ui: + ui.on_error(f"Watch error: {exc}") + ui.on_stream_end() + ui.on_state_change("error") - t = threading.Thread(target=run, daemon=True) - ws.worker_thread = t - t.start() + session_worker.send( + ws, + enqueue=_enq, + run=_run, + thread_name=f"watch-worker-{ws.id[:8]}", + ) return dispatch @@ -1634,105 +1640,108 @@ async def send_message(request: Request) -> JSONResponse: _unreserve(send_id, ws_id, attach_user_id) - # Atomically check-and-start to prevent two concurrent workers on the - # same session (ChatSession.send() is not thread-safe). # If cancel was requested, poll briefly for the worker to exit before - # rejecting. Snapshot the thread ref since force-cancel can set it to - # None concurrently. Uses async sleep to avoid blocking the event loop. - worker = ws.worker_thread - if worker and worker.is_alive() and ws.session and ws.session._cancel_event.is_set(): + # dispatching. ``_worker_running`` flips to False atomically under + # ws._lock when the worker thread reaches its finally block — gates + # the dispatch on the same flag the shared dispatcher uses. Uses + # async sleep to avoid blocking the event loop. + if ws._worker_running and ws.session and ws.session._cancel_event.is_set(): for _ in range(30): # up to 3s in 100ms steps await asyncio.sleep(0.1) - if not worker.is_alive(): + if not ws._worker_running: break - with ws._lock: - if ws.worker_thread and ws.worker_thread.is_alive(): - # Queue the message for injection at the next tool-result seam - # instead of rejecting outright. Attachments were already - # reserved above using ``send_id`` as the token — we pass - # the same id in as ``queue_msg_id`` so the queue entry, the - # reservation, and the eventual consume all share one token. - if ws.session is not None: - try: - cleaned, priority, msg_id = ws.session.queue_message( - message, - attachment_ids=list(ordered_reserved), - queue_msg_id=send_id, - ) - except queue.Full: - _release_reservation_on_fail() - return JSONResponse({"status": "queue_full"}) - ui._enqueue( - { - "type": "message_queued", - "message": cleaned, - "priority": priority, - "msg_id": msg_id, - } - ) - # Report the reservation outcome so the UI can clear - # only the chips that actually got attached, leaving - # un-reserved ones visible for retry. - dropped = [aid for aid in requested_ids if aid not in reserved_set] - return JSONResponse( - { - "status": "queued", - "priority": priority, - "msg_id": msg_id, - "attached_ids": list(ordered_reserved), - "dropped_attachment_ids": dropped, - } - ) - _release_reservation_on_fail() - ui._enqueue( - { - "type": "busy_error", - "message": "Already processing a request. Please wait.", - } + if ws.session is None: + _release_reservation_on_fail() + return JSONResponse({"error": "No session"}, status_code=500) + + session = ws.session + + # Pre-allocate the queue-path outcome dict; populated by ``_enqueue`` + # only when the shared dispatcher routes us onto a live worker. + queue_outcome: dict[str, Any] = {} + + def _enqueue() -> None: + # Reuse path: append to the live worker's pending queue. The + # ``send_id`` token threads through ``queue_msg_id`` so the + # queue entry, the attachment reservation, and the eventual + # consume all share one token. + cleaned, priority, msg_id = session.queue_message( + message, + attachment_ids=list(ordered_reserved), + queue_msg_id=send_id, + ) + queue_outcome["cleaned"] = cleaned + queue_outcome["priority"] = priority + queue_outcome["msg_id"] = msg_id + + def _run() -> None: + assert ui is not None + me = threading.current_thread() + try: + session.send( + message, + attachments=resolved_atts or None, + send_id=send_id, ) - return JSONResponse({"status": "busy"}) - session = ws.session - if session is None: + except GenerationCancelled: + # Safety net — send() normally handles this internally. + # If this thread was force-abandoned, ws.worker_thread will + # have been set to None — don't emit spurious events. _release_reservation_on_fail() - return JSONResponse({"error": "No session"}, status_code=500) + if ws.worker_thread is me: + ui.on_stream_end() + ui.on_state_change("idle") + except Exception as e: + # Release the reservation so the attachments don't stay + # soft-locked forever when the worker crashes before + # reaching the consume step. Safe-by-idempotency: once + # mark_attachments_consumed has cleared the token, a + # follow-up unreserve is a no-op. + _release_reservation_on_fail() + if ws.worker_thread is me: + ui.on_error(f"Error: {e}") + ui.on_stream_end() + ui.on_state_change("error") - def run() -> None: - assert ui is not None - me = threading.current_thread() - try: - session.send( - message, - attachments=resolved_atts or None, - send_id=send_id, - ) - except GenerationCancelled: - # Safety net — send() normally handles this internally. - # If this thread was force-abandoned, ws.worker_thread will - # have been set to None — don't emit spurious events. - _release_reservation_on_fail() - if ws.worker_thread is me: - ui.on_stream_end() - ui.on_state_change("idle") - except Exception as e: - # Release the reservation so the attachments don't stay - # soft-locked forever when the worker crashes before - # reaching the consume step. Safe-by-idempotency: once - # mark_attachments_consumed has cleared the token, a - # follow-up unreserve is a no-op. - _release_reservation_on_fail() - if ws.worker_thread is me: - ui.on_error(f"Error: {e}") - ui.on_stream_end() - ui.on_state_change("error") + ok = session_worker.send( + ws, + enqueue=_enqueue, + run=_run, + thread_name=f"send-worker-{ws.id[:8]}", + ) + if not ok: + # Returns False on queue.Full (live worker, queue at capacity) + # or session-disappeared race. Either way the message couldn't + # land — surface as queue_full so the client can retry. + _release_reservation_on_fail() + return JSONResponse({"status": "queue_full"}) - t = threading.Thread(target=run, daemon=True) - ws.worker_thread = t - t.start() + dropped = [aid for aid in requested_ids if aid not in reserved_set] + if queue_outcome: + # Reused a live worker; queue_message succeeded. + ui._enqueue( + { + "type": "message_queued", + "message": queue_outcome["cleaned"], + "priority": queue_outcome["priority"], + "msg_id": queue_outcome["msg_id"], + } + ) + return JSONResponse( + { + "status": "queued", + "priority": queue_outcome["priority"], + "msg_id": queue_outcome["msg_id"], + "attached_ids": list(ordered_reserved), + "dropped_attachment_ids": dropped, + } + ) + + # Spawned a fresh worker — count this as a new turn. _metrics.record_message_sent() with ui._ws_lock: ui._ws_messages += 1 ui._ws_turn_tool_calls = 0 - dropped = [aid for aid in requested_ids if aid not in reserved_set] return JSONResponse( { "status": "ok", @@ -1790,7 +1799,7 @@ async def cancel_generation(request: Request) -> JSONResponse: if session is None: return JSONResponse({"error": "No session"}, status_code=400) force = body.get("force", False) is True - was_running = bool(ws.worker_thread and ws.worker_thread.is_alive()) + was_running = ws._worker_running dropped = _capture_cancel_forensics(session, ui, was_running=was_running) # Only act if generation is actually in progress if was_running: @@ -1964,11 +1973,18 @@ async def command(request: Request) -> JSONResponse: ui.on_error(f"Error: {exc}") ui.on_stream_end() ui.on_state_change("error") + finally: + with ws._lock: + ws._worker_running = False + # Gate on ``_worker_running`` for parity with the shared + # session_worker dispatcher — avoids racing a /v1/api/send + # spawn with the retry spawn into two parallel workers. with ws._lock: - if ws.worker_thread and ws.worker_thread.is_alive(): + if ws._worker_running: ui.on_error("Cannot retry: workstream is busy") else: + ws._worker_running = True t = threading.Thread(target=run_retry, daemon=True) ws.worker_thread = t t.start() @@ -2679,10 +2695,18 @@ async def create_workstream(request: Request) -> JSONResponse: _fire_notify_targets(ws, last_content) except Exception: log.warning("notify_completion.hook_error", ws_id=ws.id, exc_info=True) + with ws._lock: + ws._worker_running = False - t = threading.Thread(target=_run_initial, daemon=True, name=f"ws-init-{ws.id[:8]}") - ws.worker_thread = t - t.start() + # Mark the worker live under ws._lock so a /v1/api/send + # arriving immediately after creation observes the running + # state via the shared session_worker gate instead of + # racing into a parallel worker. + with ws._lock: + ws._worker_running = True + t = threading.Thread(target=_run_initial, daemon=True, name=f"ws-init-{ws.id[:8]}") + ws.worker_thread = t + t.start() return JSONResponse( { @@ -4166,6 +4190,10 @@ async def _lifespan(app: Starlette) -> AsyncGenerator[None, None]: # Start watch runner (periodic command polling) if app.state.watch_runner: app.state.watch_runner.start() + # Start the buffered state-writer flusher + state_writer = getattr(app.state, "state_writer", None) + if state_writer is not None: + state_writer.start() # Sweep stale attachment reservations left over from process crashes # between reserve_attachments and consume/unreserve. Run once at @@ -4304,6 +4332,10 @@ async def _lifespan(app: Starlette) -> AsyncGenerator[None, None]: await tls_client.stop_renewal() if app.state.watch_runner: app.state.watch_runner.stop() + # Drain + stop the buffered state-writer + state_writer = getattr(app.state, "state_writer", None) + if state_writer is not None: + state_writer.shutdown() # Stop the orphan-reservation sweep loop _orphan_sweep_stop.set() with contextlib.suppress(asyncio.CancelledError, Exception): @@ -4361,6 +4393,7 @@ def create_app( judge_config: Any = None, config_store: Any = None, advertise_url: str = "", + state_writer: Any = None, ) -> Starlette: """Create and configure the Starlette ASGI application.""" _spec = build_server_spec() @@ -4473,6 +4506,7 @@ def create_app( lifespan=_lifespan, ) app.state.workstreams = workstreams + app.state.state_writer = state_writer app.state.global_queue = global_queue app.state.global_listeners = global_listeners app.state.global_listeners_lock = global_listeners_lock @@ -4930,11 +4964,15 @@ def main() -> None: ), session_factory=session_factory, ) + from turnstone.core.state_writer import StateWriter + + state_writer = StateWriter(_get_storage()) manager = SessionManager( interactive_adapter, storage=_get_storage(), max_active=config_store.get("server.max_workstreams"), node_id=_node_id, + state_writer=state_writer, ) interactive_adapter.attach(manager) WebUI._workstream_mgr = manager @@ -5042,6 +5080,7 @@ def main() -> None: judge_config=judge_config, config_store=config_store, advertise_url=_advertise_url, + state_writer=state_writer, ) # Wire app ref so health callbacks can access app.state for metrics