ai-ag2026andClaude Opus 4.8
4fc51072a2
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) <noreply@anthropic.com>
87 lines
3.5 KiB
Python
87 lines
3.5 KiB
Python
"""
|
|
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
|