diff --git a/headroom/cache/compression_store.py b/headroom/cache/compression_store.py index 31d6d042b..30c65fd79 100644 --- a/headroom/cache/compression_store.py +++ b/headroom/cache/compression_store.py @@ -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. diff --git a/headroom/ccr/mcp_server.py b/headroom/ccr/mcp_server.py index 2d57c3cff..806846eb8 100644 --- a/headroom/ccr/mcp_server.py +++ b/headroom/ccr/mcp_server.py @@ -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))] diff --git a/headroom/providers/proxy_routes.py b/headroom/providers/proxy_routes.py index 929973c3e..822327c90 100644 --- a/headroom/providers/proxy_routes.py +++ b/headroom/providers/proxy_routes.py @@ -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, diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 59abfd229..38cd602aa 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -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, diff --git a/headroom/proxy/helpers.py b/headroom/proxy/helpers.py index f346459a2..21b785514 100644 --- a/headroom/proxy/helpers.py +++ b/headroom/proxy/helpers.py @@ -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 diff --git a/headroom/proxy/prometheus_metrics.py b/headroom/proxy/prometheus_metrics.py index c479069f0..549b97387 100644 --- a/headroom/proxy/prometheus_metrics.py +++ b/headroom/proxy/prometheus_metrics.py @@ -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", diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index 78fcdd958..30fe69b41 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -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 [], diff --git a/headroom/proxy/ws_session_registry.py b/headroom/proxy/ws_session_registry.py index 167ce16d2..b28324cf2 100644 --- a/headroom/proxy/ws_session_registry.py +++ b/headroom/proxy/ws_session_registry.py @@ -50,6 +50,7 @@ TerminationCause = Literal[ "upstream_disconnect", "upstream_error", "client_error", + "client_cancel", "response_completed", "client_timeout", "unknown", diff --git a/headroom/transforms/kompress_compressor.py b/headroom/transforms/kompress_compressor.py index e8185c7a3..ba9dc99f9 100644 --- a/headroom/transforms/kompress_compressor.py +++ b/headroom/transforms/kompress_compressor.py @@ -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 diff --git a/tests/test_compression_store.py b/tests/test_compression_store.py index 764bd1671..38b95cb60 100644 --- a/tests/test_compression_store.py +++ b/tests/test_compression_store.py @@ -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 # ============================================================================= diff --git a/tests/test_openai_codex_ws_lifecycle.py b/tests/test_openai_codex_ws_lifecycle.py index eb3926085..d9bd3f316 100644 --- a/tests/test_openai_codex_ws_lifecycle.py +++ b/tests/test_openai_codex_ws_lifecycle.py @@ -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 diff --git a/tests/test_provider_proxy_routes.py b/tests/test_provider_proxy_routes.py index 3f4b7fae6..f4cd6263e 100644 --- a/tests/test_provider_proxy_routes.py +++ b/tests/test_provider_proxy_routes.py @@ -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: diff --git a/tests/test_proxy_codex_route_aliases.py b/tests/test_proxy_codex_route_aliases.py index 780619995..909a21129 100644 --- a/tests/test_proxy_codex_route_aliases.py +++ b/tests/test_proxy_codex_route_aliases.py @@ -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" diff --git a/tests/test_toin_feedback.py b/tests/test_toin_feedback.py index d5c537641..40d47d8b2 100644 --- a/tests/test_toin_feedback.py +++ b/tests/test_toin_feedback.py @@ -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."""