Merge pull request #5804 from nesquena/release/stage-5787

Release: pair metering begin_session with end_session teardown (#5787)

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)

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.

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