akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

Merge pull request #446 from chopratejas/codex-ws-cancel-logging

Browse source 
fix: Codex ws cancel logging
This commit is contained in:
Tejas Chopra 2026-05-10 13:21:09 -07:00 • committed by GitHub
commit 73b0e56e97
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 733 additions and 18 deletions
Download patch file Download diff file Expand all files Collapse all files

77
headroom/cache/compression_store.py vendored
View file

@ -53,6 +53,32 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
_RETRIEVAL_LOG_PREVIEW_CHARS = 4096
_SECRET_KEY_VALUE_RE = re.compile(
r"(?i)\b([A-Z0-9_-]*(?:API[_-]?KEY|TOKEN|SECRET|PASSWORD|CREDENTIAL|AUTH)[A-Z0-9_-]*)"
r"(\s*[:=]\s*)([\"']?)([^\"'\s,}]+)"
)
_AUTH_VALUE_RE = re.compile(r"(?i)\b(Bearer|Basic)\s+[A-Za-z0-9._~+/=-]{12,}")
_API_KEY_VALUE_RE = re.compile(r"\bsk-[A-Za-z0-9_-]{12,}\b")
def _redact_retrieval_log_payload(payload: str) -> str:
redacted = _SECRET_KEY_VALUE_RE.sub(r"\1\2\3[REDACTED]", payload)
redacted = _AUTH_VALUE_RE.sub(r"\1 [REDACTED]", redacted)
return _API_KEY_VALUE_RE.sub("sk-[REDACTED]", redacted)
def _payload_for_retrieval_log(payload: str) -> dict[str, Any]:
redacted = _redact_retrieval_log_payload(payload)
preview = redacted[:_RETRIEVAL_LOG_PREVIEW_CHARS]
truncated = len(redacted) > len(preview)
return {
"payload_chars": len(payload),
"payload_preview_chars": len(preview),
"payload_truncated": truncated,
"payload_preview": preview,
}
@dataclass
class CompressionEntry:
@ -337,6 +363,15 @@ class CompressionStore:
retrieval_type="full",
tool_signature_hash=entry.tool_signature_hash,
)
self._log_retrieval_payload(
hash_key=hash_key,
query=query,
retrieval_type="full",
payload=entry.original_content,
items_retrieved=entry.original_item_count,
total_items=entry.original_item_count,
entry=entry,
)
# CRITICAL: Make a deep copy to return
# (entry could be modified/evicted after lock release)
@ -442,9 +477,51 @@ class CompressionStore:
)
# Process feedback immediately to ensure TOIN learns in real-time
self.process_pending_feedback()
self._log_retrieval_payload(
hash_key=hash_key,
query=query,
retrieval_type="search",
payload=json.dumps(results, ensure_ascii=False),
items_retrieved=len(results),
total_items=len(items),
entry=entry,
)
return results
def _log_retrieval_payload(
self,
*,
hash_key: str,
query: str | None,
retrieval_type: str,
payload: str,
items_retrieved: int,
total_items: int,
entry: CompressionEntry,
) -> None:
event = {
"event": "headroom_retrieve",
"hash": hash_key,
"retrieval_type": retrieval_type,
"query": query,
"items_retrieved": items_retrieved,
"total_items": total_items,
"tool_name": entry.tool_name,
"tool_call_id": entry.tool_call_id,
"compression_strategy": entry.compression_strategy,
"tool_signature_hash": entry.tool_signature_hash,
"original_tokens": entry.original_tokens,
"compressed_tokens": entry.compressed_tokens,
"original_item_count": entry.original_item_count,
"compressed_item_count": entry.compressed_item_count,
**_payload_for_retrieval_log(payload),
}
logger.info(
"event=headroom_retrieve %s",
json.dumps(event, ensure_ascii=False, separators=(",", ":")),
)
def _search_items_from_original(self, original_content: str) -> list[Any]:
"""Normalize cached originals into searchable items.

38
headroom/ccr/mcp_server.py
View file

@ -581,22 +581,39 @@ class HeadroomMCPServer:
@self.server.call_tool()
async def call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent]:
started = time.perf_counter()
logger.info(
"event=mcp_tool_call_received tool=%s arguments=%s",
name,
json.dumps(arguments, ensure_ascii=False, default=str),
)
try:
if name == COMPRESS_TOOL_NAME:
return await self._handle_compress(arguments)
result = await self._handle_compress(arguments)
elif name == CCR_TOOL_NAME:
return await self._handle_retrieve(arguments)
result = await self._handle_retrieve(arguments)
elif name == STATS_TOOL_NAME:
return await self._handle_stats()
result = await self._handle_stats()
elif name == READ_TOOL_NAME and _READ_ENABLED:
return await self._handle_read(arguments)
result = await self._handle_read(arguments)
else:
return [
result = [
TextContent(
type="text",
text=json.dumps({"error": f"Unknown tool: {name}"}),
)
]
logger.info(
"event=mcp_tool_call_completed tool=%s duration_ms=%.2f output=%s",
name,
(time.perf_counter() - started) * 1000.0,
json.dumps(
[getattr(item, "text", str(item)) for item in result],
ensure_ascii=False,
default=str,
),
)
return result
except Exception as e:
logger.error(f"Tool {name} failed: {e}", exc_info=True)
return [
@ -635,7 +652,18 @@ class HeadroomMCPServer:
]
query = arguments.get("query")
logger.info(
"event=mcp_retrieve_started hash=%s query=%s",
hash_key,
json.dumps(query, ensure_ascii=False, default=str),
)
result = await self._retrieve_content(hash_key, query)
logger.info(
"event=mcp_retrieve_completed hash=%s query=%s result=%s",
hash_key,
json.dumps(query, ensure_ascii=False, default=str),
json.dumps(result, ensure_ascii=False, default=str),
)
return [TextContent(type="text", text=json.dumps(result, indent=2))]

51
headroom/providers/proxy_routes.py
View file

@ -46,6 +46,41 @@ def _select_passthrough_base_url(proxy: Any, headers: dict[str, str]) -> str:
return _api_target(proxy, provider_name)
async def _handle_chatgpt_model_metadata(
proxy: Any,
request: Request,
upstream_path: str,
) -> Response | None:
headers = dict(request.headers.items())
headers.pop("host", None)
headers, is_chatgpt_auth = _resolve_codex_routing_headers(headers)
if not is_chatgpt_auth:
return None
url = f"https://chatgpt.com{upstream_path}"
if request.url.query:
url = f"{url}?{request.url.query}"
body = await request.body()
try:
assert proxy.http_client is not None
resp = await proxy.http_client.request(
request.method,
url,
headers=headers,
content=body,
timeout=120.0,
)
return Response(
content=resp.content,
status_code=resp.status_code,
headers=dict(resp.headers),
)
except Exception as exc:
logger.error("Passthrough %s failed: %s", upstream_path, exc)
return Response(content=str(exc), status_code=502)
def register_provider_routes(app: FastAPI, proxy: Any) -> None:
"""Register provider-specific proxy endpoints."""
@ -207,6 +242,14 @@ def register_provider_routes(app: FastAPI, proxy: Any) -> None:
@app.get("/v1/models")
async def list_models(request: Request):
chatgpt_response = await _handle_chatgpt_model_metadata(
proxy,
request,
"/backend-api/models",
)
if chatgpt_response is not None:
return chatgpt_response
provider_name = proxy.provider_runtime.model_metadata_provider(dict(request.headers))
return await proxy.handle_passthrough(
request,
@ -217,6 +260,14 @@ def register_provider_routes(app: FastAPI, proxy: Any) -> None:
@app.get("/v1/models/{model_id}")
async def get_model(request: Request, model_id: str):
chatgpt_response = await _handle_chatgpt_model_metadata(
proxy,
request,
f"/backend-api/models/{model_id}",
)
if chatgpt_response is not None:
return chatgpt_response
provider_name = proxy.provider_runtime.model_metadata_provider(dict(request.headers))
return await proxy.handle_passthrough(
request,

122
headroom/proxy/handlers/openai.py
View file

@ -2236,6 +2236,23 @@ class OpenAIHandlerMixin:
_ws_url_obj = getattr(websocket, "url", None)
_ws_url = str(_ws_url_obj) if _ws_url_obj is not None else ""
_ws_path = getattr(_ws_url_obj, "path", "") if _ws_url_obj is not None else ""
if not _ws_path:
_ws_path = "/v1/responses"
metrics_for_inbound_ws = getattr(self, "metrics", None)
if metrics_for_inbound_ws is not None and hasattr(
metrics_for_inbound_ws, "record_inbound_request"
):
with contextlib.suppress(Exception):
metrics_for_inbound_ws.record_inbound_request(method="WS", path=_ws_path)
logger.info(
"event=proxy_inbound_websocket request_id=%s session_id=%s path=%s "
"client=%s header_count=%d",
request_id,
session_id,
_ws_path,
getattr(websocket, "client", ""),
len(ws_headers),
)
from headroom.proxy.helpers import capture_codex_wire_debug
capture_codex_wire_debug(
@ -2549,6 +2566,12 @@ class OpenAIHandlerMixin:
ws_recorded_uncached_input_tokens_total = 0
ws_recorded_tokens_saved_total = 0
ws_response_create_frames = 1
ws_client_frames_total = 1
ws_upstream_frames_total = 0
ws_cancel_frames = 0
ws_last_client_frame_type = str(body.get("type") or "unknown") if body else "unknown"
ws_last_upstream_frame_type = "unknown"
ws_client_disconnect_seen = False
_ws_bypass = self._headroom_bypass_enabled(ws_headers)
if _ws_bypass:
logger.info(
@ -3043,16 +3066,44 @@ class OpenAIHandlerMixin:
async def _client_to_upstream() -> None:
nonlocal client_relay_error, ws_response_create_frames
nonlocal ws_client_frames_total, ws_cancel_frames
nonlocal ws_last_client_frame_type, ws_client_disconnect_seen
client_frame_index = 1
try:
while True:
msg = await websocket.receive_text()
client_frame_index += 1
ws_client_frames_total += 1
if session_handle is not None:
session_handle.mark_activity()
_inbound_frame_body: Any = None
try:
_inbound_frame_body = json.loads(msg)
except json.JSONDecodeError:
_inbound_frame_body = None
ws_last_client_frame_type = (
str(_inbound_frame_body.get("type") or "unknown")
if isinstance(_inbound_frame_body, dict)
else "non_json"
)
if ws_last_client_frame_type == "response.cancel":
ws_cancel_frames += 1
logger.info(
"[%s] WS client sent response.cancel "
"session_id=%s frame=%d cancels=%d",
request_id,
session_id,
client_frame_index,
ws_cancel_frames,
)
else:
logger.debug(
"[%s] WS client frame session_id=%s frame=%d type=%s",
request_id,
session_id,
client_frame_index,
ws_last_client_frame_type,
)
capture_codex_wire_debug(
"ws_inbound_client_frame",
request_id=request_id,
@ -3115,6 +3166,17 @@ class OpenAIHandlerMixin:
logger.debug(
f"[{request_id}] WS client→upstream relay ended: {relay_err}"
)
else:
ws_client_disconnect_seen = True
logger.info(
"[%s] WS client disconnected session_id=%s "
"frames=%d cancels=%d last_type=%s",
request_id,
session_id,
ws_client_frames_total,
ws_cancel_frames,
ws_last_client_frame_type,
)
with contextlib.suppress(Exception):
await upstream.close()
@ -3147,6 +3209,7 @@ class OpenAIHandlerMixin:
nonlocal ws_recorded_cache_write_tokens_total
nonlocal ws_recorded_uncached_input_tokens_total
nonlocal ws_recorded_tokens_saved_total
nonlocal ws_upstream_frames_total, ws_last_upstream_frame_type
memory_enabled = bool(self.memory_handler and memory_user_id)
@ -3247,6 +3310,9 @@ class OpenAIHandlerMixin:
upstream_frame_index = 0
async for msg in upstream:
upstream_frame_index += 1
ws_upstream_frames_total += 1
if session_handle is not None:
session_handle.mark_activity()
if (
_first_event_started_at is not None
and "upstream_first_event" not in stage_timer
@ -3257,6 +3323,7 @@ class OpenAIHandlerMixin:
* 1000.0,
)
if isinstance(msg, bytes):
ws_last_upstream_frame_type = "binary"
capture_codex_wire_debug(
"ws_upstream_binary_frame",
request_id=request_id,
@ -3295,10 +3362,19 @@ class OpenAIHandlerMixin:
try:
event = json.loads(msg_str)
except (json.JSONDecodeError, TypeError):
ws_last_upstream_frame_type = "non_json"
await websocket.send_text(msg_str)
continue
event_type = event.get("type", "")
ws_last_upstream_frame_type = str(event_type or "unknown")
logger.debug(
"[%s] WS upstream frame session_id=%s frame=%d type=%s",
request_id,
session_id,
upstream_frame_index,
ws_last_upstream_frame_type,
)
if event_type == "response.created":
response_started_ms = time.perf_counter() * 1000.0
(
@ -3556,6 +3632,13 @@ class OpenAIHandlerMixin:
logger.debug(
f"[{request_id}] WS relay {task_name} raised: {exc!r}"
)
if (
ws_cancel_frames > 0
and not response_completed_seen
and termination_cause
in {"upstream_disconnect", "client_disconnect", "unknown"}
):
termination_cause = "client_cancel"
finally:
# In case anything above raised before the
# cancel-and-await loop ran.
@ -3566,8 +3649,19 @@ class OpenAIHandlerMixin:
await asyncio.gather(*relay_tasks, return_exceptions=True)
logger.info(
f"[{request_id}] WS /v1/responses completed "
f"(tokens_saved={tokens_saved}, cause={termination_cause})"
"[%s] WS /v1/responses completed "
"(tokens_saved=%d, cause=%s, client_frames=%d, upstream_frames=%d, "
"cancel_frames=%d, client_disconnect=%s, last_client_type=%s, "
"last_upstream_type=%s)",
request_id,
tokens_saved,
termination_cause,
ws_client_frames_total,
ws_upstream_frames_total,
ws_cancel_frames,
ws_client_disconnect_seen,
ws_last_client_frame_type,
ws_last_upstream_frame_type,
)
break
except Exception as ws_err:
@ -3658,6 +3752,13 @@ class OpenAIHandlerMixin:
"route": "chatgpt_subscription" if is_chatgpt_auth else "openai_api",
"ws_response_create_frames": str(ws_response_create_frames),
"ws_frames_compressed": str(ws_frames_compressed),
"ws_client_frames_total": str(ws_client_frames_total),
"ws_upstream_frames_total": str(ws_upstream_frames_total),
"ws_cancel_frames": str(ws_cancel_frames),
"ws_last_client_frame_type": ws_last_client_frame_type,
"ws_last_upstream_frame_type": ws_last_upstream_frame_type,
"ws_client_disconnect_seen": str(ws_client_disconnect_seen),
"ws_termination_cause": termination_cause,
"cache_read_tokens": str(ws_cache_read_tokens_total),
"cache_write_tokens": str(ws_cache_write_tokens_total),
"uncached_input_tokens": str(ws_uncached_input_tokens_total),
@ -3793,6 +3894,23 @@ class OpenAIHandlerMixin:
metrics_for_close.record_ws_session_duration(
session_duration_ms, termination_cause
)
metrics_for_ws_inbound_close = getattr(self, "metrics", None)
if metrics_for_ws_inbound_close is not None and hasattr(
metrics_for_ws_inbound_close, "record_inbound_response"
):
with contextlib.suppress(Exception):
metrics_for_ws_inbound_close.record_inbound_response(
status_code=f"ws:{termination_cause}"
)
logger.info(
"event=proxy_inbound_websocket_closed request_id=%s session_id=%s "
"path=%s cause=%s duration_ms=%.2f",
request_id,
session_id,
_ws_path,
termination_cause,
(time.perf_counter() - session_started_at) * 1000.0,
)
await emit_stage_timings_log(
path="openai_responses_ws",
request_id=request_id,

14
headroom/proxy/helpers.py
View file

@ -100,17 +100,17 @@ def _safe_event_name(event: str) -> str:
return "".join(ch if ch.isalnum() or ch in ("-", "_") else "_" for ch in event)[:80]
def _wire_debug_preview(value: Any, *, max_chars: int = 900) -> str:
"""Return a compact, human-readable preview for proxy.log.
def _wire_debug_preview(value: Any, *, max_chars: int | None = None) -> str:
"""Return the redacted wire payload for proxy.log.
This is intentionally lossy: the full redacted payload is already written
to the wire-debug JSON file. The log line should be short enough to scan
live without flooding the proxy log.
This is intentionally not truncated. During Codex WS debugging we need the
proxy log itself to show the complete frame so we can decide later where a
deliberate trim boundary belongs.
"""
try:
if isinstance(value, bytes):
text = safe_decode_for_logging(value, max_bytes=max_chars)
text = value.decode("utf-8", errors="replace")
elif isinstance(value, str):
text = value
elif value is None:
@ -121,7 +121,7 @@ def _wire_debug_preview(value: Any, *, max_chars: int = 900) -> str:
text = repr(value)
text = " ".join(text.split())
if len(text) > max_chars:
if max_chars is not None and len(text) > max_chars:
return text[: max_chars - 1] + "…"
return text

