From 19d00e3586168ac4265fa602831792d876b896d1 Mon Sep 17 00:00:00 2001 From: nesquena-hermes Date: Sat, 4 Jul 2026 21:25:16 +0000 Subject: [PATCH] fix(#5543): fire-and-forget commit_session_memory + drain-safe managed stop Move the 1-5s commit_session_memory() extraction off the POST /api/session/new request thread (was blocking + New Chat) into a fire-and-forget daemon thread, guarded by the existing per-session in-flight serialization. Data-loss prevention on managed stop: - server.py installs a SIGTERM handler that requests an orderly stop so serve_forever()'s finally block runs drain_all_on_shutdown() (the default SIGTERM disposition terminated without unwinding, losing in-flight commits). - session_lifecycle: bounded self-unregistering background-commit registry; a _draining flag so no NEW worker starts once the drain begins (the inline generation-drain commits any late arrival instead); a 30s overall drain deadline so a stuck worker cannot hang the stop path forever. Co-authored-by: luperrypf --- api/routes.py | 58 ++- api/session_lifecycle.py | 110 +++++- server.py | 30 ++ ...test_issue5543_background_memory_commit.py | 358 ++++++++++++++++++ 4 files changed, 539 insertions(+), 17 deletions(-) create mode 100644 tests/test_issue5543_background_memory_commit.py diff --git a/api/routes.py b/api/routes.py index 1881956c..e1b258af 100644 --- a/api/routes.py +++ b/api/routes.py @@ -13067,21 +13067,49 @@ def handle_post(handler, parsed) -> bool: # the new session (#5420). prev_session_id = None if prev_session_id: - try: - from api.session_lifecycle import commit_session_memory - from api.config import SESSION_AGENT_CACHE, SESSION_AGENT_CACHE_LOCK - prev_agent = None - with SESSION_AGENT_CACHE_LOCK: - _cached = SESSION_AGENT_CACHE.get(prev_session_id) - if _cached: - prev_agent = _cached[0] - commit_session_memory(prev_session_id, agent=prev_agent) - except Exception: - logger.debug( - "Lifecycle commit for prev_session %s failed", - prev_session_id, - exc_info=True, - ) + # Fire-and-forget: commit_memory_session() can take 1-5+ seconds + # (extraction call to the memory provider), and blocking the + # response here made "+ New Chat" feel slow/unresponsive. + # commit_session_memory() already serialises overlapping commits + # for a session via its own in-flight guard, so running it off + # the request thread is safe. + def _commit_prev_session_memory(_sid=prev_session_id): + try: + from api.session_lifecycle import commit_session_memory + from api.config import SESSION_AGENT_CACHE, SESSION_AGENT_CACHE_LOCK + prev_agent = None + with SESSION_AGENT_CACHE_LOCK: + _cached = SESSION_AGENT_CACHE.get(_sid) + if _cached: + prev_agent = _cached[0] + commit_session_memory(_sid, agent=prev_agent) + except Exception: + logger.warning( + "Lifecycle commit for prev_session %s failed", + _sid, + exc_info=True, + ) + finally: + # Self-unregister so the background-commit registry does + # not leak completed threads; drain only tracks live ones. + try: + from api.session_lifecycle import _unregister_background_commit_thread + _unregister_background_commit_thread(threading.current_thread()) + except Exception: + pass + + t = threading.Thread( + target=_commit_prev_session_memory, + daemon=True, + name=f"commit-memory-{prev_session_id}", + ) + from api.session_lifecycle import _register_background_commit_thread + # Refused only if shutdown draining has already begun; in that + # window the inline drain commits the pending generation instead, + # so skipping the worker start is safe (avoids a late daemon + # thread the drain snapshot already missed). + if _register_background_commit_thread(t): + t.start() s = new_session( workspace=workspace, model=model, diff --git a/api/session_lifecycle.py b/api/session_lifecycle.py index c33c32ee..275a9ef2 100644 --- a/api/session_lifecycle.py +++ b/api/session_lifecycle.py @@ -45,6 +45,56 @@ _condition = threading.Condition(_lock) _sessions: dict[str, dict] = {} +# Background commit threads spawned by fire-and-forget paths (e.g. POST /api/session/new). +# Tracked so drain_all_on_shutdown() can join them before interpreter teardown, +# preventing silent data loss when daemon threads are abandoned mid-commit. +_background_commit_threads: set[threading.Thread] = set() +_background_commit_threads_lock = threading.Lock() +# Set once drain_all_on_shutdown() begins so no NEW background worker is started +# after the drain has snapshotted the registry. A late arrival is not lost: its +# uncommitted generation is still committed inline by the generation-drain loop +# below (which re-snapshots _sessions each pass and commits with wait=True). +_draining = False + + +def _register_background_commit_thread(t: threading.Thread) -> bool: + """Register a fire-and-forget commit thread. Returns True if the caller + should start it. Returns False when shutdown draining has already begun — + the caller must NOT start a new worker in that window; the inline + generation-drain will commit the pending work instead.""" + with _background_commit_threads_lock: + if _draining: + return False + _background_commit_threads.add(t) + return True + + +def _unregister_background_commit_thread(t: threading.Thread) -> None: + # A completed worker removes itself so the registry does not grow without + # bound over the process lifetime; drain only needs threads still running. + with _background_commit_threads_lock: + _background_commit_threads.discard(t) + + +def _drain_background_commit_threads(timeout: float = 5.0, deadline_fn=None) -> None: + """Join tracked fire-and-forget commit threads. When ``deadline_fn`` is given + (returns the remaining overall shutdown budget in seconds), each join is + clamped to that budget and the phase stops early once it is exhausted, so the + join phase cannot overrun ``drain_all_on_shutdown``'s deadline.""" + with _background_commit_threads_lock: + threads = list(_background_commit_threads) + _background_commit_threads.clear() + for t in threads: + if not t.is_alive(): + continue + if deadline_fn is None: + t.join(timeout) + continue + remaining = deadline_fn() + if remaining <= 0: + return + t.join(min(timeout, remaining)) + def _new_entry() -> dict: return { @@ -225,16 +275,72 @@ def commit_session_memory(session_id: str, agent=None, *, wait: bool = False, ti return True -def drain_all_on_shutdown() -> None: +def drain_all_on_shutdown(deadline: float = 30.0) -> None: + # Signal that draining has begun so no NEW background commit worker is + # started after this point (see _register_background_commit_thread). Any + # already-registered worker is joined below; a caller that loses the race + # (registered after our snapshot, or refused a start) still has its pending + # generation committed by the inline loop. + global _draining + with _background_commit_threads_lock: + _draining = True + + # Overall wall-clock budget so a stuck worker holding in_flight cannot hang + # the shutdown path forever. Start the clock BEFORE the join phase so the + # join and the inline drain share one budget. The inline commit runs in a + # daemon worker joined only for the remaining budget — otherwise a provider + # wedged inside commit_memory_session() (whether we are waiting on an + # existing in_flight commit, OR the inline drain owns the segment and calls + # the provider directly) would block indefinitely and the deadline check + # would never run again. + started = time.monotonic() + + def _remaining() -> float: + return deadline - (time.monotonic() - started) + + def _commit_within_budget(session_id: str, budget: float) -> bool: + """Run one commit in a daemon worker, join at most `budget` seconds. + Returns True only if the commit completed within budget AND made + progress. On timeout the daemon worker is left running (it will finish + or die at interpreter exit; either way it mutates generation state only + under _condition, so nothing is corrupted) and we report no progress so + the drain can honor its deadline instead of hanging on a wedged + provider.""" + holder = {"ok": False} + + def _run(): + holder["ok"] = commit_session_memory(session_id, wait=True) + + worker = threading.Thread(target=_run, daemon=True, name=f"drain-commit-{session_id}") + worker.start() + worker.join(budget) + return holder["ok"] if not worker.is_alive() else False + + # Drain in-flight background commit threads first (fire-and-forget from + # POST /api/session/new) so their commit work completes before we drain any + # remaining uncommitted generations inline. Each join consumes the shared + # budget so it cannot overrun the overall deadline. + _drain_background_commit_threads(deadline_fn=_remaining) + while True: with _lock: snapshot = [sid for sid, entry in _sessions.items() if entry["generation"] > entry["committed_generation"]] if not snapshot: return + remaining = _remaining() + if remaining <= 0: + logger.warning("drain_all_on_shutdown: deadline hit with uncommitted sessions: %s", sorted(snapshot)) + return made_progress = False for sid in snapshot: - if commit_session_memory(sid, wait=True): + budget = _remaining() + if budget <= 0: + logger.warning("drain_all_on_shutdown: deadline hit with uncommitted sessions: %s", sorted(snapshot)) + return + # Budget-bounded (covers both the wait-for-in_flight path AND the + # unbounded provider call when the inline drain owns the segment). + if _commit_within_budget(sid, budget): made_progress = True if not made_progress: logger.debug("drain_all_on_shutdown: stopped with uncommitted sessions: %s", sorted(snapshot)) diff --git a/server.py b/server.py index 29a8c220..a74141c6 100644 --- a/server.py +++ b/server.py @@ -688,6 +688,36 @@ def main() -> None: print(f' Remote access: ssh -N -L {PORT}:127.0.0.1:{PORT} @', flush=True) print(f' Then open: {scheme}://localhost:{PORT}', flush=True) print('', flush=True) + + # ctl.sh stops the WebUI with SIGTERM. Python's default SIGTERM handler + # terminates the process WITHOUT unwinding the try/finally around + # serve_forever(), so drain_all_on_shutdown() (which flushes in-flight + # fire-and-forget memory commits) would never run on the normal managed + # stop. Install a handler that requests an orderly shutdown so + # serve_forever() returns and the existing `finally` block drains cleanly. + # + # httpd.shutdown() blocks until serve_forever() has exited and MUST NOT be + # called from the thread running serve_forever() (it would deadlock), so we + # dispatch it from a short-lived helper thread. The handler is idempotent + # and guards against double-shutdown (e.g. repeated SIGTERM/SIGINT). + _shutdown_requested = threading.Event() + + def _request_shutdown(signum, _frame): + if _shutdown_requested.is_set(): + return + _shutdown_requested.set() + threading.Thread( + target=httpd.shutdown, + name="webui-sigterm-shutdown", + daemon=True, + ).start() + + try: + signal.signal(signal.SIGTERM, _request_shutdown) + except (ValueError, OSError): + # Not on the main thread (e.g. embedded/test harness); skip handler. + logger.debug("Could not install SIGTERM handler", exc_info=True) + try: httpd.serve_forever() finally: diff --git a/tests/test_issue5543_background_memory_commit.py b/tests/test_issue5543_background_memory_commit.py new file mode 100644 index 00000000..9bb99766 --- /dev/null +++ b/tests/test_issue5543_background_memory_commit.py @@ -0,0 +1,358 @@ +"""Tests for issue #5543: fire-and-forget commit_session_memory + shutdown drain. + +PR #5543 (@luperrypf) moves commit_session_memory() off the POST +/api/session/new request thread into a daemon "background commit" thread so +"+ New Chat" returns immediately instead of blocking 1-5s on memory-provider +extraction. The gate-certifier required three guarantees this file locks in: + +1. Fire-and-forget: spawning the background commit does not block the caller, + even while the underlying commit is still running. +2. Registry hygiene: a completed background-commit worker unregisters itself so + ``_background_commit_threads`` does not grow without bound (Finding 2). +3. Shutdown drain: drain_all_on_shutdown() joins in-flight background commit + threads AND flushes the pending generation, and the server's SIGTERM handler + routes the managed ``ctl.sh stop`` through the serve_forever() ``finally`` so + the drain actually runs (Finding 1). +""" + +from __future__ import annotations + +import ast +import importlib +import socket +import threading +import time +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path + + +REPO = Path(__file__).resolve().parents[1] + + +def _fresh_lifecycle(): + """Import/reload lifecycle module and clear process-global test state.""" + lifecycle = importlib.import_module("api.session_lifecycle") + lifecycle = importlib.reload(lifecycle) + reset = getattr(lifecycle, "_reset_for_tests", None) + if callable(reset): + reset() + # Ensure the background-commit registry starts empty for each test. + with lifecycle._background_commit_threads_lock: + lifecycle._background_commit_threads.clear() + return lifecycle + + +class RecordingAgent: + """Agent whose commit blocks on an event so we can hold it "in flight".""" + + def __init__(self): + self.calls = 0 + self.entered = threading.Event() + self.release = threading.Event() + + def commit_memory_session(self): + self.calls += 1 + self.entered.set() + # Block until the test releases us (bounded so a buggy test can't hang). + self.release.wait(timeout=5) + + +def _make_background_worker(lifecycle, sid, agent): + """Mirror the routes.py fire-and-forget worker: commit then self-unregister.""" + + def _worker(): + try: + lifecycle.commit_session_memory(sid, agent=agent) + finally: + lifecycle._unregister_background_commit_thread(threading.current_thread()) + + t = threading.Thread(target=_worker, daemon=True, name=f"commit-memory-{sid}") + lifecycle._register_background_commit_thread(t) + return t + + +# --------------------------------------------------------------------------- # +# Finding 2: registry hygiene # +# --------------------------------------------------------------------------- # + + +def test_register_and_unregister_roundtrip(): + lifecycle = _fresh_lifecycle() + t = threading.Thread(target=lambda: None) + + # Registration returns True (caller should start) when not draining. + assert lifecycle._register_background_commit_thread(t) is True + with lifecycle._background_commit_threads_lock: + assert t in lifecycle._background_commit_threads + + lifecycle._unregister_background_commit_thread(t) + with lifecycle._background_commit_threads_lock: + assert t not in lifecycle._background_commit_threads + + # Idempotent: unregistering an absent thread is a harmless no-op. + lifecycle._unregister_background_commit_thread(t) + + +def test_register_refused_once_draining(): + """After drain begins, registration is refused (returns False) and the + thread is NOT added — the caller must skip start(); the inline generation + drain commits the pending work instead. This closes Codex's + register-before-start shutdown window.""" + lifecycle = _fresh_lifecycle() + # Nothing pending, so drain_all_on_shutdown returns quickly but still sets + # the _draining flag. + lifecycle.drain_all_on_shutdown() + + t = threading.Thread(target=lambda: None) + assert lifecycle._register_background_commit_thread(t) is False + with lifecycle._background_commit_threads_lock: + assert t not in lifecycle._background_commit_threads + + +def test_completed_worker_self_unregisters(): + """A finished background-commit worker must drop itself from the registry.""" + lifecycle = _fresh_lifecycle() + agent = RecordingAgent() + agent.release.set() # commit completes immediately + sid = "self-unregister" + lifecycle.register_agent(sid, agent) + lifecycle.mark_turn_completed(sid, agent=agent) + + t = _make_background_worker(lifecycle, sid, agent) + t.start() + t.join(timeout=5) + assert not t.is_alive() + + with lifecycle._background_commit_threads_lock: + assert t not in lifecycle._background_commit_threads + assert len(lifecycle._background_commit_threads) == 0 + assert agent.calls == 1 + + +# --------------------------------------------------------------------------- # +# Fire-and-forget: the caller does not block # +# --------------------------------------------------------------------------- # + + +def test_fire_and_forget_does_not_block_caller(): + """Starting the background commit returns immediately, even mid-commit.""" + lifecycle = _fresh_lifecycle() + agent = RecordingAgent() # blocks in commit until released + sid = "non-blocking" + lifecycle.register_agent(sid, agent) + lifecycle.mark_turn_completed(sid, agent=agent) + + t = _make_background_worker(lifecycle, sid, agent) + + start = time.monotonic() + t.start() + # Caller "returns" as soon as the thread is started; measure that dispatch. + elapsed = time.monotonic() - start + + # The commit is genuinely in flight (proves we didn't run it inline). + assert agent.entered.wait(timeout=2) + assert elapsed < 1.0, f"caller blocked for {elapsed:.3f}s" + assert lifecycle.has_uncommitted_work(sid) is True + + # Let the in-flight commit finish and clean up. + agent.release.set() + t.join(timeout=5) + + +# --------------------------------------------------------------------------- # +# Finding 1: shutdown drain flushes the pending commit # +# --------------------------------------------------------------------------- # + + +def test_drain_joins_inflight_background_commit_and_flushes(): + """drain_all_on_shutdown() waits for the in-flight worker, then commits.""" + lifecycle = _fresh_lifecycle() + agent = RecordingAgent() # held in flight until we release it + sid = "drain-inflight" + lifecycle.register_agent(sid, agent) + lifecycle.mark_turn_completed(sid, agent=agent) + + t = _make_background_worker(lifecycle, sid, agent) + t.start() + assert agent.entered.wait(timeout=2) # commit is in flight + + drained = threading.Event() + + def _do_drain(): + lifecycle.drain_all_on_shutdown() + drained.set() + + dt = threading.Thread(target=_do_drain, daemon=True) + dt.start() + + # While the background commit is still in flight, drain must be blocked + # joining it — it cannot have finished yet. + time.sleep(0.1) + assert not drained.is_set() + + # Release the commit; drain should now join the worker and finish. + agent.release.set() + assert drained.wait(timeout=5) + dt.join(timeout=2) + + assert agent.calls == 1 + assert lifecycle.has_uncommitted_work(sid) is False + with lifecycle._background_commit_threads_lock: + assert t not in lifecycle._background_commit_threads + + +def test_sigterm_style_shutdown_runs_finally_drain(): + """The SIGTERM handler mechanism drives serve_forever() through its finally. + + Replicates server.py's approach: a signal handler requests shutdown by + dispatching httpd.shutdown() from a short-lived helper thread (you cannot + call it from serve_forever()'s own thread). serve_forever() returns, the + ``finally`` runs, and drain_all_on_shutdown() flushes the pending commit. + """ + lifecycle = _fresh_lifecycle() + agent = RecordingAgent() + agent.release.set() # commit completes immediately during drain + sid = "sigterm-drain" + lifecycle.register_agent(sid, agent) + lifecycle.mark_turn_completed(sid, agent=agent) + + class _Handler(BaseHTTPRequestHandler): + def log_message(self, *args): # silence test noise + pass + + httpd = ThreadingHTTPServer(("127.0.0.1", 0), _Handler) + finally_ran = threading.Event() + + def _serve(): + try: + httpd.serve_forever() + finally: + httpd.server_close() + lifecycle.drain_all_on_shutdown() + finally_ran.set() + + serve_thread = threading.Thread(target=_serve, daemon=True) + serve_thread.start() + + # Wait until the server is actually accepting connections before shutdown. + addr = httpd.server_address + for _ in range(100): + try: + with socket.create_connection(addr, timeout=0.1): + break + except OSError: + time.sleep(0.02) + + # --- exact handler mechanism from server.py --- + shutdown_requested = threading.Event() + + def _request_shutdown(): + if shutdown_requested.is_set(): + return + shutdown_requested.set() + threading.Thread(target=httpd.shutdown, daemon=True).start() + + _request_shutdown() + _request_shutdown() # idempotent: a repeated signal must not double-shutdown + + serve_thread.join(timeout=5) + assert not serve_thread.is_alive(), "serve_forever() did not exit on shutdown" + assert finally_ran.wait(timeout=5), "finally block / drain did not run" + assert agent.calls == 1 + assert lifecycle.has_uncommitted_work(sid) is False + + +# --------------------------------------------------------------------------- # +# Static wiring guards (the inline handler in main() cannot be imported) # +# --------------------------------------------------------------------------- # + + +def test_server_installs_sigterm_handler_and_drains_in_finally(): + """server.py must wire a SIGTERM handler and drain in serve_forever's finally.""" + src = (REPO / "server.py").read_text(encoding="utf-8") + tree = ast.parse(src) + + # A SIGTERM handler is installed via signal.signal(signal.SIGTERM, ...). + installs_sigterm = any( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "signal" + and node.args + and ast.dump(node.args[0]).find("SIGTERM") != -1 + for node in ast.walk(tree) + ) + assert installs_sigterm, "server.py does not install a SIGTERM handler" + + # httpd.shutdown() is dispatched (from a helper thread) rather than the + # default SIGTERM behavior that would skip the finally. + assert "httpd.shutdown" in src + # The existing shutdown drain is still wired. + assert "drain_all_on_shutdown" in src + # SIGPIPE handling must be preserved. + assert "SIGPIPE" in src + + +# --------------------------------------------------------------------------- # +# Deadline enforcement: a wedged worker must not hang the managed-stop drain # +# --------------------------------------------------------------------------- # + + +def test_drain_deadline_returns_despite_stuck_worker(): + """A background commit wedged in commit_memory_session() (e.g. a hung memory + provider) must NOT hang drain_all_on_shutdown() past its deadline. The inline + drain runs each commit in a budget-bounded daemon worker, so a stuck in_flight + session cannot block the drain. Regression for the dual-gate finding that the + top-of-loop deadline check never fired while blocked inside a stuck commit.""" + lifecycle = _fresh_lifecycle() + agent = RecordingAgent() # NEVER released → stays wedged in commit + sid = "stuck-drain" + lifecycle.register_agent(sid, agent) + lifecycle.mark_turn_completed(sid, agent=agent) + + # Put the worker in flight, but keep it OUT of the registry so the 5s join + # phase does not dominate the timing — we are isolating the inline-drain + # deadline path. + worker = threading.Thread( + target=lambda: lifecycle.commit_session_memory(sid, agent=agent), + daemon=True, + name="stuck-commit", + ) + worker.start() + assert agent.entered.wait(timeout=2) # worker is now wedged holding in_flight + + started = time.monotonic() + lifecycle.drain_all_on_shutdown(deadline=0.5) + elapsed = time.monotonic() - started + + # Without the timeout fix this blocks until the worker is released (never); + # with it, the drain returns within roughly the deadline budget. + assert elapsed < 3.0, f"drain hung {elapsed:.1f}s past its 0.5s deadline" + + # Let the wedged worker unwind so the test process exits cleanly. + agent.release.set() + worker.join(timeout=5) + + +def test_drain_deadline_returns_when_inline_provider_hangs_no_prior_worker(): + """Deeper case (2nd-round gate finding): a DIRTY session with NO pre-existing + in_flight worker, whose provider call hangs when the inline drain owns the + segment and calls commit_memory_session() DIRECTLY. The timeout= param only + bounds waiting on an already-in_flight commit; the direct provider call must + also be budget-bounded (run in a joined daemon worker), or the drain hangs.""" + lifecycle = _fresh_lifecycle() + agent = RecordingAgent() # NEVER released → the DIRECT provider call wedges + sid = "inline-hang" + lifecycle.register_agent(sid, agent) + lifecycle.mark_turn_completed(sid, agent=agent) + # No background worker started: the session is simply dirty + # (generation > committed_generation) with in_flight=False, so the inline + # drain itself acquires the segment and calls the hanging provider directly. + + started = time.monotonic() + lifecycle.drain_all_on_shutdown(deadline=0.5) + elapsed = time.monotonic() - started + + assert elapsed < 3.0, f"inline-drain hung {elapsed:.1f}s past its 0.5s deadline" + + # Let the wedged inline commit worker unwind cleanly. + agent.release.set()