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