"""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()