From 4fc51072a29223fc8dcc24594c8658c71f56dd11 Mon Sep 17 00:00:00 2001 From: ai-ag2026 <261867348+ai-ag2026@users.noreply.github.com> Date: Wed, 8 Jul 2026 20:25:10 +0200 Subject: [PATCH 1/2] fix(streaming): pair metering begin_session with end_session teardown (#4633, #2476) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The streaming worker called meter().begin_session(stream_id) but never a paired end_session. A turn that produced zero output tokens (pre-flight cancel, or a setup exception before the first token) left a _SessionMeter in GlobalMeter._sessions forever: get_stats() only prunes sessions with first_token_ts > 0, so a zero-token session is never reclaimed, and each leaked entry inflates the `active` count sent over the SSE metering event. begin_session() and the metering ticker's .start() are now registered as the first statements inside the worker's outer `try`, and paired with an idempotent meter().end_session(stream_id, 0) plus a deterministic _metering_stop.set() in that try's outer `finally` (the same block that pops STREAMS/CANCEL_FLAGS) — so every exit path tears the session down. Deferring .start() until after put() is defined also closes a latent start-before-put ordering window in the ticker closure. The metering payload is unchanged; end_session only pops the session. Verification: tests/test_metering_session_lifecycle.py (leak reproduction, end_session reclaim, idempotency, begin->cancel-before-token->empty) plus a plain-assert run on python3.11 (conftest gates pytest to 3.11-3.13; system python is 3.14). py_compile clean. WARNING streaming-contract: adjusts the streaming worker's begin/end teardown ordering. Needs RFC/contract review against docs/rfcs/ (session-sse-contract-v1.md) before merge — draft PR. Co-Authored-By: Claude Opus 4.8 (1M context) --- api/streaming.py | 35 ++++++++-- tests/test_metering_session_lifecycle.py | 86 ++++++++++++++++++++++++ 2 files changed, 117 insertions(+), 4 deletions(-) create mode 100644 tests/test_metering_session_lifecycle.py diff --git a/api/streaming.py b/api/streaming.py index d8046bb6..46cff87b 100644 --- a/api/streaming.py +++ b/api/streaming.py @@ -6711,12 +6711,18 @@ def _run_agent_streaming( return _usage - # Register this stream with the global streaming meter - meter().begin_session(stream_id) - # Metering ticker — emits a metering event at 1 Hz while sessions are active. # When get_interval() returns >= 10.0 (no active sessions), the ticker exits # so no idle readings are emitted and the SSE consumer sees nothing. + # + # #4633/#2476: begin_session() and the ticker .start() are deferred into the + # outer `try` below so the outer `finally` (which pops STREAMS/CANCEL_FLAGS) + # always runs its paired end_session()/_metering_stop.set() teardown. A raise + # between here and that `try` would otherwise leak the _sessions[stream_id] + # entry — get_stats() only prunes sessions with first_token_ts > 0, so a + # zero-token turn (pre-flight cancel, setup raise) is never reclaimed and its + # count inflates the SSE `active` field. Deferring .start() until after `put` + # is defined also removes a latent start-before-put ordering window. _metering_stop = threading.Event() def _metering_ticker(): @@ -6732,7 +6738,6 @@ def _run_agent_streaming( put('metering', stats) _metering_thread = threading.Thread(target=_metering_ticker, daemon=True) - _metering_thread.start() def put(event, data): # If cancelled, drop all further events except the cancel event itself @@ -6809,6 +6814,11 @@ def _run_agent_streaming( _ckpt_thread = None _agent_lock = None try: + # Register this stream with the global streaming meter and start the 1 Hz + # metering ticker. Kept INSIDE the outer try so the outer `finally`'s + # end_session()/_metering_stop.set() teardown is always paired (#4633/#2476). + meter().begin_session(stream_id) + _metering_thread.start() # Bind THIS turn's session identity to the worker thread/context BEFORE # any agent work (so every mid-turn notify_on_complete background spawn # captures THIS session, not a concurrent turn's process-global env). @@ -9865,6 +9875,23 @@ def _run_agent_streaming( _error_payload['old_session_id'] = session_id put('apperror', _error_payload) finally: + # #4633/#2476: symmetric metering teardown. begin_session() (top of the + # outer try) had no paired end_session(), so zero-token turns leaked a + # _sessions[stream_id] entry that get_stats() pruning never reclaims (its + # criterion requires first_token_ts > 0). end_session() is idempotent — + # it just pops _sessions[stream_id]; the metering payload is unchanged. + # _metering_stop.set() deterministically stops the ticker (the inner + # finally also sets it on the normal path; setting twice is harmless). + try: + # 0: end_session() currently ignores final_output_tokens — it only + # pops _sessions[stream_id]. If it is ever extended to consume the + # count (e.g. persisting final output tokens to a billing ledger), + # this teardown caller will need to supply the real total; the outer + # finally doesn't have easy access to it today. + meter().end_session(stream_id, 0) + except Exception: + logger.debug("Failed to end metering session for stream %s", stream_id, exc_info=True) + _metering_stop.set() # Stop the periodic checkpoint thread before the final recovery path. # The checkpoint thread also uses the per-session lock; joining it first # avoids contending with checkpoint writes during stale-pending repair. diff --git a/tests/test_metering_session_lifecycle.py b/tests/test_metering_session_lifecycle.py new file mode 100644 index 00000000..78fd0a3f --- /dev/null +++ b/tests/test_metering_session_lifecycle.py @@ -0,0 +1,86 @@ +""" +Regression tests for #4633 / #2476 — streaming meter session lifecycle. + +Before this fix, ``_run_agent_streaming`` called ``meter().begin_session(stream_id)`` +but never called a paired ``end_session``. A streaming turn that produced zero +output tokens (pre-flight cancel, a setup exception before the first token) left +a ``_SessionMeter`` in ``GlobalMeter._sessions`` forever: + + * ``get_stats()`` only prunes sessions with ``first_token_ts > 0`` — a + zero-token session has ``first_token_ts == 0.0`` and is therefore NEVER + reclaimed. + * The leaked session inflates ``get_stats()['active']`` (``len(_sessions)``), + which is reported to the browser via the SSE ``metering`` event. + +The fix registers ``begin_session`` as the first statement inside the worker's +outer ``try`` and pairs it with ``meter().end_session(stream_id, 0)`` in that +try's outer ``finally`` (the same block that pops ``STREAMS``/``CANCEL_FLAGS``), +so every exit path — normal completion, exception, or pre-flight cancel — tears +the session down. ``end_session`` only pops ``_sessions[stream_id]``; the +metering payload is unchanged. + +These tests exercise the ``GlobalMeter`` contract directly, since that is where +the leak invariant lives. +""" +import pytest + +from api.metering import GlobalMeter + + +@pytest.fixture() +def m(): + """A fresh, isolated meter per test (avoids touching the module singleton).""" + return GlobalMeter() + + +class TestZeroTokenLeak: + + def test_get_stats_does_not_prune_zero_token_session(self, m): + """Documents the leak: a session with no token is never pruned by get_stats().""" + m.begin_session('s1') + # No record_token() -> first_token_ts stays 0.0. + m.get_stats() # pruning pass + assert 's1' in m._sessions + assert m.get_stats()['active'] == 1 + + def test_end_session_reclaims_zero_token_session(self, m): + """The fix's mechanism: end_session pops the leaked zero-token session.""" + m.begin_session('s1') + m.end_session('s1', 0) + assert 's1' not in m._sessions + assert m.get_stats()['active'] == 0 + + def test_end_session_is_idempotent(self, m): + """The outer finally may run after the inner path already ended the session.""" + m.begin_session('s1') + m.end_session('s1', 0) + m.end_session('s1', 0) # must not raise + assert 's1' not in m._sessions + + def test_end_session_on_never_begun_stream_is_safe(self, m): + """finally runs even if begin_session never executed (early raise).""" + m.end_session('never', 0) # must not raise / KeyError + assert m._sessions == {} + + +class TestBeginCancelBeforeToken: + + def test_begin_then_cancel_before_token_leaves_no_session(self, m): + """begin -> cancel-before-first-token -> end_session -> _sessions empty.""" + stream_id = 'cancel_preflight' + m.begin_session(stream_id) + assert len(m._sessions) == 1 + # No token recorded (cancelled before the model emitted anything). + # The worker's outer finally calls end_session(stream_id, 0): + m.end_session(stream_id, 0) + assert len(m._sessions) == 0 + + def test_normal_turn_still_reclaimed(self, m): + """A turn that DID produce tokens is also reclaimed by end_session.""" + stream_id = 'normal' + m.begin_session(stream_id) + m.record_token(stream_id, 10) + m.record_token(stream_id, 42) + assert m._sessions[stream_id].output_tokens == 42 + m.end_session(stream_id, 42) + assert len(m._sessions) == 0 From 2c46e9e251b023d3f417aa0848c77a12f9b972a7 Mon Sep 17 00:00:00 2001 From: nesquena-hermes Date: Thu, 9 Jul 2026 00:51:59 +0000 Subject: [PATCH 2/2] docs(changelog): #5787 metering session begin/end pairing --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index af558915..0d9d8b84 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,8 @@ ### Fixed +- **Zero-token turns (a pre-flight cancel or setup error) no longer leak a streaming-meter session.** The streaming meter's `begin_session()` and its 1 Hz ticker were started before the outer `try`, so a raise or cancel before the first token left the `_sessions[stream_id]` entry unreclaimed (`get_stats()` only prunes sessions that emitted a token), inflating the SSE `active` count over a long-lived server. `begin_session()` + the ticker start now live inside the outer `try` so the paired `end_session()` / ticker-stop teardown in the `finally` always runs. Thanks @ai-ag2026. (#5787, #4633, #2476) + - **Deleting a session now evicts its run-journal writer locks instead of leaking them.** The per-run-journal writer-lock cache (`_WRITER_LOCKS`) kept an entry after a session's journal was deleted, so the cache grew unbounded, and a later journal reusing the same path could be handed a stale lock. `delete_run_journal` now evicts that session's lock-cache entries — but only after confirming the directory was actually removed, so a Windows/permission transient that leaves files (locks still live) never has its lock wrongly dropped out from under a writer. Thanks @ai-ag2026. (#5784) - **The background reaper and drain daemon threads can no longer be started twice under a race.** `start_session_channel_reaper` and `start_drain_thread` used a check-then-act (`is_alive()` then create) with no lock, so two concurrent callers could both see "not running" and each spawn a thread — the loser's thread was never referenced by the module and ran forever, un-joinable. Both start paths now perform the liveness check, thread creation, and global assignment atomically under a dedicated lifecycle lock. Stop/join still run outside the lock, so there's no deadlock. Thanks @ai-ag2026. (#5783)