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