Spaces:
Build error
Build error
File size: 7,553 Bytes
95cb372 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 | """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
|