mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
Merge pull request #420 from chopratejas/fix-codex-responses-pyo3-and-frozen-count
fix: re-enable Codex /v1/responses compression + fix prose-format over-freeze
This commit is contained in:
Tejas Chopra
GitHub
commit
7aaa4ac48f
14 changed files with 1362 additions and 55 deletions
|
|
@ -5,14 +5,14 @@
|
|||
},
|
||||
"metadata": {
|
||||
"description": "Headroom marketplace for Claude Code and GitHub Copilot CLI plugins.",
|
||||
"version": "0.21.4"
|
||||
"version": "0.21.5"
|
||||
},
|
||||
"plugins": [
|
||||
{
|
||||
"name": "headroom",
|
||||
"source": "./plugins/headroom-agent-hooks",
|
||||
"description": "Headroom startup hooks for Claude Code and GitHub Copilot CLI.",
|
||||
"version": "0.21.4",
|
||||
"version": "0.21.5",
|
||||
"author": {
|
||||
"name": "Headroom Contributors",
|
||||
"url": "https://github.com/chopratejas/headroom"
|
||||
|
|
|
|||
4
.github/plugin/marketplace.json
vendored
4
.github/plugin/marketplace.json
vendored
|
|
@ -5,14 +5,14 @@
|
|||
},
|
||||
"metadata": {
|
||||
"description": "Headroom marketplace for Claude Code and GitHub Copilot CLI plugins.",
|
||||
"version": "0.21.4"
|
||||
"version": "0.21.5"
|
||||
},
|
||||
"plugins": [
|
||||
{
|
||||
"name": "headroom",
|
||||
"source": "./plugins/headroom-agent-hooks",
|
||||
"description": "Headroom startup hooks for Claude Code and GitHub Copilot CLI.",
|
||||
"version": "0.21.4",
|
||||
"version": "0.21.5",
|
||||
"author": {
|
||||
"name": "Headroom Contributors",
|
||||
"url": "https://github.com/chopratejas/headroom"
|
||||
|
|
|
|||
|
|
@ -390,6 +390,42 @@ impl CompressionManifest {
|
|||
.iter()
|
||||
.any(|b| matches!(b.action, BlockAction::Compressed { .. }))
|
||||
}
|
||||
|
||||
/// Aggregate `original_tokens − compressed_tokens` across every
|
||||
/// `BlockAction::Compressed` outcome. Zero when no block was
|
||||
/// rewritten. Saturating subtraction guards against the
|
||||
/// theoretically-impossible case where a `Compressed` variant
|
||||
/// reports compressed > original (the dispatcher's
|
||||
/// `RejectedNotSmaller` gate should make this unreachable, but the
|
||||
/// saturating arithmetic keeps callers panic-free).
|
||||
pub fn tokens_saved(&self) -> usize {
|
||||
self.block_outcomes
|
||||
.iter()
|
||||
.filter_map(|b| match &b.action {
|
||||
BlockAction::Compressed {
|
||||
original_tokens,
|
||||
compressed_tokens,
|
||||
..
|
||||
} => Some(original_tokens.saturating_sub(*compressed_tokens)),
|
||||
_ => None,
|
||||
})
|
||||
.sum()
|
||||
}
|
||||
|
||||
/// Distinct compressor strategies that actually produced rewritten
|
||||
/// output, in first-seen order. Mirrors what the proxy logs as
|
||||
/// `transforms_applied`. Empty when no block was rewritten.
|
||||
pub fn transforms_applied(&self) -> Vec<&'static str> {
|
||||
let mut seen: Vec<&'static str> = Vec::new();
|
||||
for b in &self.block_outcomes {
|
||||
if let BlockAction::Compressed { strategy, .. } = &b.action {
|
||||
if !seen.contains(strategy) {
|
||||
seen.push(*strategy);
|
||||
}
|
||||
}
|
||||
}
|
||||
seen
|
||||
}
|
||||
}
|
||||
|
||||
/// Outcome of dispatching the live zone.
|
||||
|
|
@ -1554,6 +1590,111 @@ mod tests {
|
|||
assert_eq!(manifest.messages_below_frozen_floor, 1);
|
||||
assert_eq!(manifest.latest_user_message_index, None);
|
||||
}
|
||||
|
||||
// ─── Manifest accessor helpers (consumed by PyO3 binding) ─────────
|
||||
|
||||
fn make_manifest(actions: Vec<BlockAction>) -> CompressionManifest {
|
||||
CompressionManifest {
|
||||
messages_total: actions.len(),
|
||||
messages_below_frozen_floor: 0,
|
||||
latest_user_message_index: None,
|
||||
block_outcomes: actions
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.map(|(i, a)| BlockOutcome {
|
||||
message_index: i,
|
||||
block_index: None,
|
||||
block_type: "test".to_string(),
|
||||
action: a,
|
||||
})
|
||||
.collect(),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tokens_saved_zero_for_empty_manifest() {
|
||||
let m = CompressionManifest::empty();
|
||||
assert_eq!(m.tokens_saved(), 0);
|
||||
assert!(m.transforms_applied().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tokens_saved_sums_compressed_outcomes_only() {
|
||||
let m = make_manifest(vec![
|
||||
BlockAction::Compressed {
|
||||
strategy: "smart_crusher",
|
||||
original_bytes: 0,
|
||||
compressed_bytes: 0,
|
||||
original_tokens: 100,
|
||||
compressed_tokens: 30,
|
||||
},
|
||||
BlockAction::NoCompressionApplied {
|
||||
content_type: "image".to_string(),
|
||||
},
|
||||
BlockAction::Compressed {
|
||||
strategy: "log_compressor",
|
||||
original_bytes: 0,
|
||||
compressed_bytes: 0,
|
||||
original_tokens: 200,
|
||||
compressed_tokens: 50,
|
||||
},
|
||||
BlockAction::RejectedNotSmaller {
|
||||
strategy: "smart_crusher",
|
||||
original_bytes: 0,
|
||||
compressed_bytes: 0,
|
||||
original_tokens: 80,
|
||||
compressed_tokens: 90,
|
||||
},
|
||||
]);
|
||||
// 70 + 150 = 220; rejected variant must not contribute.
|
||||
assert_eq!(m.tokens_saved(), 220);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn transforms_applied_dedup_first_seen_order() {
|
||||
let m = make_manifest(vec![
|
||||
BlockAction::Compressed {
|
||||
strategy: "log_compressor",
|
||||
original_bytes: 0,
|
||||
compressed_bytes: 0,
|
||||
original_tokens: 50,
|
||||
compressed_tokens: 10,
|
||||
},
|
||||
BlockAction::Compressed {
|
||||
strategy: "smart_crusher",
|
||||
original_bytes: 0,
|
||||
compressed_bytes: 0,
|
||||
original_tokens: 50,
|
||||
compressed_tokens: 10,
|
||||
},
|
||||
BlockAction::Compressed {
|
||||
strategy: "log_compressor",
|
||||
original_bytes: 0,
|
||||
compressed_bytes: 0,
|
||||
original_tokens: 50,
|
||||
compressed_tokens: 10,
|
||||
},
|
||||
]);
|
||||
assert_eq!(
|
||||
m.transforms_applied(),
|
||||
vec!["log_compressor", "smart_crusher"]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tokens_saved_saturates_when_compressed_exceeds_original() {
|
||||
// Defensive — the dispatcher's RejectedNotSmaller gate should
|
||||
// make this unreachable, but the helper must not panic if a
|
||||
// future caller hand-constructs such a manifest.
|
||||
let m = make_manifest(vec![BlockAction::Compressed {
|
||||
strategy: "smart_crusher",
|
||||
original_bytes: 0,
|
||||
compressed_bytes: 0,
|
||||
original_tokens: 10,
|
||||
compressed_tokens: 50,
|
||||
}]);
|
||||
assert_eq!(m.tokens_saved(), 0);
|
||||
}
|
||||
}
|
||||
|
||||
// ─── OpenAI Chat Completions live-zone dispatcher (Phase C PR-C2) ────────
|
||||
|
|
|
|||
|
|
@ -1498,7 +1498,7 @@ fn compress_openai_responses_live_zone(
|
|||
body: &[u8],
|
||||
auth_mode: &str,
|
||||
model: &str,
|
||||
) -> (Py<PyBytes>, bool) {
|
||||
) -> (Py<PyBytes>, bool, u64, Vec<String>) {
|
||||
let mode = match auth_mode.to_ascii_lowercase().as_str() {
|
||||
"payg" => RustLiveZoneAuthMode::Payg,
|
||||
"oauth" => RustLiveZoneAuthMode::OAuth,
|
||||
|
|
@ -1512,17 +1512,41 @@ fn compress_openai_responses_live_zone(
|
|||
};
|
||||
|
||||
match rust_compress_openai_responses_live_zone(body, mode, model_str) {
|
||||
Ok(LiveZoneOutcome::NoChange { .. }) => (PyBytes::new_bound(py, body).unbind(), false),
|
||||
Ok(LiveZoneOutcome::Modified { new_body, .. }) => {
|
||||
Ok(LiveZoneOutcome::NoChange { manifest }) => {
|
||||
let saved = manifest.tokens_saved() as u64;
|
||||
let transforms: Vec<String> = manifest
|
||||
.transforms_applied()
|
||||
.into_iter()
|
||||
.map(String::from)
|
||||
.collect();
|
||||
(
|
||||
PyBytes::new_bound(py, body).unbind(),
|
||||
false,
|
||||
saved,
|
||||
transforms,
|
||||
)
|
||||
}
|
||||
Ok(LiveZoneOutcome::Modified { new_body, manifest }) => {
|
||||
// `RawValue::get` returns the underlying serialized JSON
|
||||
// as `&str`; bytes are valid UTF-8 by construction.
|
||||
let bytes = new_body.get().as_bytes();
|
||||
(PyBytes::new_bound(py, bytes).unbind(), true)
|
||||
let saved = manifest.tokens_saved() as u64;
|
||||
let transforms: Vec<String> = manifest
|
||||
.transforms_applied()
|
||||
.into_iter()
|
||||
.map(String::from)
|
||||
.collect();
|
||||
(
|
||||
PyBytes::new_bound(py, bytes).unbind(),
|
||||
true,
|
||||
saved,
|
||||
transforms,
|
||||
)
|
||||
}
|
||||
Err(_) => {
|
||||
// BodyNotJson / NoMessagesArray are non-fatal: nothing to
|
||||
// compress, fall through to passthrough byte-for-byte.
|
||||
(PyBytes::new_bound(py, body).unbind(), false)
|
||||
(PyBytes::new_bound(py, body).unbind(), false, 0, Vec::new())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
16
headroom/cache/compression_cache.py
vendored
16
headroom/cache/compression_cache.py
vendored
|
|
@ -228,6 +228,18 @@ class CompressionCache:
|
|||
an assistant message with tool_use blocks, or a tool_result whose
|
||||
content hash is already in the cache. The first unstable tool_result
|
||||
(cache miss) stops the count.
|
||||
|
||||
The trailing message is *always* excluded from the frozen prefix
|
||||
(cap of ``len(messages) - 1``). The trailing message represents
|
||||
the just-arrived turn — by definition it has not yet been sent
|
||||
upstream and therefore cannot be in any provider prefix cache.
|
||||
Without this cap, prose-format clients (Cline, OpenClaude, Aider,
|
||||
any client that does not use OpenAI-native or Anthropic-native
|
||||
tool messages) would have every message marked stable, making
|
||||
the live zone empty and producing zero compression. See issue
|
||||
observed 2026-05-07 with Cline+DeepSeek (`Pipeline: freezing
|
||||
first N/N messages` followed by ``Transform content_router:
|
||||
X -> X tokens (saved 0)`` on every request).
|
||||
"""
|
||||
with self._lock:
|
||||
count = 0
|
||||
|
|
@ -244,7 +256,9 @@ class CompressionCache:
|
|||
# Regular user/assistant/system messages and assistant+tool_use
|
||||
# are always stable — fall through.
|
||||
count += 1
|
||||
return count
|
||||
# Reserve the trailing message as the live zone. `max(0, ...)`
|
||||
# handles the empty-list edge case cleanly.
|
||||
return min(count, max(0, len(messages) - 1))
|
||||
|
||||
def apply_cached(self, messages: list[dict]) -> list[dict]:
|
||||
"""Return a new list with cached compressions swapped into tool results.
|
||||
|
|
|
|||
|
|
@ -1091,6 +1091,16 @@ class OpenAIHandlerMixin:
|
|||
|
||||
get_codex_rate_limit_state().update_from_headers(dict(response.headers))
|
||||
|
||||
# Tag the metric/log with auth_mode + endpoint so the
|
||||
# dashboard can break down by client class (PAYG vs
|
||||
# subscription vs OAuth) without re-classifying.
|
||||
_auth_mode_chat = getattr(request.state, "auth_mode", None)
|
||||
_chat_log_tags = {
|
||||
**(tags or {}),
|
||||
"auth_mode": _auth_mode_chat.value if _auth_mode_chat else "payg",
|
||||
"endpoint": "chat_completions",
|
||||
}
|
||||
|
||||
await self.metrics.record_request(
|
||||
provider="openai",
|
||||
model=model,
|
||||
|
|
@ -1105,6 +1115,41 @@ class OpenAIHandlerMixin:
|
|||
uncached_input_tokens=uncached_input_tokens,
|
||||
)
|
||||
|
||||
# Per-request log entry for /transformations/feed +
|
||||
# /stats `recent_requests`. Without this the dashboard
|
||||
# only shows aggregates for non-streaming OpenAI traffic.
|
||||
# Mirror of the streaming.py + anthropic.py wiring.
|
||||
if getattr(self, "logger", None) is not None:
|
||||
from headroom.proxy.helpers import compute_turn_id
|
||||
from headroom.proxy.models import RequestLog
|
||||
|
||||
self.logger.log(
|
||||
RequestLog(
|
||||
request_id=request_id,
|
||||
timestamp=datetime.now().isoformat(),
|
||||
provider="openai",
|
||||
model=model,
|
||||
input_tokens_original=original_tokens,
|
||||
input_tokens_optimized=optimized_tokens,
|
||||
output_tokens=output_tokens,
|
||||
tokens_saved=tokens_saved,
|
||||
savings_percent=(tokens_saved / original_tokens * 100)
|
||||
if original_tokens > 0
|
||||
else 0,
|
||||
optimization_latency_ms=optimization_latency,
|
||||
total_latency_ms=total_latency,
|
||||
tags=_chat_log_tags,
|
||||
cache_hit=False,
|
||||
transforms_applied=transforms_applied,
|
||||
request_messages=body.get("messages")
|
||||
if getattr(self.config, "log_full_messages", False)
|
||||
else None,
|
||||
turn_id=compute_turn_id(
|
||||
model, body.get("system"), body.get("messages")
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
if tokens_saved > 0:
|
||||
logger.info(
|
||||
f"[{request_id}] {model}: {original_tokens:,} → {optimized_tokens:,} "
|
||||
|
|
@ -1493,18 +1538,34 @@ class OpenAIHandlerMixin:
|
|||
)
|
||||
|
||||
_input_bytes = json.dumps(body).encode("utf-8")
|
||||
_new_bytes, _modified = _rust_compress_responses(
|
||||
(
|
||||
_new_bytes,
|
||||
_modified,
|
||||
_rust_tokens_saved,
|
||||
_rust_transforms,
|
||||
) = _rust_compress_responses(
|
||||
_input_bytes,
|
||||
auth_mode.value,
|
||||
model,
|
||||
)
|
||||
if _modified:
|
||||
body = json.loads(_new_bytes)
|
||||
transforms_applied = list(transforms_applied) + ["openai_responses_live_zone"]
|
||||
tokens_saved = int(_rust_tokens_saved)
|
||||
optimized_tokens = max(0, original_tokens - tokens_saved)
|
||||
transforms_applied = [
|
||||
"openai_responses_live_zone",
|
||||
*_rust_transforms,
|
||||
*list(transforms_applied),
|
||||
]
|
||||
logger.info(
|
||||
f"[{request_id}] /v1/responses compressed "
|
||||
f"{len(_input_bytes):,}→{len(_new_bytes):,} bytes "
|
||||
f"(auth_mode={auth_mode.value})"
|
||||
"[%s] /v1/responses compressed %d→%d bytes "
|
||||
"(%d tokens saved, auth_mode=%s, transforms=%s)",
|
||||
request_id,
|
||||
len(_input_bytes),
|
||||
len(_new_bytes),
|
||||
tokens_saved,
|
||||
auth_mode.value,
|
||||
transforms_applied,
|
||||
)
|
||||
except Exception as _e:
|
||||
logger.warning(
|
||||
|
|
@ -1625,6 +1686,12 @@ class OpenAIHandlerMixin:
|
|||
if self.cost_tracker:
|
||||
self.cost_tracker.record_tokens(model, tokens_saved, total_input_tokens)
|
||||
|
||||
_resp_log_tags = {
|
||||
**(tags or {}),
|
||||
"auth_mode": auth_mode.value if auth_mode else "payg",
|
||||
"endpoint": "responses_http",
|
||||
}
|
||||
|
||||
await self.metrics.record_request(
|
||||
provider="openai",
|
||||
model=model,
|
||||
|
|
@ -1635,6 +1702,39 @@ class OpenAIHandlerMixin:
|
|||
overhead_ms=optimization_latency,
|
||||
)
|
||||
|
||||
# Per-request log entry for /transformations/feed +
|
||||
# /stats `recent_requests`. Mirror of streaming.py /
|
||||
# anthropic.py wiring; without this the dashboard's
|
||||
# per-request feed misses every Codex HTTP turn.
|
||||
if getattr(self, "logger", None) is not None:
|
||||
from headroom.proxy.helpers import compute_turn_id
|
||||
from headroom.proxy.models import RequestLog
|
||||
|
||||
self.logger.log(
|
||||
RequestLog(
|
||||
request_id=request_id,
|
||||
timestamp=datetime.now().isoformat(),
|
||||
provider="openai",
|
||||
model=model,
|
||||
input_tokens_original=original_tokens,
|
||||
input_tokens_optimized=optimized_tokens,
|
||||
output_tokens=output_tokens,
|
||||
tokens_saved=tokens_saved,
|
||||
savings_percent=(tokens_saved / original_tokens * 100)
|
||||
if original_tokens > 0
|
||||
else 0,
|
||||
optimization_latency_ms=optimization_latency,
|
||||
total_latency_ms=total_latency,
|
||||
tags=_resp_log_tags,
|
||||
cache_hit=False,
|
||||
transforms_applied=transforms_applied,
|
||||
request_messages=messages
|
||||
if getattr(self.config, "log_full_messages", False)
|
||||
else None,
|
||||
turn_id=compute_turn_id(model, body.get("instructions"), messages),
|
||||
)
|
||||
)
|
||||
|
||||
logger.info(f"[{request_id}] /v1/responses {model}: {total_input_tokens:,} tokens")
|
||||
|
||||
# Capture Codex rate-limit window data from response headers
|
||||
|
|
@ -1707,6 +1807,13 @@ class OpenAIHandlerMixin:
|
|||
|
||||
# Forward client headers to upstream, adding required OpenAI-Beta header
|
||||
ws_headers = dict(websocket.headers)
|
||||
# Extract per-request tags from headers up front so the
|
||||
# session-end RequestLog can attach them. `_extract_tags` is
|
||||
# the same helper the HTTP handlers use; on a WebSocket the
|
||||
# tags come from `x-headroom-tag-*` headers in the upgrade
|
||||
# handshake. Returns `{}` when no tags are present.
|
||||
_extract_ws_tags = getattr(self, "_extract_tags", None)
|
||||
ws_tags = _extract_ws_tags(ws_headers) if callable(_extract_ws_tags) else {}
|
||||
|
||||
# Extract subprotocol from client — this is an application-level negotiation
|
||||
# that MUST be forwarded end-to-end (unlike sec-websocket-key which is per-connection).
|
||||
|
|
@ -1921,10 +2028,14 @@ class OpenAIHandlerMixin:
|
|||
# via uvicorn). Hot-fix follow-up to PR #406: the first frame
|
||||
# is now compressed via the inline PyO3 binding right before
|
||||
# upstream send (see the compression block further below).
|
||||
# Subsequent client→upstream frames in the relay loop remain
|
||||
# unmodified — multi-frame compression is a separate follow-up.
|
||||
# Subsequent client→upstream frames are now ALSO compressed
|
||||
# via `_maybe_compress_response_create_frame` in
|
||||
# `_client_to_upstream` so long-lived subscription Codex
|
||||
# sessions get savings on every turn, not just the first.
|
||||
body: dict[str, Any] = {}
|
||||
tokens_saved = 0
|
||||
transforms_applied: list[str] = []
|
||||
ws_frames_compressed = 0
|
||||
try:
|
||||
body = json.loads(first_msg_raw)
|
||||
except json.JSONDecodeError:
|
||||
|
|
@ -2111,7 +2222,12 @@ class OpenAIHandlerMixin:
|
|||
_model = (_inner.get("model") if isinstance(_inner, dict) else None) or ""
|
||||
|
||||
_inner_bytes = json.dumps(_inner).encode("utf-8")
|
||||
_new_bytes, _modified = _rust_compress_responses(
|
||||
(
|
||||
_new_bytes,
|
||||
_modified,
|
||||
_ws_rust_saved,
|
||||
_ws_rust_transforms,
|
||||
) = _rust_compress_responses(
|
||||
_inner_bytes,
|
||||
_ws_auth_mode.value,
|
||||
_model,
|
||||
|
|
@ -2127,10 +2243,23 @@ class OpenAIHandlerMixin:
|
|||
else:
|
||||
_send_body = _new_inner
|
||||
first_msg_raw = json.dumps(_send_body)
|
||||
tokens_saved += int(_ws_rust_saved)
|
||||
for _t in (
|
||||
"openai_responses_ws_live_zone",
|
||||
*list(_ws_rust_transforms),
|
||||
):
|
||||
if _t not in transforms_applied:
|
||||
transforms_applied.append(_t)
|
||||
logger.info(
|
||||
f"[{request_id}] WS /v1/responses compressed "
|
||||
f"{len(_inner_bytes):,}→{len(_new_bytes):,} bytes "
|
||||
f"(auth_mode={_ws_auth_mode.value})"
|
||||
"[%s] WS /v1/responses compressed "
|
||||
"%d→%d bytes (%d tokens saved, "
|
||||
"auth_mode=%s, transforms=%s)",
|
||||
request_id,
|
||||
len(_inner_bytes),
|
||||
len(_new_bytes),
|
||||
int(_ws_rust_saved),
|
||||
_ws_auth_mode.value,
|
||||
transforms_applied,
|
||||
)
|
||||
except Exception as _ce:
|
||||
logger.warning(
|
||||
|
|
@ -2192,11 +2321,107 @@ class OpenAIHandlerMixin:
|
|||
upstream_relay_error: BaseException | None = None
|
||||
client_relay_error: BaseException | None = None
|
||||
|
||||
async def _maybe_compress_response_create_frame(
|
||||
raw_msg: str,
|
||||
) -> str:
|
||||
"""Compress a single client→upstream frame
|
||||
when its `type` is `response.create`. Other
|
||||
event types (response.cancel, session.update,
|
||||
etc.) pass through unchanged. Errors are
|
||||
warned and the original frame is returned —
|
||||
fail loud in logs, fail safe on the wire.
|
||||
Updates outer-scope ``tokens_saved``,
|
||||
``transforms_applied``, and
|
||||
``ws_frames_compressed`` so the session-end
|
||||
log reports cumulative savings across all
|
||||
frames in the WS session.
|
||||
"""
|
||||
nonlocal tokens_saved, transforms_applied
|
||||
nonlocal ws_frames_compressed
|
||||
if not self.config.optimize:
|
||||
return raw_msg
|
||||
try:
|
||||
parsed_frame = json.loads(raw_msg)
|
||||
except json.JSONDecodeError:
|
||||
return raw_msg
|
||||
if (
|
||||
not isinstance(parsed_frame, dict)
|
||||
or parsed_frame.get("type") != "response.create"
|
||||
):
|
||||
return raw_msg
|
||||
wrapped_frame = isinstance(parsed_frame.get("response"), dict)
|
||||
inner_payload = (
|
||||
parsed_frame["response"] if wrapped_frame else parsed_frame
|
||||
)
|
||||
if not isinstance(inner_payload, dict):
|
||||
return raw_msg
|
||||
try:
|
||||
from headroom._core import (
|
||||
compress_openai_responses_live_zone as _ws_frame_compress,
|
||||
)
|
||||
|
||||
inner_bytes = json.dumps(inner_payload).encode("utf-8")
|
||||
model_for_frame = inner_payload.get("model") or ""
|
||||
_frame_auth_mode = classify_auth_mode(ws_headers)
|
||||
(
|
||||
new_inner_bytes,
|
||||
modified,
|
||||
rust_saved,
|
||||
rust_transforms,
|
||||
) = _ws_frame_compress(
|
||||
inner_bytes,
|
||||
_frame_auth_mode.value,
|
||||
model_for_frame,
|
||||
)
|
||||
except Exception as _frame_err:
|
||||
logger.warning(
|
||||
"[%s] WS /v1/responses frame compression "
|
||||
"failed; forwarding original: %s: %s",
|
||||
request_id,
|
||||
type(_frame_err).__name__,
|
||||
_frame_err,
|
||||
)
|
||||
return raw_msg
|
||||
if not modified:
|
||||
return raw_msg
|
||||
try:
|
||||
new_inner = json.loads(new_inner_bytes)
|
||||
except json.JSONDecodeError:
|
||||
return raw_msg
|
||||
if not isinstance(new_inner, dict):
|
||||
return raw_msg
|
||||
if wrapped_frame:
|
||||
parsed_frame["response"] = new_inner
|
||||
rewritten = json.dumps(parsed_frame)
|
||||
else:
|
||||
rewritten = json.dumps(new_inner)
|
||||
tokens_saved += int(rust_saved)
|
||||
for t in (
|
||||
"openai_responses_ws_live_zone",
|
||||
*list(rust_transforms),
|
||||
):
|
||||
if t not in transforms_applied:
|
||||
transforms_applied.append(t)
|
||||
ws_frames_compressed += 1
|
||||
logger.info(
|
||||
"[%s] WS /v1/responses frame compressed "
|
||||
"%d→%d bytes (%d tokens saved, "
|
||||
"auth_mode=%s, frame=%d)",
|
||||
request_id,
|
||||
len(inner_bytes),
|
||||
len(new_inner_bytes),
|
||||
int(rust_saved),
|
||||
_frame_auth_mode.value,
|
||||
ws_frames_compressed,
|
||||
)
|
||||
return rewritten
|
||||
|
||||
async def _client_to_upstream() -> None:
|
||||
nonlocal client_relay_error
|
||||
try:
|
||||
while True:
|
||||
msg = await websocket.receive_text()
|
||||
msg = await _maybe_compress_response_create_frame(msg)
|
||||
await upstream.send(msg)
|
||||
except asyncio.CancelledError:
|
||||
# Explicit cancel from the outer
|
||||
|
|
@ -2577,16 +2802,80 @@ class OpenAIHandlerMixin:
|
|||
websocket, body, first_msg_raw, upstream_headers, request_id
|
||||
)
|
||||
|
||||
# Record metrics
|
||||
if tokens_saved > 0:
|
||||
model_name = body.get("model", "unknown") if isinstance(body, dict) else "unknown"
|
||||
await self.metrics.record_request(
|
||||
provider="openai",
|
||||
model=model_name,
|
||||
input_tokens=0,
|
||||
output_tokens=0,
|
||||
tokens_saved=tokens_saved,
|
||||
latency_ms=0,
|
||||
# ── WS session-end metric + RequestLog ──────────────────
|
||||
#
|
||||
# Unconditional (was previously gated on `tokens_saved>0`,
|
||||
# which made first-frame no-changes invisible). We record
|
||||
# one entry per WS session that aggregates `tokens_saved`
|
||||
# across every `response.create` frame compressed by the
|
||||
# first-frame block + `_maybe_compress_response_create_frame`.
|
||||
# The RequestLog entry mirrors the streaming.py /
|
||||
# anthropic.py shape so /transformations/feed surfaces
|
||||
# Codex WS turns.
|
||||
ws_session_duration_ms = (time.perf_counter() - session_started_at) * 1000.0
|
||||
ws_inner_for_telemetry: dict[str, Any] = (
|
||||
body.get("response", body) if isinstance(body, dict) else {}
|
||||
)
|
||||
if not isinstance(ws_inner_for_telemetry, dict):
|
||||
ws_inner_for_telemetry = {}
|
||||
model_name = (
|
||||
ws_inner_for_telemetry.get("model")
|
||||
or (body.get("model") if isinstance(body, dict) else None)
|
||||
or "unknown"
|
||||
)
|
||||
_final_auth_mode = classify_auth_mode(ws_headers)
|
||||
ws_session_tags = {
|
||||
**(ws_tags or {}),
|
||||
"auth_mode": _final_auth_mode.value,
|
||||
"endpoint": "responses_ws",
|
||||
"ws_frames_compressed": str(ws_frames_compressed),
|
||||
}
|
||||
await self.metrics.record_request(
|
||||
provider="openai",
|
||||
model=model_name,
|
||||
input_tokens=0,
|
||||
output_tokens=0,
|
||||
tokens_saved=tokens_saved,
|
||||
latency_ms=ws_session_duration_ms,
|
||||
)
|
||||
if getattr(self, "logger", None) is not None:
|
||||
from headroom.proxy.helpers import compute_turn_id
|
||||
from headroom.proxy.models import RequestLog
|
||||
|
||||
ws_messages_for_log: list[dict[str, Any]] = []
|
||||
ws_input_for_log = ws_inner_for_telemetry.get("input")
|
||||
ws_instructions_for_log = ws_inner_for_telemetry.get("instructions")
|
||||
if isinstance(ws_instructions_for_log, str) and ws_instructions_for_log:
|
||||
ws_messages_for_log.append(
|
||||
{"role": "system", "content": ws_instructions_for_log}
|
||||
)
|
||||
if isinstance(ws_input_for_log, str) and ws_input_for_log:
|
||||
ws_messages_for_log.append({"role": "user", "content": ws_input_for_log})
|
||||
self.logger.log(
|
||||
RequestLog(
|
||||
request_id=request_id,
|
||||
timestamp=datetime.now().isoformat(),
|
||||
provider="openai",
|
||||
model=model_name,
|
||||
input_tokens_original=0,
|
||||
input_tokens_optimized=0,
|
||||
output_tokens=0,
|
||||
tokens_saved=tokens_saved,
|
||||
savings_percent=0.0,
|
||||
optimization_latency_ms=0.0,
|
||||
total_latency_ms=ws_session_duration_ms,
|
||||
tags=ws_session_tags,
|
||||
cache_hit=False,
|
||||
transforms_applied=transforms_applied,
|
||||
request_messages=ws_messages_for_log
|
||||
if getattr(self.config, "log_full_messages", False)
|
||||
else None,
|
||||
turn_id=compute_turn_id(
|
||||
model_name,
|
||||
ws_instructions_for_log,
|
||||
ws_messages_for_log,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
{
|
||||
"name": "headroom",
|
||||
"version": "0.21.4",
|
||||
"version": "0.21.5",
|
||||
"description": "Headroom startup hooks for Claude Code and GitHub Copilot CLI.",
|
||||
"author": {
|
||||
"name": "Headroom Contributors",
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
{
|
||||
"name": "headroom",
|
||||
"version": "0.21.4",
|
||||
"version": "0.21.5",
|
||||
"description": "Headroom startup hooks for Claude Code and GitHub Copilot CLI.",
|
||||
"author": {
|
||||
"name": "Headroom Contributors",
|
||||
|
|
|
|||
390
tests/e2e_real_compression.py
Normal file
390
tests/e2e_real_compression.py
Normal file
|
|
@ -0,0 +1,390 @@
|
|||
"""End-to-end compression verification with realistic multi-turn payloads.
|
||||
|
||||
Headroom only compresses content the model has already seen — assistant
|
||||
turns, tool results, and Responses-API output items. Fresh user prompts
|
||||
are *intentionally* skipped (the model needs them verbatim, and they're
|
||||
in the live-zone tail anyway). A conversation that contains nothing but
|
||||
a single user prompt will produce 0 tokens saved by design — that is
|
||||
not a bug; it's the live-zone-only invariant.
|
||||
|
||||
This script exercises every (provider × endpoint × streaming) combination
|
||||
with a payload large enough to trigger compression. Pass criteria:
|
||||
|
||||
* tokens_saved > 0 for at least one chat-completions case
|
||||
* tokens_saved > 0 for at least one /v1/messages case
|
||||
* tokens_saved > 0 for the /v1/responses case
|
||||
* tokens_saved > 0 for streaming variants
|
||||
* No proxy errors, no compression-failed warnings on happy paths
|
||||
|
||||
Reads keys from .env. Run via:
|
||||
|
||||
.venv/bin/python tests/e2e_real_compression.py
|
||||
|
||||
# Note on auth-header construction
|
||||
# The API keys are read from `os.environ` *inside* `_post` and never
|
||||
# stored as local variables in the test runner's main scope. This
|
||||
# breaks the CodeQL taint flow that would otherwise flag every
|
||||
# diagnostic `print()` in the loop as
|
||||
# `py/clear-text-logging-sensitive-data` because credentials live in
|
||||
# the same scope.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import socket
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parent.parent
|
||||
|
||||
|
||||
def load_env_into_environ() -> None:
|
||||
"""Read REPO_ROOT/.env and merge into os.environ. Keys are never
|
||||
returned to the caller — see module docstring."""
|
||||
p = REPO_ROOT / ".env"
|
||||
if not p.exists():
|
||||
return
|
||||
for line in p.read_text().splitlines():
|
||||
line = line.strip()
|
||||
if not line or line.startswith("#") or "=" not in line:
|
||||
continue
|
||||
k, v = line.split("=", 1)
|
||||
k = k.strip()
|
||||
v = v.strip().strip('"').strip("'")
|
||||
os.environ.setdefault(k, v)
|
||||
|
||||
|
||||
def have_required_keys() -> tuple[bool, str]:
|
||||
"""Sentinel check without exposing the keys themselves to local scope."""
|
||||
missing = [n for n in ("OPENAI_API_KEY", "ANTHROPIC_API_KEY") if not os.environ.get(n)]
|
||||
if missing:
|
||||
return False, ", ".join(missing)
|
||||
return True, ""
|
||||
|
||||
|
||||
def free_port() -> int:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
s.bind(("127.0.0.1", 0))
|
||||
return s.getsockname()[1]
|
||||
|
||||
|
||||
def wait_ready(port: int, timeout_s: float = 60.0) -> None:
|
||||
deadline = time.time() + timeout_s
|
||||
while time.time() < deadline:
|
||||
try:
|
||||
with urllib.request.urlopen(f"http://127.0.0.1:{port}/livez", timeout=2) as r:
|
||||
if r.status == 200:
|
||||
return
|
||||
except Exception:
|
||||
time.sleep(0.5)
|
||||
raise TimeoutError("proxy not ready")
|
||||
|
||||
|
||||
def _post(url: str, body: dict, *, provider: str, stream: bool = False) -> tuple[int, Any]:
|
||||
"""Make a POST request, building auth headers from os.environ at
|
||||
call time. The credential never appears in the caller's local
|
||||
scope, which keeps CodeQL's taint analysis happy."""
|
||||
if provider == "openai":
|
||||
headers = {
|
||||
"Authorization": "Bearer " + (os.environ.get("OPENAI_API_KEY") or ""),
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
elif provider == "anthropic":
|
||||
headers = {
|
||||
"x-api-key": os.environ.get("ANTHROPIC_API_KEY") or "",
|
||||
"anthropic-version": "2023-06-01",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
else:
|
||||
raise ValueError(f"unknown provider: {provider!r}")
|
||||
|
||||
req = urllib.request.Request(
|
||||
url, data=json.dumps(body).encode("utf-8"), headers=headers, method="POST"
|
||||
)
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=120) as r:
|
||||
raw = r.read()
|
||||
if stream:
|
||||
return r.status, raw.decode("utf-8", errors="replace")
|
||||
try:
|
||||
return r.status, json.loads(raw)
|
||||
except json.JSONDecodeError:
|
||||
return r.status, raw.decode("utf-8", errors="replace")
|
||||
except urllib.error.HTTPError as e:
|
||||
try:
|
||||
return e.code, json.loads(e.read())
|
||||
except Exception:
|
||||
return e.code, str(e)
|
||||
|
||||
|
||||
# ── Payload builders ──────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def long_build_log() -> str:
|
||||
"""~24 KB of structured BuildOutput-style content. The Rust
|
||||
LogCompressor recognizes this and compresses aggressively."""
|
||||
return "".join(
|
||||
f"[2024-01-01 00:00:{i % 60:02d}] INFO compile.rs:42 building module foo_{i} "
|
||||
f"(crate=workspace-{i // 10}, deps=[serde={i}, tokio={i}, regex={i % 7}])\n"
|
||||
for i in range(400)
|
||||
)
|
||||
|
||||
|
||||
def anthropic_messages_payload(streaming: bool = False) -> dict:
|
||||
return {
|
||||
"model": "claude-haiku-4-5-20251001",
|
||||
"max_tokens": 30,
|
||||
"stream": streaming,
|
||||
"tools": [
|
||||
{
|
||||
"name": "shell",
|
||||
"description": "Run a shell command",
|
||||
"input_schema": {
|
||||
"type": "object",
|
||||
"properties": {"command": {"type": "string"}},
|
||||
"required": ["command"],
|
||||
},
|
||||
}
|
||||
],
|
||||
"messages": [
|
||||
{"role": "user", "content": "Run cargo build and tell me if it succeeded."},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "text", "text": "Running it now."},
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": "toolu_e2e_1",
|
||||
"name": "shell",
|
||||
"input": {"command": "cargo build --release"},
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_e2e_1",
|
||||
"content": long_build_log(),
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "user", "content": "One word: pass or fail?"},
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def openai_chat_payload(streaming: bool = False) -> dict:
|
||||
return {
|
||||
"model": "gpt-4o-mini",
|
||||
"max_tokens": 30,
|
||||
"stream": streaming,
|
||||
"messages": [
|
||||
{"role": "user", "content": "Run cargo build and report if it succeeded."},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_e2e_1",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "shell",
|
||||
"arguments": '{"command": "cargo build --release"}',
|
||||
},
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": "call_e2e_1",
|
||||
"content": long_build_log(),
|
||||
},
|
||||
{"role": "user", "content": "One word: pass or fail?"},
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def openai_responses_payload(streaming: bool = False) -> dict:
|
||||
return {
|
||||
"model": "gpt-4o-mini",
|
||||
"max_output_tokens": 30,
|
||||
"stream": streaming,
|
||||
"input": [
|
||||
{
|
||||
"type": "message",
|
||||
"role": "user",
|
||||
"content": [{"type": "input_text", "text": "Run cargo build and report."}],
|
||||
},
|
||||
{
|
||||
"type": "function_call",
|
||||
"call_id": "call_resp_e2e_1",
|
||||
"name": "shell",
|
||||
"arguments": '{"command": "cargo build --release"}',
|
||||
},
|
||||
{
|
||||
"type": "function_call_output",
|
||||
"call_id": "call_resp_e2e_1",
|
||||
"output": long_build_log(),
|
||||
},
|
||||
{
|
||||
"type": "message",
|
||||
"role": "user",
|
||||
"content": [{"type": "input_text", "text": "One word: pass or fail?"}],
|
||||
},
|
||||
],
|
||||
"instructions": "You read shell output and reply tersely.",
|
||||
}
|
||||
|
||||
|
||||
# ── Test runner ───────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def main() -> int:
|
||||
load_env_into_environ()
|
||||
ok, missing = have_required_keys()
|
||||
if not ok:
|
||||
print(f"FAIL: missing keys: {missing}", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
port = free_port()
|
||||
print(f"[e2e] starting proxy on :{port}")
|
||||
log_fp = open("/tmp/e2e_real_proxy.log", "w")
|
||||
proc = subprocess.Popen(
|
||||
[
|
||||
str(REPO_ROOT / ".venv/bin/headroom"),
|
||||
"proxy",
|
||||
"--port",
|
||||
str(port),
|
||||
"--no-telemetry",
|
||||
],
|
||||
env={**os.environ, "HEADROOM_REQUIRE_RUST_CORE": "true"},
|
||||
stdout=log_fp,
|
||||
stderr=subprocess.STDOUT,
|
||||
cwd=str(REPO_ROOT),
|
||||
)
|
||||
|
||||
failures: list[str] = []
|
||||
try:
|
||||
wait_ready(port)
|
||||
print("[e2e] proxy ready")
|
||||
|
||||
# Cases carry only structural info: name, path, provider tag,
|
||||
# body, stream. Auth headers are built inside `_post` from
|
||||
# os.environ — see module docstring.
|
||||
cases: list[tuple[str, str, str, dict, bool]] = [
|
||||
(
|
||||
"anthropic_messages_nonstream",
|
||||
"/v1/messages",
|
||||
"anthropic",
|
||||
anthropic_messages_payload(streaming=False),
|
||||
False,
|
||||
),
|
||||
(
|
||||
"anthropic_messages_stream",
|
||||
"/v1/messages",
|
||||
"anthropic",
|
||||
anthropic_messages_payload(streaming=True),
|
||||
True,
|
||||
),
|
||||
(
|
||||
"openai_chat_nonstream",
|
||||
"/v1/chat/completions",
|
||||
"openai",
|
||||
openai_chat_payload(streaming=False),
|
||||
False,
|
||||
),
|
||||
(
|
||||
"openai_chat_stream",
|
||||
"/v1/chat/completions",
|
||||
"openai",
|
||||
openai_chat_payload(streaming=True),
|
||||
True,
|
||||
),
|
||||
(
|
||||
"openai_responses_nonstream",
|
||||
"/v1/responses",
|
||||
"openai",
|
||||
openai_responses_payload(streaming=False),
|
||||
False,
|
||||
),
|
||||
]
|
||||
|
||||
for name, path, provider, body, stream in cases:
|
||||
url = f"http://127.0.0.1:{port}{path}"
|
||||
print(f"[e2e] {name}: POST {path}")
|
||||
status, _ = _post(url, body, provider=provider, stream=stream)
|
||||
if status != 200:
|
||||
failures.append(f"{name}: HTTP {status}")
|
||||
continue
|
||||
print(" ok status=200")
|
||||
|
||||
# ── Scrape proxy log for compression evidence ────────────
|
||||
time.sleep(1.5)
|
||||
canonical_log = Path.home() / ".headroom" / "logs" / "proxy.log"
|
||||
if canonical_log.exists():
|
||||
log_lines = canonical_log.read_text(errors="replace").splitlines()[-5000:]
|
||||
else:
|
||||
log_lines = Path("/tmp/e2e_real_proxy.log").read_text(errors="replace").splitlines()
|
||||
|
||||
compressed_evidence = [
|
||||
line
|
||||
for line in log_lines
|
||||
if "compressed" in line and ("tokens" in line.lower() or "bytes" in line.lower())
|
||||
][-30:]
|
||||
if compressed_evidence:
|
||||
print("\n[e2e] compression evidence (last 10 lines):")
|
||||
for line in compressed_evidence[-10:]:
|
||||
idx = line.find("] ")
|
||||
print(" ", line[idx + 2 :] if idx > 0 else line)
|
||||
else:
|
||||
print("\n[e2e] no compression evidence in canonical log")
|
||||
|
||||
saved_pattern = re.compile(r"saved (\d[\d,]*) tokens?", re.IGNORECASE)
|
||||
total_saved = 0
|
||||
for line in log_lines[-2000:]:
|
||||
m = saved_pattern.search(line)
|
||||
if m:
|
||||
num = int(m.group(1).replace(",", ""))
|
||||
if num > 0:
|
||||
total_saved += num
|
||||
|
||||
print(
|
||||
f"\n[e2e] aggregate tokens saved across cases (~last 2000 log lines): {total_saved:,}"
|
||||
)
|
||||
if total_saved == 0:
|
||||
failures.append(f"no compression evidence — check {canonical_log}")
|
||||
|
||||
joined = "\n".join(log_lines)
|
||||
if "compression failed" in joined:
|
||||
failures.append("proxy log contains 'compression failed' — see canonical log")
|
||||
|
||||
finally:
|
||||
print("\n[e2e] terminating proxy")
|
||||
proc.terminate()
|
||||
try:
|
||||
proc.wait(timeout=10)
|
||||
except subprocess.TimeoutExpired:
|
||||
proc.kill()
|
||||
log_fp.close()
|
||||
|
||||
if failures:
|
||||
print("\n=== E2E FAILURES ===")
|
||||
for f in failures:
|
||||
print(" -", f)
|
||||
return 1
|
||||
print("\n=== E2E ALL GREEN ===")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
279
tests/e2e_ws_responses_compression.py
Normal file
279
tests/e2e_ws_responses_compression.py
Normal file
|
|
@ -0,0 +1,279 @@
|
|||
"""End-to-end verification that /v1/responses WebSocket compression fires.
|
||||
|
||||
We can't reach OpenAI's WS endpoint without the `responses_websockets`
|
||||
beta enabled on the test API key, so this test does the next-best
|
||||
thing: it spins up a *fake upstream* WebSocket server, points the
|
||||
proxy at it via OPENAI_API_URL, and connects a client to the proxy.
|
||||
|
||||
Verifies:
|
||||
1. First-frame compression: the client sends a `response.create`
|
||||
event with a 24 KB output_item; the fake upstream receives the
|
||||
COMPRESSED frame (much smaller than what was sent).
|
||||
2. Multi-frame compression: a second `response.create` on the same
|
||||
WS session is also compressed (the new behavior — was previously
|
||||
first-frame-only).
|
||||
3. Other event types (e.g. `response.cancel`) pass through
|
||||
unchanged.
|
||||
4. Proxy log surfaces both compression events with token-saved
|
||||
numbers.
|
||||
|
||||
Run via:
|
||||
|
||||
.venv/bin/python tests/e2e_ws_responses_compression.py
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import socket
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
|
||||
import websockets
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parent.parent
|
||||
|
||||
|
||||
def free_port() -> int:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
s.bind(("127.0.0.1", 0))
|
||||
return s.getsockname()[1]
|
||||
|
||||
|
||||
def wait_ready(port: int, timeout_s: float = 60.0) -> None:
|
||||
deadline = time.time() + timeout_s
|
||||
while time.time() < deadline:
|
||||
try:
|
||||
with urllib.request.urlopen(f"http://127.0.0.1:{port}/livez", timeout=2) as r:
|
||||
if r.status == 200:
|
||||
return
|
||||
except Exception:
|
||||
time.sleep(0.5)
|
||||
raise TimeoutError("proxy not ready")
|
||||
|
||||
|
||||
def long_build_log() -> str:
|
||||
return "".join(
|
||||
f"[2024-01-01 00:00:{i % 60:02d}] INFO compile.rs:42 building module foo_{i} "
|
||||
f"(crate=workspace-{i // 10}, deps=[serde={i}, tokio={i}])\n"
|
||||
for i in range(400)
|
||||
)
|
||||
|
||||
|
||||
def make_response_create_payload(turn_no: int) -> dict:
|
||||
"""A wire-shape `response.create` envelope, with a long
|
||||
function_call_output that the dispatcher will compress."""
|
||||
return {
|
||||
"type": "response.create",
|
||||
"response": {
|
||||
"model": "gpt-4o-mini",
|
||||
"input": [
|
||||
{
|
||||
"type": "message",
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "input_text",
|
||||
"text": f"Turn {turn_no} — please summarize the build output.",
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
"type": "function_call",
|
||||
"call_id": f"call_e2e_ws_{turn_no}",
|
||||
"name": "shell",
|
||||
"arguments": '{"command": "cargo build --release"}',
|
||||
},
|
||||
{
|
||||
"type": "function_call_output",
|
||||
"call_id": f"call_e2e_ws_{turn_no}",
|
||||
"output": long_build_log(),
|
||||
},
|
||||
],
|
||||
"instructions": "You read shell output and reply tersely.",
|
||||
"max_output_tokens": 30,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
# ── Fake upstream WS server ─────────────────────────────────────────────
|
||||
#
|
||||
# Captures every text frame the proxy forwards so we can assert what
|
||||
# arrived upstream actually got compressed.
|
||||
|
||||
|
||||
class FakeUpstream:
|
||||
def __init__(self) -> None:
|
||||
self.received_frames: list[str] = []
|
||||
self.server: websockets.server.WebSocketServer | None = None
|
||||
self.port: int = 0
|
||||
|
||||
async def _handler(self, ws):
|
||||
try:
|
||||
async for msg in ws:
|
||||
if isinstance(msg, str):
|
||||
self.received_frames.append(msg)
|
||||
# Echo a minimal completion event so the proxy doesn't
|
||||
# think upstream is hung.
|
||||
await ws.send(
|
||||
json.dumps(
|
||||
{
|
||||
"type": "response.completed",
|
||||
"response": {"id": "fake_resp", "output": []},
|
||||
}
|
||||
)
|
||||
)
|
||||
except websockets.exceptions.ConnectionClosed:
|
||||
pass
|
||||
|
||||
async def start(self) -> int:
|
||||
self.port = free_port()
|
||||
self.server = await websockets.serve(self._handler, "127.0.0.1", self.port)
|
||||
return self.port
|
||||
|
||||
async def stop(self) -> None:
|
||||
if self.server:
|
||||
self.server.close()
|
||||
await self.server.wait_closed()
|
||||
|
||||
|
||||
async def main_async() -> int:
|
||||
fake = FakeUpstream()
|
||||
upstream_port = await fake.start()
|
||||
upstream_url = f"http://127.0.0.1:{upstream_port}"
|
||||
|
||||
proxy_port = free_port()
|
||||
print(f"[ws-e2e] fake upstream at ws://127.0.0.1:{upstream_port}")
|
||||
print(f"[ws-e2e] starting proxy on :{proxy_port}")
|
||||
|
||||
log_fp = open("/tmp/e2e_ws_proxy.log", "w")
|
||||
proc = subprocess.Popen(
|
||||
[
|
||||
str(REPO_ROOT / ".venv/bin/headroom"),
|
||||
"proxy",
|
||||
"--port",
|
||||
str(proxy_port),
|
||||
"--no-telemetry",
|
||||
# Point /v1/responses upstream at our fake server instead
|
||||
# of api.openai.com.
|
||||
"--openai-api-url",
|
||||
upstream_url,
|
||||
],
|
||||
env={
|
||||
**os.environ,
|
||||
# Need *some* OpenAI key value so the proxy doesn't refuse;
|
||||
# the fake upstream ignores it.
|
||||
"OPENAI_API_KEY": os.environ.get("OPENAI_API_KEY", "sk-fake-for-test"),
|
||||
"ANTHROPIC_API_KEY": os.environ.get("ANTHROPIC_API_KEY", "sk-ant-fake"),
|
||||
"HEADROOM_REQUIRE_RUST_CORE": "true",
|
||||
},
|
||||
stdout=log_fp,
|
||||
stderr=subprocess.STDOUT,
|
||||
cwd=str(REPO_ROOT),
|
||||
)
|
||||
|
||||
failures: list[str] = []
|
||||
try:
|
||||
wait_ready(proxy_port)
|
||||
print("[ws-e2e] proxy ready")
|
||||
|
||||
# ── Connect WS client to the proxy ───────────────────────
|
||||
proxy_ws_url = f"ws://127.0.0.1:{proxy_port}/v1/responses"
|
||||
async with websockets.connect(
|
||||
proxy_ws_url,
|
||||
additional_headers={
|
||||
"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY', 'sk-fake')}",
|
||||
"OpenAI-Beta": "responses_websockets=2026-02-06",
|
||||
},
|
||||
) as ws:
|
||||
# Frame 1: response.create with large content
|
||||
payload_1 = make_response_create_payload(1)
|
||||
payload_1_bytes = len(json.dumps(payload_1).encode("utf-8"))
|
||||
print(f"[ws-e2e] sending frame 1 ({payload_1_bytes:,} bytes)")
|
||||
await ws.send(json.dumps(payload_1))
|
||||
|
||||
# Wait for the fake upstream to receive (or timeout)
|
||||
await asyncio.sleep(2.0)
|
||||
|
||||
# Frame 2: ANOTHER response.create on the same session
|
||||
payload_2 = make_response_create_payload(2)
|
||||
payload_2_bytes = len(json.dumps(payload_2).encode("utf-8"))
|
||||
print(f"[ws-e2e] sending frame 2 ({payload_2_bytes:,} bytes)")
|
||||
await ws.send(json.dumps(payload_2))
|
||||
await asyncio.sleep(2.0)
|
||||
|
||||
# Frame 3: a non-response.create event — should pass through
|
||||
cancel = {"type": "response.cancel"}
|
||||
print("[ws-e2e] sending frame 3 (response.cancel — passthrough)")
|
||||
await ws.send(json.dumps(cancel))
|
||||
await asyncio.sleep(1.0)
|
||||
|
||||
# ── Inspect what arrived at the fake upstream ────────────
|
||||
print(f"\n[ws-e2e] fake upstream received {len(fake.received_frames)} frames")
|
||||
for i, frame in enumerate(fake.received_frames):
|
||||
print(f" frame {i + 1}: {len(frame.encode()):,} bytes")
|
||||
|
||||
# First two frames should be MUCH smaller than what we sent.
|
||||
# The third (response.cancel) should be a small fixed size.
|
||||
if len(fake.received_frames) < 2:
|
||||
failures.append(f"expected ≥2 frames at upstream, got {len(fake.received_frames)}")
|
||||
else:
|
||||
f1 = len(fake.received_frames[0].encode())
|
||||
f2 = len(fake.received_frames[1].encode())
|
||||
if f1 >= payload_1_bytes // 2:
|
||||
failures.append(
|
||||
f"frame 1 not compressed: arrived {f1:,} bytes (sent {payload_1_bytes:,})"
|
||||
)
|
||||
if f2 >= payload_2_bytes // 2:
|
||||
failures.append(
|
||||
f"frame 2 not compressed (multi-frame regression): "
|
||||
f"arrived {f2:,} bytes (sent {payload_2_bytes:,})"
|
||||
)
|
||||
|
||||
# ── Scrape proxy log for compression evidence ────────────
|
||||
await asyncio.sleep(1.0)
|
||||
canonical = Path.home() / ".headroom" / "logs" / "proxy.log"
|
||||
log_lines = canonical.read_text(errors="replace").splitlines()[-1000:]
|
||||
ws_compressed = [
|
||||
line for line in log_lines if "WS /v1/responses" in line and "compressed" in line
|
||||
]
|
||||
print("\n[ws-e2e] WS compression log lines (last few):")
|
||||
for line in ws_compressed[-6:]:
|
||||
idx = line.find("] ")
|
||||
print(" ", line[idx + 2 :] if idx > 0 else line)
|
||||
if len(ws_compressed) < 2:
|
||||
failures.append(
|
||||
f"expected ≥2 WS compression log entries (first frame + multi-frame), "
|
||||
f"saw {len(ws_compressed)}"
|
||||
)
|
||||
|
||||
finally:
|
||||
print("\n[ws-e2e] terminating proxy")
|
||||
proc.terminate()
|
||||
try:
|
||||
proc.wait(timeout=10)
|
||||
except subprocess.TimeoutExpired:
|
||||
proc.kill()
|
||||
log_fp.close()
|
||||
await fake.stop()
|
||||
|
||||
if failures:
|
||||
print("\n=== WS E2E FAILURES ===")
|
||||
for f in failures:
|
||||
print(" -", f)
|
||||
return 1
|
||||
print("\n=== WS E2E ALL GREEN ===")
|
||||
return 0
|
||||
|
||||
|
||||
def main() -> int:
|
||||
return asyncio.run(main_async())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
|
|
@ -113,15 +113,20 @@ class TestCompressionCacheFrozenCount:
|
|||
def test_empty_cache_returns_zero(self, cache: CompressionCache) -> None:
|
||||
assert cache.compute_frozen_count([]) == 0
|
||||
|
||||
def test_user_assistant_always_stable(self, cache: CompressionCache) -> None:
|
||||
def test_user_assistant_stable_with_live_zone_cap(self, cache: CompressionCache) -> None:
|
||||
"""Plain user/assistant turns are individually stable, but the
|
||||
trailing message is reserved as the live zone — the new turn
|
||||
cannot be in any provider prefix cache. See docstring on
|
||||
``CompressionCache.compute_frozen_count``."""
|
||||
messages = [
|
||||
{"role": "user", "content": "hello"},
|
||||
{"role": "assistant", "content": "hi there"},
|
||||
{"role": "user", "content": "how are you"},
|
||||
]
|
||||
assert cache.compute_frozen_count(messages) == 3
|
||||
# 3 messages structurally stable; cap clamps to len-1 = 2.
|
||||
assert cache.compute_frozen_count(messages) == 2
|
||||
|
||||
def test_tool_result_with_cache_hit_is_stable(self, cache: CompressionCache) -> None:
|
||||
def test_tool_result_with_cache_hit_capped_at_live_zone(self, cache: CompressionCache) -> None:
|
||||
tool_content = "tool output data"
|
||||
h = CompressionCache.content_hash(tool_content)
|
||||
cache.store_compressed(h, "compressed tool output", tokens_saved=5)
|
||||
|
|
@ -137,7 +142,9 @@ class TestCompressionCacheFrozenCount:
|
|||
"content": [{"type": "tool_result", "tool_use_id": "t1", "content": tool_content}],
|
||||
},
|
||||
]
|
||||
assert cache.compute_frozen_count(messages) == 3
|
||||
# All 3 stable; cap clamps to len-1 = 2 (trailing tool_result is
|
||||
# the live zone).
|
||||
assert cache.compute_frozen_count(messages) == 2
|
||||
|
||||
def test_tool_result_cache_miss_stops_frozen(self, cache: CompressionCache) -> None:
|
||||
messages = [
|
||||
|
|
@ -188,9 +195,10 @@ class TestCompressionCacheFrozenCount:
|
|||
},
|
||||
{"role": "user", "content": "follow up"},
|
||||
]
|
||||
# Without mark_stable, this would stop at msg[1] → frozen=1.
|
||||
# With stable hash, the walk continues past msg[1] → frozen=3.
|
||||
assert cache.compute_frozen_count(messages) == 3
|
||||
# Without mark_stable, the walk would stop at msg[1] → frozen=1.
|
||||
# With stable hash, the walk continues past msg[1]; structural
|
||||
# count = 3, then capped at len-1 = 2 (live-zone reservation).
|
||||
assert cache.compute_frozen_count(messages) == 2
|
||||
|
||||
def test_update_from_result_identical_content_marks_stable(
|
||||
self, cache: CompressionCache
|
||||
|
|
@ -217,7 +225,8 @@ class TestCompressionCacheFrozenCount:
|
|||
h = CompressionCache.content_hash(tool_content)
|
||||
assert h in cache._stable_hashes
|
||||
|
||||
# Frozen count should now walk past this tool_result
|
||||
# Frozen count walks past this tool_result (its hash is stable),
|
||||
# but the trailing message is still reserved as live zone.
|
||||
messages = [
|
||||
{"role": "user", "content": "hello"},
|
||||
{
|
||||
|
|
@ -226,7 +235,7 @@ class TestCompressionCacheFrozenCount:
|
|||
},
|
||||
{"role": "user", "content": "more stuff"},
|
||||
]
|
||||
assert cache.compute_frozen_count(messages) == 3
|
||||
assert cache.compute_frozen_count(messages) == 2
|
||||
|
||||
def test_mark_stable_from_messages(self, cache: CompressionCache) -> None:
|
||||
"""mark_stable_from_messages records hashes for tool_results."""
|
||||
|
|
|
|||
|
|
@ -48,21 +48,21 @@ class TestPassthroughCases:
|
|||
def test_not_json_passthrough(self):
|
||||
compress = _ensure_binding()
|
||||
body = b"this is not JSON at all"
|
||||
out, modified = compress(body, "payg", "gpt-4o-mini")
|
||||
out, modified, _saved, _transforms = compress(body, "payg", "gpt-4o-mini")
|
||||
assert out == body
|
||||
assert modified is False
|
||||
|
||||
def test_no_input_array_passthrough(self):
|
||||
compress = _ensure_binding()
|
||||
body = json.dumps({"model": "gpt-4o-mini"}).encode()
|
||||
out, modified = compress(body, "payg", "gpt-4o-mini")
|
||||
out, modified, _saved, _transforms = compress(body, "payg", "gpt-4o-mini")
|
||||
assert out == body
|
||||
assert modified is False
|
||||
|
||||
def test_empty_input_array_passthrough(self):
|
||||
compress = _ensure_binding()
|
||||
body = json.dumps({"model": "gpt-4o-mini", "input": []}).encode()
|
||||
out, modified = compress(body, "payg", "gpt-4o-mini")
|
||||
out, modified, _saved, _transforms = compress(body, "payg", "gpt-4o-mini")
|
||||
assert out == body
|
||||
assert modified is False
|
||||
|
||||
|
|
@ -76,7 +76,7 @@ class TestPassthroughCases:
|
|||
"input": [{"type": "message", "role": "user", "content": "hi"}],
|
||||
}
|
||||
).encode()
|
||||
out, modified = compress(body, "payg", "gpt-4o-mini")
|
||||
out, modified, _saved, _transforms = compress(body, "payg", "gpt-4o-mini")
|
||||
assert modified is False
|
||||
# Body should be byte-equal (passthrough, not re-serialized).
|
||||
assert out == body
|
||||
|
|
@ -94,7 +94,7 @@ class TestAuthModeAccepted:
|
|||
compress = _ensure_binding()
|
||||
body = json.dumps({"model": "gpt-4o-mini", "input": []}).encode()
|
||||
# Should not raise on any string input.
|
||||
out, modified = compress(body, auth_mode, "gpt-4o-mini")
|
||||
out, modified, _saved, _transforms = compress(body, auth_mode, "gpt-4o-mini")
|
||||
assert isinstance(out, bytes)
|
||||
assert modified is False
|
||||
|
||||
|
|
@ -105,7 +105,7 @@ class TestModelDefault:
|
|||
def test_empty_model_uses_default(self):
|
||||
compress = _ensure_binding()
|
||||
body = json.dumps({"input": []}).encode()
|
||||
out, modified = compress(body, "payg", "")
|
||||
out, modified, _saved, _transforms = compress(body, "payg", "")
|
||||
assert isinstance(out, bytes)
|
||||
assert modified is False
|
||||
|
||||
|
|
@ -118,12 +118,78 @@ class TestNoExceptionsLeak:
|
|||
|
||||
def test_garbage_bytes_no_raise(self):
|
||||
compress = _ensure_binding()
|
||||
out, modified = compress(b"\xff\xfe\x00\xff", "payg", "gpt-4o-mini")
|
||||
out, modified, _saved, _transforms = compress(b"\xff\xfe\x00\xff", "payg", "gpt-4o-mini")
|
||||
assert modified is False
|
||||
assert out == b"\xff\xfe\x00\xff"
|
||||
|
||||
def test_empty_body_no_raise(self):
|
||||
compress = _ensure_binding()
|
||||
out, modified = compress(b"", "payg", "gpt-4o-mini")
|
||||
out, modified, _saved, _transforms = compress(b"", "payg", "gpt-4o-mini")
|
||||
assert modified is False
|
||||
assert out == b""
|
||||
|
||||
|
||||
class TestTelemetryFields:
|
||||
"""The 4-tuple return surfaces ``tokens_saved`` (sum of
|
||||
`original_tokens − compressed_tokens` across the manifest's
|
||||
Compressed outcomes) and ``transforms_applied`` (deduplicated list
|
||||
of compressor strategy names). The Python proxy uses these to
|
||||
populate /transformations/feed and the dashboard's per-request log
|
||||
without recounting tokens. See `crates/headroom-core/src/transforms/
|
||||
live_zone.rs::CompressionManifest::tokens_saved` /
|
||||
`::transforms_applied`."""
|
||||
|
||||
def test_no_change_returns_zero_savings_and_empty_transforms(self):
|
||||
compress = _ensure_binding()
|
||||
body = json.dumps({"model": "gpt-4o-mini", "input": []}).encode()
|
||||
out, modified, saved, transforms = compress(body, "payg", "gpt-4o-mini")
|
||||
assert modified is False
|
||||
assert out == body
|
||||
assert saved == 0
|
||||
assert transforms == []
|
||||
|
||||
def test_field_types(self):
|
||||
"""Pin the wire shape so downstream callers don't break."""
|
||||
compress = _ensure_binding()
|
||||
body = json.dumps({"model": "gpt-4o-mini", "input": []}).encode()
|
||||
result = compress(body, "payg", "gpt-4o-mini")
|
||||
assert isinstance(result, tuple)
|
||||
assert len(result) == 4
|
||||
out, modified, saved, transforms = result
|
||||
assert isinstance(out, bytes)
|
||||
assert isinstance(modified, bool)
|
||||
assert isinstance(saved, int)
|
||||
assert isinstance(transforms, list)
|
||||
assert all(isinstance(t, str) for t in transforms)
|
||||
|
||||
def test_large_local_shell_output_compresses_with_telemetry(self):
|
||||
"""End-to-end check: a payload large enough to clear the
|
||||
per-item byte threshold produces ``modified=True`` plus a
|
||||
non-zero ``tokens_saved`` and a populated ``transforms``
|
||||
list. Mirrors the shape in the Rust crate's
|
||||
``large_log_output_compressed`` test."""
|
||||
compress = _ensure_binding()
|
||||
log_body = "".join(
|
||||
f"[2024-01-01 00:00:00] INFO compile.rs:42 building module foo_{i}\n"
|
||||
for i in range(400)
|
||||
)
|
||||
assert len(log_body) > 2048
|
||||
body = json.dumps(
|
||||
{
|
||||
"model": "gpt-4o",
|
||||
"input": [
|
||||
{
|
||||
"type": "local_shell_call_output",
|
||||
"call_id": "c1",
|
||||
"output": log_body,
|
||||
}
|
||||
],
|
||||
}
|
||||
).encode()
|
||||
out, modified, saved, transforms = compress(body, "payg", "gpt-4o")
|
||||
assert modified is True
|
||||
assert saved > 0
|
||||
assert transforms, "expected at least one strategy in transforms"
|
||||
new_doc = json.loads(out)
|
||||
assert new_doc["input"][0]["type"] == "local_shell_call_output"
|
||||
assert len(new_doc["input"][0]["output"]) < len(log_body)
|
||||
|
|
|
|||
|
|
@ -60,7 +60,7 @@ def _ws_compress_first_frame(
|
|||
model = (inner.get("model") if isinstance(inner, dict) else None) or ""
|
||||
|
||||
inner_bytes = json.dumps(inner).encode("utf-8")
|
||||
new_bytes, modified = compress(inner_bytes, auth_mode_value, model)
|
||||
new_bytes, modified, _saved, _transforms = compress(inner_bytes, auth_mode_value, model)
|
||||
if not modified:
|
||||
return first_msg_raw, False
|
||||
|
||||
|
|
|
|||
|
|
@ -87,8 +87,13 @@ class TestMultiTurnCompression:
|
|||
_make_user_msg("now edit it"),
|
||||
]
|
||||
frozen = cache.compute_frozen_count(messages)
|
||||
# All 4 messages stable (user, tool_use, tool_result cached, user)
|
||||
assert frozen == 4
|
||||
# First 3 stable; trailing user message ("now edit it") is the
|
||||
# live zone by construction — it has not been sent upstream
|
||||
# before, so it cannot be in any provider prefix cache. Cap at
|
||||
# len - 1 prevents the over-freeze pattern that produced 0 %
|
||||
# compression for prose-format clients (issue observed
|
||||
# 2026-05-07 with Cline+DeepSeek).
|
||||
assert frozen == 3
|
||||
|
||||
# apply_cached should swap the content
|
||||
result = cache.apply_cached(messages)
|
||||
|
|
@ -129,9 +134,10 @@ class TestMultiTurnCompression:
|
|||
# Now cache C too
|
||||
cache.store_compressed(CompressionCache.content_hash(code_c), "cc", tokens_saved=100)
|
||||
|
||||
# Turn 3: all cached
|
||||
# Turn 3: all 6 messages structurally stable, but the trailing
|
||||
# message is reserved as live zone. Frozen prefix = 5.
|
||||
frozen = cache.compute_frozen_count(messages)
|
||||
assert frozen == 6 # all stable
|
||||
assert frozen == 5
|
||||
|
||||
|
||||
class TestNoMessageInjection:
|
||||
|
|
@ -305,3 +311,92 @@ class TestUpdateFromResult:
|
|||
msg = _make_user_msg("same content")
|
||||
cache.update_from_result([msg], [msg])
|
||||
assert cache.get_stats()["entries"] == 0
|
||||
|
||||
|
||||
class TestProseFormatLiveZoneInvariant:
|
||||
"""Cline / OpenClaude / Aider — prose-format clients send tool calls
|
||||
embedded in plain assistant text and tool results pasted into plain
|
||||
user messages. There are no `tool_use`, `tool_result`, or
|
||||
``role: "tool"`` blocks anywhere in the conversation.
|
||||
|
||||
Pre-fix: ``compute_frozen_count`` walked all messages and found no
|
||||
"unstable" boundary, returning ``len(messages)``. The pipeline then
|
||||
froze every message — including the brand-new user turn — leaving
|
||||
the live zone empty. ContentRouter saw ``saved 0`` on every request.
|
||||
Bug observed 2026-05-07 with Cline+DeepSeek over /v1/chat/completions
|
||||
in token mode.
|
||||
|
||||
Post-fix: cap at ``len(messages) - 1`` always reserves the trailing
|
||||
message as the live zone. These tests lock that invariant.
|
||||
"""
|
||||
|
||||
def test_pure_user_assistant_turns_leave_live_zone(self):
|
||||
cache = CompressionCache()
|
||||
# A 6-turn Cline-shaped conversation: alternating user/assistant
|
||||
# plain-text. No tool blocks of any kind.
|
||||
messages = [
|
||||
_make_user_msg("system instructions baked into first user msg"),
|
||||
_make_assistant_msg("<execute_command>ls</execute_command>"),
|
||||
_make_user_msg("[tool_result]\nfile1.py\nfile2.py\n[/tool_result]"),
|
||||
_make_assistant_msg("<read_file>file1.py</read_file>"),
|
||||
_make_user_msg("[tool_result]\n<contents...>\n[/tool_result]"),
|
||||
_make_user_msg("now please refactor it"),
|
||||
]
|
||||
frozen = cache.compute_frozen_count(messages)
|
||||
# Pre-fix would return 6 (every plain message is "stable").
|
||||
# Post-fix: 6 messages stable, capped at len-1 = 5.
|
||||
assert frozen == 5
|
||||
assert frozen < len(messages), (
|
||||
"Live zone must never be empty; trailing user message must "
|
||||
"always be available for compression"
|
||||
)
|
||||
|
||||
def test_single_message_yields_zero_frozen(self):
|
||||
# Edge case: only the user's first message, nothing to freeze.
|
||||
cache = CompressionCache()
|
||||
messages = [_make_user_msg("first turn")]
|
||||
assert cache.compute_frozen_count(messages) == 0
|
||||
|
||||
def test_empty_messages_yields_zero(self):
|
||||
cache = CompressionCache()
|
||||
assert cache.compute_frozen_count([]) == 0
|
||||
|
||||
def test_two_messages_first_is_frozen_second_is_live(self):
|
||||
cache = CompressionCache()
|
||||
messages = [
|
||||
_make_user_msg("turn 1 content"),
|
||||
_make_user_msg("turn 2 — the live zone"),
|
||||
]
|
||||
# First message structurally stable; trailing is live → 1 frozen.
|
||||
assert cache.compute_frozen_count(messages) == 1
|
||||
|
||||
def test_anthropic_format_last_tool_result_is_still_live(self):
|
||||
"""Even when the trailing message is a tool_result whose content
|
||||
IS in the cache, it stays in the live zone. Trailing == live by
|
||||
construction; the cache-read decision is upstream's job."""
|
||||
cache = CompressionCache()
|
||||
code = _large_code_content(50)
|
||||
cache.store_compressed(
|
||||
CompressionCache.content_hash(code), "compressed code", tokens_saved=200
|
||||
)
|
||||
messages = [
|
||||
_make_user_msg("hi"),
|
||||
_make_assistant_msg("ok"),
|
||||
_make_tool_result_msg("t1", code),
|
||||
]
|
||||
frozen = cache.compute_frozen_count(messages)
|
||||
# Walk gets to 3, cap clamps to 2 (= len-1).
|
||||
assert frozen == 2
|
||||
|
||||
def test_openai_format_last_tool_msg_is_live(self):
|
||||
cache = CompressionCache()
|
||||
content = "tool output " * 100
|
||||
cache.store_compressed(
|
||||
CompressionCache.content_hash(content), "compressed", tokens_saved=300
|
||||
)
|
||||
messages = [
|
||||
_make_user_msg("run cmd"),
|
||||
_make_openai_tool_msg("tc1", content),
|
||||
]
|
||||
# Walk: user (stable, 1), tool (cached, 2). Cap → 1.
|
||||
assert cache.compute_frozen_count(messages) == 1
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue