diff --git a/headroom/proxy/cost.py b/headroom/proxy/cost.py index 0bb772f92..37670c309 100644 --- a/headroom/proxy/cost.py +++ b/headroom/proxy/cost.py @@ -745,9 +745,7 @@ class CostTracker: # clamp it so `total_tokens_removed` reflects actually-forwarded bytes # instead of surfacing spurious negatives (verified clean on the wire). if tokens_saved < 0: - import logging as _lg - - _lg.getLogger(__name__).debug( + logger.debug( "record_tokens: clamping negative tokens_saved=%d to 0 for %s (artifact; wire not inflated)", tokens_saved, model, diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 86a12ce43..2823de0b9 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -1070,7 +1070,7 @@ class OpenAIHandlerMixin: return base_frozen_count final_idx = len(messages) - 1 if messages[final_idx].get("role") in ("user", "tool", "function"): - return max(base_frozen_count, final_idx) + return final_idx return len(messages) @staticmethod @@ -2363,8 +2363,8 @@ class OpenAIHandlerMixin: openai_frozen_count = openai_prefix_tracker.get_frozen_message_count() if is_cache_mode(self.config.mode): - openai_frozen_count = self._strict_previous_turn_frozen_count( - original_client_messages, + openai_frozen_count = OpenAIHandlerMixin._strict_previous_turn_frozen_count( + messages, openai_frozen_count, ) @@ -2399,8 +2399,14 @@ class OpenAIHandlerMixin: # Zone 1: Swap cached compressed versions working_messages = comp_cache.apply_cached(messages) - # Re-freeze boundary - openai_frozen_count = comp_cache.compute_frozen_count(messages) + # Re-freeze boundary. Token mode can use the compression + # cache's positional frozen count. Cache mode must keep the + # latest observation mutable even when the compression + # cache has no compressible entry for it yet; otherwise + # OpenAI-compatible tool-call clients freeze the entire + # conversation and report near-zero savings. + if not is_cache_mode(self.config.mode): + openai_frozen_count = comp_cache.compute_frozen_count(messages) result = await self._run_compression_in_executor( lambda: self.openai_pipeline.apply( @@ -2408,7 +2414,14 @@ class OpenAIHandlerMixin: model=model, model_limit=context_limit, context=extract_user_query(working_messages), - frozen_message_count=openai_frozen_count, + frozen_message_count=( + OpenAIHandlerMixin._strict_previous_turn_frozen_count( + working_messages, + openai_frozen_count, + ) + if is_cache_mode(self.config.mode) + else openai_frozen_count + ), biases=_hook_biases, compression_policy=compression_policy, # Thread the savings-profile knobs (e.g. @@ -2434,13 +2447,21 @@ class OpenAIHandlerMixin: # so tokens_saved captures both Zone 1 + Zone 2 savings. optimized_tokens = result.tokens_after else: + apply_frozen_count = ( + OpenAIHandlerMixin._strict_previous_turn_frozen_count( + messages, + openai_frozen_count, + ) + if is_cache_mode(self.config.mode) + else openai_frozen_count + ) result = await self._run_compression_in_executor( lambda: self.openai_pipeline.apply( messages=messages, model=model, model_limit=context_limit, context=extract_user_query(messages), - frozen_message_count=openai_frozen_count, + frozen_message_count=apply_frozen_count, biases=_hook_biases, compression_policy=compression_policy, # Same savings-profile threading as the token-mode diff --git a/headroom/proxy/prometheus_metrics.py b/headroom/proxy/prometheus_metrics.py index 25033b6a0..b6ab41b15 100644 --- a/headroom/proxy/prometheus_metrics.py +++ b/headroom/proxy/prometheus_metrics.py @@ -597,9 +597,7 @@ class PrometheusMetrics: # model. Clamp so total_tokens_removed / avg_compression_pct reflect the # actually-forwarded bytes instead of surfacing spurious negatives. if tokens_saved < 0: - import logging as _lg - - _lg.getLogger(__name__).debug( + logger.debug( "metrics.record: clamping negative tokens_saved=%d to 0 for %s (artifact; wire not inflated)", tokens_saved, model, diff --git a/tests/test_cost_tracker_counterfactual.py b/tests/test_cost_tracker_counterfactual.py index f49c974db..63cc123cd 100644 --- a/tests/test_cost_tracker_counterfactual.py +++ b/tests/test_cost_tracker_counterfactual.py @@ -80,6 +80,19 @@ def test_savings_zero_when_no_tokens_saved(): assert stats["total_tokens_saved"] == 0 +def test_negative_token_savings_are_clamped_to_zero(): + """Estimator artifacts must not reduce cumulative savings below reality.""" + from headroom.proxy.server import CostTracker + + ct = CostTracker() + + ct.record_tokens("openai-compatible", tokens_saved=-500, tokens_sent=5_000) + stats = ct.stats() + + assert stats["total_tokens_saved"] == 0 + assert stats["per_model"]["openai-compatible"]["tokens_saved"] == 0 + + def test_multi_model_savings(): """Savings across multiple models use each model's own list price.""" from headroom.proxy.server import CostTracker diff --git a/tests/test_observability_metrics.py b/tests/test_observability_metrics.py index f49c3b71c..6b7e09890 100644 --- a/tests/test_observability_metrics.py +++ b/tests/test_observability_metrics.py @@ -198,3 +198,20 @@ async def test_prometheus_metrics_reads_late_configured_otel_metrics() -> None: assert spy.rate_limited_calls == [{"provider": "anthropic", "model": "claude-sonnet"}] finally: reset_otel_metrics() + + +@pytest.mark.asyncio +async def test_prometheus_metrics_clamps_negative_token_savings() -> None: + metrics = PrometheusMetrics() + + await metrics.record_request( + provider="openai", + model="openai-compatible", + input_tokens=100, + output_tokens=5, + tokens_saved=-25, + latency_ms=1.0, + ) + + assert metrics.tokens_saved_total == 0 + assert metrics.savings_history[-1][1] == 0 diff --git a/tests/test_proxy_handler_helpers.py b/tests/test_proxy_handler_helpers.py index e55662c75..2609d5fa1 100644 --- a/tests/test_proxy_handler_helpers.py +++ b/tests/test_proxy_handler_helpers.py @@ -221,6 +221,27 @@ def test_openai_handler_prefix_helpers_cover_edge_cases() -> None: ) == 1 ) + assert ( + OpenAIHandlerMixin._strict_previous_turn_frozen_count( + [{"role": "assistant"}, {"role": "tool", "content": "observation"}], + 0, + ) + == 1 + ) + assert ( + OpenAIHandlerMixin._strict_previous_turn_frozen_count( + [{"role": "user"}, {"role": "assistant"}, {"role": "tool", "content": "obs"}], + 3, + ) + == 2 + ) + assert ( + OpenAIHandlerMixin._strict_previous_turn_frozen_count( + [{"role": "assistant"}, {"role": "function", "content": "legacy observation"}], + 0, + ) + == 1 + ) assert ( OpenAIHandlerMixin._strict_previous_turn_frozen_count( [{"role": "user"}, {"role": "assistant"}], diff --git a/tests/test_proxy_openai_cache_stability.py b/tests/test_proxy_openai_cache_stability.py index 44339c91a..4cbd18d5d 100644 --- a/tests/test_proxy_openai_cache_stability.py +++ b/tests/test_proxy_openai_cache_stability.py @@ -112,6 +112,85 @@ def test_openai_cache_mode_freezes_previous_turns() -> None: assert captured["frozen_message_count"] == 2 +@pytest.mark.parametrize("tail_role", ["tool", "function"]) +def test_openai_cache_mode_keeps_final_tool_observation_mutable(tail_role: str) -> None: + captured = {} + with _make_proxy_client() as client: + proxy = client.app.state.proxy + proxy.config.optimize = True + proxy.config.mode = "cache" + + fake_tracker = _FakePrefixTracker(frozen_count=0) + proxy.session_tracker_store.compute_session_id = lambda request, model, messages: ( + "stable-session" + ) + proxy.session_tracker_store.get_or_create = lambda session_id, provider: fake_tracker + + def _fake_apply(**kwargs): + captured.setdefault("calls", []).append( + { + "frozen_message_count": kwargs.get("frozen_message_count"), + "roles": [msg.get("role") for msg in kwargs["messages"]], + "mode": proxy.config.mode, + } + ) + return SimpleNamespace( + messages=kwargs["messages"], + transforms_applied=["test:compress-tail"], + timing={}, + tokens_before=120, + tokens_after=80, + waste_signals=None, + ) + + proxy.openai_pipeline.apply = _fake_apply + + async def _fake_retry(method, url, headers, body, stream=False, **kwargs): # noqa: ANN001 + return httpx.Response( + 200, + json={ + "id": "chatcmpl_tool_tail", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 80, "completion_tokens": 3, "total_tokens": 83}, + }, + ) + + proxy._retry_request = _fake_retry + + tail = { + "role": tail_role, + "content": "large command observation " * 200, + } + if tail_role == "tool": + tail["tool_call_id"] = "call_1" + else: + tail["name"] = "bash" + + response = client.post( + "/v1/chat/completions", + headers={"authorization": "Bearer test-key"}, + json={ + "model": "gpt-4o-mini", + "messages": [ + {"role": "user", "content": "turn1"}, + {"role": "assistant", "content": "run command"}, + tail, + ], + }, + ) + + assert response.status_code == 200 + assert any(call["frozen_message_count"] == 2 for call in captured["calls"]), captured[ + "calls" + ] + + def test_openai_cache_mode_restores_mutated_frozen_prefix() -> None: captured = {} with _make_proxy_client() as client: