tests/test_compression/test_per_request_segments.py

"""Tests for per-request compression segments capture and propagation.

Flow under test:
    ContentRouter observer → `record_compression` → thread-local
    accumulator (when inside `request_segments_scope`) →
    `TransformResult.compression_segments` → `RequestLog.compression_segments`
    → `/stats` `recent_requests`.
"""

from __future__ import annotations

from headroom.proxy.models import RequestLog
from headroom.proxy.prometheus_metrics import PrometheusMetrics
from headroom.transforms.observability import (
    current_request_segments,
    request_segments_scope,
)


def test_record_compression_appends_to_active_scope():
    """PrometheusMetrics.record_compression appends to the active
    per-request scope and still increments global counters."""
    m = PrometheusMetrics()
    # Pre-scope: accumulator is None, observer is a no-op for segments
    m.record_compression("smart_crusher", original_tokens=100, compressed_tokens=50)
    assert m.compressions_by_strategy["smart_crusher"] == 1
    assert m.tokens_original_by_strategy["smart_crusher"] == 100

    with request_segments_scope() as segs:
        assert current_request_segments() is segs
        m.record_compression("smart_crusher", original_tokens=300, compressed_tokens=80)
        m.record_compression("code_aware", original_tokens=200, compressed_tokens=120)
        assert len(segs) == 2
        assert segs[0] == {
            "strategy": "smart_crusher",
            "original_tokens": 300,
            "compressed_tokens": 80,
        }
        assert segs[1] == {
            "strategy": "code_aware",
            "original_tokens": 200,
            "compressed_tokens": 120,
        }

    # Global counters kept accumulating across the scope:
    assert m.compressions_by_strategy["smart_crusher"] == 2
    assert m.tokens_original_by_strategy["code_aware"] == 200

    # After scope, observer no longer appends anywhere:
    m.record_compression("diff", original_tokens=50, compressed_tokens=30)
    assert len(segs) == 2  # unchanged


def test_nested_scopes_isolate_per_request_state():
    """Nested scopes don't leak. The outer scope's list is restored
    after the inner scope exits."""
    with request_segments_scope() as outer:
        outer.append({"strategy": "outer", "original_tokens": 10, "compressed_tokens": 5})
        with request_segments_scope() as inner:
            assert current_request_segments() is inner
            inner.append({"strategy": "inner", "original_tokens": 20, "compressed_tokens": 10})
            assert len(outer) == 1  # inner append didn't touch outer
            assert len(inner) == 1
        # Inner scope exited — outer is restored as the current slot
        assert current_request_segments() is outer
        outer.append({"strategy": "outer2", "original_tokens": 30, "compressed_tokens": 15})
        assert len(outer) == 2


def test_exception_inside_scope_still_clears_slot():
    """Scope cleanup must run even if the body raises — otherwise a
    crashed request leaks segments into the next request on the same
    worker thread."""
    try:
        with request_segments_scope() as segs:
            segs.append({"strategy": "smart_crusher", "original_tokens": 5, "compressed_tokens": 2})
            raise RuntimeError("simulated pipeline failure")
    except RuntimeError:
        pass

    # Slot restored to None — no leak to whatever runs next on this thread.
    assert current_request_segments() is None


def test_request_log_carries_compression_segments():
    """RequestLog serialization roundtrip preserves the per-request
    routing log so `/stats` `recent_requests` consumers see it."""
    segments = [
        {"strategy": "smart_crusher", "original_tokens": 500, "compressed_tokens": 120},
        {"strategy": "code_aware", "original_tokens": 300, "compressed_tokens": 90},
    ]
    log = RequestLog(
        timestamp="2026-06-06T12:00:00Z",
        request_id="req-abc",
        provider="anthropic",
        model="claude-opus-4-7",
        input_tokens_original=800,
        input_tokens_optimized=210,
        output_tokens=0,
        tokens_saved=590,
        savings_percent=73.75,
        optimization_latency_ms=5.0,
        total_latency_ms=42.0,
        tags={},
        cache_hit=False,
        transforms_applied=["content_router"],
        compression_segments=segments,
    )

    from dataclasses import asdict

    d = asdict(log)
    assert d["compression_segments"] == segments
    # asdict is what request_logger uses for Redis serialization →
    # confirms the field survives the trip to `_merged_recent_requests`.


def test_request_log_default_compression_segments_is_none():
    """Legacy call sites that don't pass compression_segments get
    None — graceful degradation, no KeyError, no panic."""
    log = RequestLog(
        timestamp="2026-06-06T12:00:00Z",
        request_id="req-xyz",
        provider="openai",
        model="gpt-4",
        input_tokens_original=100,
        input_tokens_optimized=100,
        output_tokens=0,
        tokens_saved=0,
        savings_percent=0.0,
        optimization_latency_ms=1.0,
        total_latency_ms=10.0,
        tags={},
        cache_hit=False,
        transforms_applied=[],
    )
    assert log.compression_segments is None


def test_pipeline_apply_propagates_segments_through_transform_result():
    """End-to-end: TransformPipeline.apply wraps itself in
    request_segments_scope, so any observer-bearing transform that
    runs during apply() contributes segments to TransformResult."""
    from headroom.config import HeadroomConfig
    from headroom.transforms.pipeline import TransformPipeline

    # Build a minimal pipeline with a single ContentRouter that has
    # a real PrometheusMetrics observer attached. ContentRouter
    # emits one observer event per routing decision.
    cfg = HeadroomConfig()
    metrics = PrometheusMetrics()
    pipeline = TransformPipeline(cfg, observer=metrics)

    # Fabricate messages with enough content that the router will
    # actually compress (not passthrough). The exact strategy doesn't
    # matter — we just need at least one record_compression event.
    messages = [
        {"role": "system", "content": "You are helpful."},
        {
            "role": "user",
            "content": (
                "Here is a large JSON array of logs:\n"
                + "\n".join(f'{{"id": {i}, "level": "info", "msg": "event"}}' for i in range(200))
            ),
        },
    ]

    result = pipeline.apply(
        messages,
        model="claude-opus-4-7",
        model_limit=200000,
    )

    # If the router compressed anything, we should see segments both
    # on the result and in the global per-strategy counters. If the
    # router passthrough'd everything, both are empty — the invariant
    # is that they agree, not that compression must occur.
    assert result.compression_segments is None or isinstance(
        result.compression_segments, list
    )
    if result.compression_segments:
        # Every entry has the observer shape
        for entry in result.compression_segments:
            assert set(entry.keys()) == {
                "strategy",
                "original_tokens",
                "compressed_tokens",
            }
        # Global counters agree on totals
        total_original = sum(e["original_tokens"] for e in result.compression_segments)
        total_in_metrics = sum(metrics.tokens_original_by_strategy.values())
        assert total_original == total_in_metrics