59
headroom/proxy/prometheus_metrics.py
View file

@ -76,6 +76,12 @@ class PrometheusMetrics:
self.requests_cached = 0
self.requests_rate_limited = 0
self.requests_failed = 0
self.inbound_requests_total = 0
self.inbound_requests_completed = 0
self.inbound_requests_active = 0
self.inbound_requests_by_method: dict[str, int] = defaultdict(int)
self.inbound_requests_by_path: dict[str, int] = defaultdict(int)
self.inbound_responses_by_status: dict[str, int] = defaultdict(int)
self.tokens_input_total = 0
self.tokens_output_total = 0
@ -203,6 +209,12 @@ class PrometheusMetrics:
self.requests_cached = 0
self.requests_rate_limited = 0
self.requests_failed = 0
self.inbound_requests_total = 0
self.inbound_requests_completed = 0
self.inbound_requests_active = 0
self.inbound_requests_by_method.clear()
self.inbound_requests_by_path.clear()
self.inbound_responses_by_status.clear()
self.tokens_input_total = 0
self.tokens_output_total = 0
@ -339,6 +351,32 @@ class PrometheusMetrics:
if saved > 0:
self.tokens_saved_by_strategy[strategy] += saved
def record_inbound_request(self, *, method: str, path: str) -> None:
self.inbound_requests_total += 1
self.inbound_requests_active += 1
self.inbound_requests_by_method[method.upper()] += 1
self.inbound_requests_by_path[path] += 1
def record_inbound_response(self, *, status_code: int | str) -> None:
self.inbound_requests_completed += 1
self.inbound_requests_active = max(0, self.inbound_requests_active - 1)
self.inbound_responses_by_status[str(status_code)] += 1
def record_inbound_aborted(self, *, reason: str) -> None:
self.inbound_requests_completed += 1
self.inbound_requests_active = max(0, self.inbound_requests_active - 1)
self.inbound_responses_by_status[f"aborted:{reason}"] += 1
def inbound_snapshot(self) -> dict[str, object]:
return {
"total": self.inbound_requests_total,
"completed": self.inbound_requests_completed,
"active": self.inbound_requests_active,
"by_method": dict(self.inbound_requests_by_method),
"by_path": dict(self.inbound_requests_by_path),
"by_status": dict(self.inbound_responses_by_status),
}
async def record_request(
self,
provider: str,
@ -599,6 +637,27 @@ class PrometheusMetrics:
help_text="Failed requests",
value=self.requests_failed,
)
_append_metric(
lines,
name="headroom_inbound_requests_total",
metric_type="counter",
help_text="All inbound HTTP requests accepted by the proxy",
value=self.inbound_requests_total,
)
_append_metric(
lines,
name="headroom_inbound_requests_completed_total",
metric_type="counter",
help_text="Inbound HTTP requests completed or aborted by the proxy",
value=self.inbound_requests_completed,
)
_append_metric(
lines,
name="headroom_inbound_requests_active",
metric_type="gauge",
help_text="Inbound HTTP requests currently active in the proxy",
value=self.inbound_requests_active,
)
_append_metric(
lines,
name="headroom_tokens_input_total",

79
headroom/proxy/server.py
View file

@ -1556,6 +1556,39 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
# outermost and we don't count requests they reject.
@app.middleware("http")
async def _record_headroom_stack(request, call_next):
started = time.perf_counter()
inbound_id = f"inbound-{time.time_ns()}"
path = request.url.path
method = request.method
query = request.url.query
headers = dict(request.headers.items())
client = getattr(request, "client", None)
client_addr = ""
if client is not None:
client_host = getattr(client, "host", None)
client_port = getattr(client, "port", None)
client_addr = f"{client_host}:{client_port}" if client_port else str(client_host)
try:
proxy.metrics.record_inbound_request(method=method, path=path)
except Exception:
logger.debug("record_inbound_request failed", exc_info=True)
try:
from headroom.proxy.helpers import redact_for_wire_debug
safe_headers = redact_for_wire_debug(headers)
except Exception:
safe_headers = {"redaction_error": True}
logger.info(
"event=proxy_inbound_request id=%s method=%s path=%s query=%s client=%s "
"content_length=%s headers=%s",
inbound_id,
method,
path,
query,
client_addr,
request.headers.get("content-length", ""),
json.dumps(safe_headers, ensure_ascii=False, default=str),
)
if request.url.path.startswith("/v1/"):
stack = request.headers.get("x-headroom-stack")
if stack:
@ -1563,7 +1596,50 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
proxy.metrics.record_stack(stack)
except Exception:
logger.debug("record_stack failed", exc_info=True)
return await call_next(request)
try:
response = await call_next(request)
except asyncio.CancelledError:
try:
proxy.metrics.record_inbound_aborted(reason="cancelled")
except Exception:
logger.debug("record_inbound_aborted failed", exc_info=True)
logger.info(
"event=proxy_inbound_request_aborted id=%s method=%s path=%s reason=cancelled "
"duration_ms=%.2f",
inbound_id,
method,
path,
(time.perf_counter() - started) * 1000.0,
)
raise
except Exception as exc:
try:
proxy.metrics.record_inbound_aborted(reason=type(exc).__name__)
except Exception:
logger.debug("record_inbound_aborted failed", exc_info=True)
logger.info(
"event=proxy_inbound_request_aborted id=%s method=%s path=%s reason=%s "
"duration_ms=%.2f",
inbound_id,
method,
path,
type(exc).__name__,
(time.perf_counter() - started) * 1000.0,
)
raise
try:
proxy.metrics.record_inbound_response(status_code=response.status_code)
except Exception:
logger.debug("record_inbound_response failed", exc_info=True)
logger.info(
"event=proxy_inbound_response id=%s method=%s path=%s status=%s duration_ms=%.2f",
inbound_id,
method,
path,
response.status_code,
(time.perf_counter() - started) * 1000.0,
)
return response
# Third-party proxy extensions (Enterprise, custom plugins). Discovered via
# the `headroom.proxy_extension` entry-point group, but **opt-in only**:
@ -1905,6 +1981,7 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
},
"toin": get_toin().get_stats(),
"cli_filtering": cli_filtering_stats,
"proxy_inbound": proxy.metrics.inbound_snapshot(),
"cache": await proxy.cache.stats() if proxy.cache else None,
"rate_limiter": await proxy.rate_limiter.stats() if proxy.rate_limiter else None,
"recent_requests": proxy.logger.get_recent(10) if proxy.logger else [],

1
headroom/proxy/ws_session_registry.py
View file

@ -50,6 +50,7 @@ TerminationCause = Literal[
"upstream_disconnect",
"upstream_error",
"client_error",
"client_cancel",
"response_completed",
"client_timeout",
"unknown",

79
headroom/transforms/kompress_compressor.py
View file

@ -14,7 +14,9 @@ Usage:
from __future__ import annotations
import contextlib
import gc
import hashlib
import logging
import threading
from dataclasses import dataclass
@ -36,6 +38,62 @@ _kompress_cache: dict[str, tuple[Any, Any, str]] = {}
_kompress_lock = threading.Lock()
def _bucket_count(value: int) -> str:
"""Return a coarse, privacy-preserving size bucket."""
if value <= 0:
return "0"
lower = 1 << (value.bit_length() - 1)
upper = lower << 1
return f"{lower}-{upper}"
def _kompress_content_signature(content: str) -> Any:
"""Create a first-class TOIN signature for Kompress/plain-text content.
This intentionally keys on shape, not values. Retrieval pressure should
teach TOIN about this class of compressed content without storing the
content or treating it as an anonymous fallback.
"""
from ..telemetry.models import ToolSignature
words = content.split()
line_count = content.count("\n") + 1 if content else 0
nonempty_lines = [line for line in content.splitlines() if line.strip()]
avg_line_chars = (
sum(len(line) for line in nonempty_lines) // len(nonempty_lines) if nonempty_lines else 0
)
has_paths = "/" in content or "\\" in content
has_assignment_like_tokens = any("=" in word for word in words[:200])
has_brackets = any(ch in content for ch in "{}[]()")
has_error_terms = any(
term in content.lower() for term in ("error", "exception", "traceback", "failed", "fatal")
)
shape = "|".join(
(
"kompress-text",
f"chars:{_bucket_count(len(content))}",
f"words:{_bucket_count(len(words))}",
f"lines:{_bucket_count(line_count)}",
f"avg_line:{_bucket_count(avg_line_chars)}",
f"paths:{int(has_paths)}",
f"assign:{int(has_assignment_like_tokens)}",
f"brackets:{int(has_brackets)}",
f"errors:{int(has_error_terms)}",
)
)
structure_hash = hashlib.sha256(shape.encode()).hexdigest()[:24]
return ToolSignature(
structure_hash=structure_hash,
field_count=0,
has_nested_objects=False,
has_arrays=False,
max_depth=0,
string_field_count=1,
has_error_like_field=has_error_terms,
has_message_like_field=True,
)
def _is_onnx_available() -> bool:
"""Check if ONNX Runtime is available (lightweight, no torch needed)."""
try:
@ -800,13 +858,30 @@ class KompressCompressor(Transform):
try:
from ..cache.compression_store import get_compression_store
signature = _kompress_content_signature(original)
compressed_tokens = len(compressed.split())
store = get_compression_store()
return store.store(
cache_key = store.store(
original,
compressed,
original_tokens=original_tokens,
compressed_tokens=len(compressed.split()),
compressed_tokens=compressed_tokens,
original_item_count=original_tokens,
compressed_item_count=compressed_tokens,
tool_signature_hash=signature.structure_hash,
compression_strategy="kompress",
)
with contextlib.suppress(Exception):
from ..telemetry import get_toin
get_toin().record_compression(
tool_signature=signature,
original_count=original_tokens,
compressed_count=compressed_tokens,
original_tokens=original_tokens,
compressed_tokens=compressed_tokens,
strategy="kompress",
)
return cache_key
except Exception:
return None

96
tests/test_compression_store.py
View file

@ -17,8 +17,10 @@ from __future__ import annotations
import hashlib
import json
import logging
import threading
import time
from contextlib import contextmanager
from typing import Any
from unittest.mock import MagicMock, patch
@ -32,6 +34,100 @@ from headroom.cache.compression_store import (
reset_compression_store,
)
@contextmanager
def _capture_headroom_retrieve_events():
events: list[dict[str, Any]] = []
prefix = "event=headroom_retrieve "
class _Handler(logging.Handler):
def emit(self, record: logging.LogRecord) -> None:
message = record.getMessage()
if prefix in message:
events.append(json.loads(message.split(prefix, 1)[1]))
logger = logging.getLogger("headroom.cache.compression_store")
previous_level = logger.level
handler = _Handler(level=logging.INFO)
logger.addHandler(handler)
logger.setLevel(logging.INFO)
try:
yield events
finally:
logger.removeHandler(handler)
logger.setLevel(previous_level)
def test_retrieve_logs_payload_preview():
store = CompressionStore(enable_feedback=False)
hash_key = store.store(
original="secret-ish payload for operator debugging",
compressed="payload",
original_tokens=8,
compressed_tokens=1,
original_item_count=1,
compressed_item_count=1,
tool_name="tool_a",
)
with _capture_headroom_retrieve_events() as events:
entry = store.retrieve(hash_key)
assert entry is not None
assert len(events) == 1
assert events[0]["hash"] == hash_key
assert events[0]["retrieval_type"] == "full"
assert events[0]["payload_preview"] == "secret-ish payload for operator debugging"
assert "payload" not in events[0]
assert events[0]["payload_truncated"] is False
assert events[0]["tool_name"] == "tool_a"
def test_retrieve_log_redacts_secret_payload_values():
store = CompressionStore(enable_feedback=False)
hash_key = store.store(
original="OPENAI_API_KEY=sk-proj-secret1234567890 Authorization: Bearer token123456789",
compressed="payload",
)
with _capture_headroom_retrieve_events() as events:
entry = store.retrieve(hash_key)
assert entry is not None
assert len(events) == 1
assert "sk-proj-secret1234567890" not in events[0]["payload_preview"]
assert "Bearer token123456789" not in events[0]["payload_preview"]
assert "OPENAI_API_KEY=[REDACTED]" in events[0]["payload_preview"]
assert "Authorization: [REDACTED]" in events[0]["payload_preview"]
def test_search_logs_retrieved_payload_preview():
store = CompressionStore(enable_feedback=False)
items = [
{"id": 1, "text": "alpha target"},
{"id": 2, "text": "beta other"},
]
hash_key = store.store(
original=json.dumps(items),
compressed="[]",
original_item_count=2,
compressed_item_count=0,
tool_name="search_tool",
)
with _capture_headroom_retrieve_events() as events:
results = store.search(hash_key, "alpha", score_threshold=0.0)
assert results
assert len(events) == 1
assert events[0]["hash"] == hash_key
assert events[0]["retrieval_type"] == "search"
assert events[0]["query"] == "alpha"
assert events[0]["payload_preview"] == json.dumps(results, ensure_ascii=False)
assert events[0]["payload_preview_chars"] == len(json.dumps(results, ensure_ascii=False))
assert events[0]["payload_truncated"] is False
# =============================================================================
# Fixtures
# =============================================================================

40
tests/test_openai_codex_ws_lifecycle.py
View file

@ -424,6 +424,46 @@ async def test_upstream_error_mid_stream_classifies_as_upstream_error():
assert handler.metrics.termination_causes[-1] == "upstream_error"
@pytest.mark.asyncio
async def test_response_cancel_frame_is_logged_as_client_cancel_lifecycle():
"""A Codex Ctrl-C maps to response.cancel on the WS stream.
The proxy should relay it upstream and classify the lifecycle as a
client-side cancel when no response.completed event follows.
"""
cancel_frame = json.dumps({"type": "response.cancel", "response_id": "r_1"})
upstream = _FakeUpstream([], hold_after_events=True)
fake_ws_mod = _make_fake_websockets_module(upstream)
client_ws = _FakeWebSocket(
frames=[_first_frame(), cancel_frame],
hold_after_initial=True,
)
handler = _DummyOpenAIHandler()
async def _trigger() -> None:
await asyncio.sleep(0.05)
client_ws.trigger_disconnect()
with patch.dict(sys.modules, {"websockets": fake_ws_mod}):
trigger_task = asyncio.create_task(_trigger())
try:
await asyncio.wait_for(
handler.handle_openai_responses_ws(client_ws),
timeout=2.0,
)
finally:
trigger_task.cancel()
try:
await trigger_task
except asyncio.CancelledError:
pass
assert cancel_frame in upstream.sent
assert handler.metrics.termination_causes[-1] == "client_cancel"
assert handler.ws_sessions.active_count() == 0
@pytest.mark.asyncio
async def test_upstream_connect_failure_still_deregisters_cleanly():
"""Handshake-phase leak must be impossible: if upstream connect

5
tests/test_provider_proxy_routes.py
View file

@ -2,6 +2,7 @@ from __future__ import annotations
import importlib
from typing import Any
from unittest.mock import patch
import httpx
from fastapi.responses import JSONResponse
@ -252,10 +253,12 @@ def test_openai_response_subpath_passthrough_returns_502_on_http_failure() -> No
with TestClient(_app()) as client:
client.app.state.proxy.http_client = FailingAsyncClient()
response = client.post("/v1/responses/compact?trace=1", json={"model": "gpt-4o"})
with patch("headroom.providers.proxy_routes.logger") as logger:
response = client.post("/v1/responses/compact?trace=1", json={"model": "gpt-4o"})
assert response.status_code == 502
assert "boom: POST https://api.openai.test/v1/responses/compact?trace=1" in response.text
logger.error.assert_called_once()
def test_openai_response_subpath_passthrough_uses_openai_target() -> None:

53
tests/test_proxy_codex_route_aliases.py
View file

@ -157,3 +157,56 @@ def test_codex_responses_subpath_passthrough_derives_chatgpt_routing_from_jwt(pa
assert url == expected_url
assert headers["authorization"] == f"Bearer {token}"
assert headers["ChatGPT-Account-ID"] == "acct-from-jwt"
@pytest.mark.parametrize(
("path", "expected_url"),
[
(
"/v1/models?client_version=0.130.0",
"https://chatgpt.com/backend-api/models?client_version=0.130.0",
),
(
"/v1/models/gpt-5.3-codex?client_version=0.130.0",
"https://chatgpt.com/backend-api/models/gpt-5.3-codex?client_version=0.130.0",
),
],
)
def test_codex_model_metadata_routes_to_chatgpt_backend_for_subscription_auth(
path,
expected_url,
):
class FakeAsyncClient:
def __init__(self):
self.calls: list[tuple[str, str, dict[str, str]]] = []
async def request(self, method, url, **kwargs): # type: ignore[no-untyped-def]
self.calls.append((method, url, dict(kwargs.get("headers", {}))))
return httpx.Response(200, json={"method": method, "url": url})
async def aclose(self):
return None
token = _jwt(
{
"https://api.openai.com/auth": {
"chatgpt_account_id": "acct-from-jwt",
}
}
)
with TestClient(create_app(ProxyConfig())) as client:
fake_http_client = FakeAsyncClient()
client.app.state.proxy.http_client = fake_http_client
client.app.state.proxy.OPENAI_API_URL = "https://api.openai.test"
response = client.get(path, headers={"Authorization": f"Bearer {token}"})
assert response.status_code == 200
assert len(fake_http_client.calls) == 1
method, url, headers = fake_http_client.calls[0]
assert method == "GET"
assert url == expected_url
assert headers["authorization"] == f"Bearer {token}"
assert headers["ChatGPT-Account-ID"] == "acct-from-jwt"

37
tests/test_toin_feedback.py
View file

@ -6,13 +6,16 @@ from unittest.mock import MagicMock, patch
import pytest
from headroom.cache.compression_store import get_compression_store, reset_compression_store
from headroom.telemetry import (
TOINConfig,
ToolIntelligenceNetwork,
ToolPattern,
ToolSignature,
get_toin,
reset_toin,
)
from headroom.transforms.kompress_compressor import KompressCompressor
@pytest.fixture(autouse=True)
@ -21,7 +24,9 @@ def reset_globals(monkeypatch, tmp_path):
temp_toin_path = str(tmp_path / "toin_test.json")
monkeypatch.setenv("HEADROOM_TOIN_PATH", temp_toin_path)
reset_toin()
reset_compression_store()
yield
reset_compression_store()
reset_toin()
@ -45,6 +50,38 @@ def _make_signature(structure_hash: str = "test_hash_123") -> ToolSignature:
)
def test_kompress_ccr_retrieval_updates_toin():
"""Kompress CCR entries should be first-class TOIN patterns."""
original = "\n".join(
[
"HEADROOM_MODE=debug PATH=/tmp/headroom",
"ordinary line without the target token",
"another ordinary line",
]
)
compressed = "HEADROOM_MODE=debug"
compressor = KompressCompressor()
hash_key = compressor._store_in_ccr(
original,
compressed,
original_tokens=len(original.split()),
)
assert hash_key is not None
store = get_compression_store()
entry = store.retrieve(hash_key)
assert entry is not None
assert entry.tool_signature_hash is not None
results = store.search(hash_key, "HEADROOM", score_threshold=0.0)
assert results
stats = get_toin().get_stats()
assert stats["total_compressions"] == 1
assert stats["total_retrievals"] == 2
@pytest.mark.skip(reason="PR-B5: observations counter and request-time hint API retired")
class TestGetRecommendationObservations:
"""Bug 1: get_recommendation() should increment observations counter."""