diff --git a/CHANGELOG.md b/CHANGELOG.md index ece1d2d7e..71b686432 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -68,6 +68,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Features +* **proxy:** add output shaping for OpenAI Responses traffic on `/v1/responses` HTTP requests and Codex WebSocket `response.create` frames, with stable output-savings holdout keys and counted WS token strata for the experiment. * **wrap:** `headroom wrap claude --1m` preserves the 1M context window. Behind a custom `ANTHROPIC_BASE_URL` (the proxy) Claude Code drops the `context-1m` beta header and caps the window at 200k for entitled subscription users; the opt-in flag sets `ANTHROPIC_MODEL=[1m]` on the launched process so the 1M window activates through Headroom. A model already selected via `ANTHROPIC_MODEL` is preserved (only the `[1m]` suffix is appended) ([#1158](https://github.com/chopratejas/headroom/issues/1158)). * **learn:** weight loops in `headroom learn`. A new loop detector (`headroom/learn/loops.py`) recognizes repeated tool-call patterns — including RTK re-fetch loops, where RTK's output truncation makes the agent re-run larger-limit variants of a *successful* command — collapses output-limit variants to one signature, measures the wasted tokens, surfaces loops as a highest-priority digest section, and weights loop guardrails above one-off rules by their measured waste. Previously loops had no special weight and a no-failure re-fetch loop was skipped entirely. Adds an RTK-loop eval (`benchmarks/rtk_loop_learn_eval.py`) that reproduces a loop, runs it through Learn, and asserts the generated guardrail ranks first and prevents re-triggering. * **learn:** write per-project learnings to the personal, gitignored `CLAUDE.local.md` by default instead of the team-shared `CLAUDE.md`, matching Claude Code's memory convention so machine-specific paths and tool-discovery byproducts no longer pollute the shared file. Adds a `--target` flag to override the destination (e.g. `--target CLAUDE.md` to opt back into the shared file, or any custom path), and auto-migrates a stale learned-patterns block out of an existing `CLAUDE.md` into `CLAUDE.local.md` with a warning ([#1072](https://github.com/chopratejas/headroom/issues/1072)). @@ -82,6 +83,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Bug Fixes +* **proxy:** forward Codex Desktop `/v1/responses` posts byte-faithfully so they stop returning upstream `400 {"detail":"Bad Request"}`. `handle_openai_responses` decoded the inbound body to inspect it but always re-serialized a canonical body on the way out, and it never stripped the inbound `content-encoding` header — so a `content-encoding: zstd` Codex Desktop request was forwarded as already-decoded JSON still advertising `zstd`, and the upstream ChatGPT Codex endpoint rejected it. The handler now keeps the original decoded bytes and forwards them verbatim whenever nothing (compression or memory injection) mutated the request, and drops the stale `content-encoding` header, mirroring the byte-faithful passthrough the chat and Anthropic paths already use ([#1542](https://github.com/headroomlabs-ai/headroom/issues/1542)). * **wrap/codex:** `headroom unwrap codex` now removes the Headroom rtk instruction block from the Codex global `AGENTS.md`. `wrap codex` injects it there, but unwrap only restored `config.toml` and MCP state, so a plain `codex` launch kept following the "prefix shell commands with `rtk`" guidance and failed once the managed rtk binary was off PATH. Unwrap now strips the marker-fenced block (preserving the rest of the file), mirroring `unwrap copilot` ([#1421](https://github.com/headroomlabs-ai/headroom/issues/1421)). * **proxy/auth:** classify real Anthropic OAuth tokens correctly. `classify_auth_mode` matched OAuth on the `sk-ant-oat-` prefix, but real access tokens are `sk-ant-oat01-...` (a version number, no dash after `oat`), so every real subscription/OAuth token fell through to the `sk-` branch and was tagged `PAYG` — enabling aggressive lossy compression, auto `cache_control`, and `prompt_cache_key` injection on subscription-bound requests the classifier is meant to route to the passthrough-prefer path. The prefix is now the dash-less `sk-ant-oat` (still matches the legacy dashed shape). The existing parity tests only passed because they used a synthetic `sk-ant-oat-01-` fixture; a regression test now covers the real `sk-ant-oat01-` format. * **install:** stop leaking a file descriptor on every `headroom install start`. `start_detached_agent()` opened the agent log file and handed it to `subprocess.Popen` but never closed the parent's copy, so each call leaked one fd (and pinned the log file open against rotation). The parent now closes its copy in a `try/finally` once the child has inherited it — the close also runs if `Popen` raises ([#1554](https://github.com/headroomlabs-ai/headroom/issues/1554)). diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 72593b4f3..480490930 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -557,6 +557,133 @@ def _responses_input_to_waste_messages(instructions: Any, input_data: Any) -> li return messages +def _output_shaping_holdout_fraction() -> float: + from headroom.proxy import runtime_env + + try: + return float(runtime_env.getenv("HEADROOM_OUTPUT_HOLDOUT", "0") or "0") + except ValueError: + return 0.0 + + +def _shape_openai_responses_for_output( + payload: dict[str, Any], + *, + input_tokens: int, + model: str, + conversation_key: str | None = None, +) -> Any: + """Apply OpenAI Responses output shaping and attach holdout labels.""" + from headroom.proxy.output_savings import ( + assign_arm, + conversation_key_from_body, + stratum_key, + stratum_label, + ) + from headroom.proxy.output_shaper import ( + OutputShaperSettings, + ShapeResult, + classify_openai_responses_input, + resolve_verbosity_level, + shape_openai_responses_request, + ) + + settings = OutputShaperSettings.from_env() + result = ShapeResult() + if not settings.enabled: + return result + + assert result.labels is not None + key = conversation_key or conversation_key_from_body(payload) + arm = assign_arm(key, _output_shaping_holdout_fraction()) + turn_kind = classify_openai_responses_input(payload.get("input")).value + stratum = stratum_key( + turn_kind=turn_kind, + input_tokens=input_tokens, + model=model or str(payload.get("model") or ""), + has_tools=bool(payload.get("tools")), + ) + result.labels.append(stratum_label(arm, stratum)) + if arm == "control": + return result + + level, _source = resolve_verbosity_level(settings) + shaped = shape_openai_responses_request( + payload, + settings=settings, + level_override=level, + ) + shaped.labels = [*result.labels, *(shaped.labels or [])] + return shaped + + +def _append_unique_transforms(transforms: list[str], labels: list[str] | None) -> None: + for label in labels or []: + if label not in transforms: + transforms.append(label) + + +def _openai_responses_payload_input_tokens( + payload: dict[str, Any], + token_provider: Any, +) -> int: + try: + tokenizer = token_provider.get_token_counter(str(payload.get("model") or "")) + return max(0, int(tokenizer.count_text(_json_debug_dumps(payload)))) + except Exception: + return max(0, _json_byte_len(payload) // 4) + + +def _openai_response_create_frame_input_tokens( + raw_msg: str, + token_provider: Any, +) -> int: + try: + parsed = json.loads(raw_msg) + except json.JSONDecodeError: + return 0 + if not isinstance(parsed, dict) or parsed.get("type") != "response.create": + return 0 + payload = parsed.get("response") if isinstance(parsed.get("response"), dict) else parsed + if not isinstance(payload, dict): + return 0 + return _openai_responses_payload_input_tokens(payload, token_provider) + + +def _shape_openai_response_create_frame( + raw_msg: str, + *, + input_tokens: int, + conversation_key: str | None = None, +) -> tuple[str, bool, list[str], str | None]: + try: + parsed = json.loads(raw_msg) + except json.JSONDecodeError: + return raw_msg, False, [], "non_json" + if not isinstance(parsed, dict) or parsed.get("type") != "response.create": + return raw_msg, False, [], "not_response_create" + + wrapped = isinstance(parsed.get("response"), dict) + payload = parsed["response"] if wrapped else parsed + if not isinstance(payload, dict): + return raw_msg, False, [], "invalid_inner_payload" + + result = _shape_openai_responses_for_output( + payload, + input_tokens=input_tokens, + model=str(payload.get("model") or ""), + conversation_key=conversation_key, + ) + labels = list(result.labels or []) + if not result.changed: + return raw_msg, False, labels, None + + if wrapped: + parsed["response"] = payload + return json.dumps(parsed), True, labels, None + return json.dumps(payload), True, labels, None + + def _openai_responses_context_budget(payload: dict[str, Any]) -> dict[str, Any]: payload_bytes = _json_byte_len(payload) buckets: dict[str, int] = {} @@ -3041,7 +3168,11 @@ class OpenAIHandlerMixin: }, ) - # Parse request + # Parse request. Keep the original (post-content-decoding) bytes so a + # request we never mutate is forwarded byte-for-byte instead of being + # canonically re-serialized. Codex Desktop posts whose body differs + # from our re-serialization are rejected upstream with HTTP 400 + # (#1542); byte-faithful passthrough avoids that. try: body, original_body_bytes = await read_request_json_with_bytes(request) except (json.JSONDecodeError, ValueError) as e: @@ -3107,6 +3238,11 @@ class OpenAIHandlerMixin: # Cloudflare Workers forward "br, zstd" which OpenAI may honor; # if httpx lacks brotli support the response body is undecipherable → 502. headers.pop("accept-encoding", None) + # Strip content-encoding: read_request_json_with_bytes already decoded + # the inbound body (zstd/gzip/deflate/br), so the bytes we forward are + # plain. Leaving a stale content-encoding header makes the upstream try + # to decompress already-decoded JSON and reject it with HTTP 400 (#1542). + headers.pop("content-encoding", None) tags = extract_tags(headers) client = classify_client(headers) # PR-A5 (P5-49): strip internal x-headroom-* from upstream-bound @@ -3520,10 +3656,31 @@ class OpenAIHandlerMixin: }, ) from _e - capture_codex_wire_debug( - "http_upstream_request", - request_id=request_id, - transport="http", + if not _bypass: + _http_conversation_key = request.headers.get("x-headroom-session-id") + _shape_result = _shape_openai_responses_for_output( + body, + input_tokens=original_tokens, + model=str(model or ""), + conversation_key=( + f"header:x-headroom-session-id:{_http_conversation_key}" + if _http_conversation_key + else None + ), + ) + _append_unique_transforms(transforms_applied, _shape_result.labels) + if _shape_result.changed: + body_mutation_tracker.mark_mutated("responses_output_shaping") + logger.info( + "[%s] /v1/responses output shaping labels=%s", + request_id, + _shape_result.labels, + ) + + capture_codex_wire_debug( + "http_upstream_request", + request_id=request_id, + transport="http", direction="headroom_to_upstream", method="POST", url=url, @@ -4667,6 +4824,7 @@ class OpenAIHandlerMixin: # internal errors), but we wrap the call site in try/except # anyway so a JSON-shape edge case can never break the WS # session. + first_frame_rewritten = False if self.config.optimize and not _ws_bypass: _first_frame_compression_elapsed_ms = 0.0 try: @@ -4763,6 +4921,7 @@ class OpenAIHandlerMixin: transforms_applied, ) ws_frames_compressed += 1 + first_frame_rewritten = True else: _log_ws_passthrough( _ws_reason or "no_compression", @@ -4858,6 +5017,32 @@ class OpenAIHandlerMixin: else "unknown", ) + if not _ws_bypass: + ( + first_msg_raw, + _shape_modified, + _shape_labels, + _shape_reason, + ) = _shape_openai_response_create_frame( + first_msg_raw, + input_tokens=_openai_response_create_frame_input_tokens( + first_msg_raw, + self.openai_provider, + ), + conversation_key=f"ws:{session_id}", + ) + _append_unique_transforms(transforms_applied, _shape_labels) + if _shape_modified: + if not first_frame_rewritten: + ws_frames_compressed += 1 + first_frame_rewritten = True + logger.info( + "[%s] WS /v1/responses output shaping frame=%d labels=%s", + request_id, + 1, + _shape_labels, + ) + _first_upstream_body: Any = None try: _first_upstream_body = json.loads(first_msg_raw) @@ -5105,6 +5290,7 @@ 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_frames_compressed nonlocal ws_last_client_frame_type, ws_client_disconnect_seen client_frame_index = 1 try: @@ -5166,6 +5352,35 @@ class OpenAIHandlerMixin: msg, frame_index=client_frame_index, ) + if not _ws_bypass: + ( + msg, + _shape_modified, + _shape_labels, + _shape_reason, + ) = _shape_openai_response_create_frame( + msg, + input_tokens=_openai_response_create_frame_input_tokens( + msg, + self.openai_provider, + ), + conversation_key=f"ws:{session_id}", + ) + _append_unique_transforms( + transforms_applied, + _shape_labels, + ) + if _shape_modified: + if not _frame_modified: + ws_frames_compressed += 1 + _frame_modified = True + logger.info( + "[%s] WS /v1/responses output shaping frame=%d labels=%s", + request_id, + client_frame_index, + _shape_labels, + ) + _outbound_frame_body: Any = None try: _outbound_frame_body = json.loads(msg) diff --git a/headroom/proxy/output_savings.py b/headroom/proxy/output_savings.py index bcff07c7e..ebe5ba44a 100644 --- a/headroom/proxy/output_savings.py +++ b/headroom/proxy/output_savings.py @@ -42,7 +42,7 @@ import hashlib import json import math from dataclasses import asdict, dataclass, field -from typing import Any +from typing import Any, cast # Coarse input-token buckets. Coarse on purpose: too many strata make # per-stratum baselines sparse and noisy. Boundaries in tokens. @@ -97,6 +97,50 @@ def stratum_key( ) +def _unwrap_response_create_body(body: dict[str, Any]) -> dict[str, Any]: + response = body.get("response") + if body.get("type") == "response.create" and isinstance(response, dict): + return cast("dict[str, Any]", response) + return body + + +def _stable_response_identifier(body: dict[str, Any]) -> str: + def _string_value(value: Any) -> str: + if isinstance(value, str): + return value + if isinstance(value, dict): + for key in ("id", "conversation_id", "session_id", "thread_id"): + nested = value.get(key) + if isinstance(nested, str) and nested: + return nested + return "" + + for key in ("conversation", "conversation_id", "session_id", "thread_id"): + value = _string_value(body.get(key)) + if value and value.lower() != "auto": + return f"{key}:{value}" + + for container_key in ("client_metadata", "metadata"): + container = body.get(container_key) + if not isinstance(container, dict): + continue + for key in ( + "conversation_id", + "conversation_key", + "session_id", + "thread_id", + "codex_session_id", + ): + value = _string_value(container.get(key)) + if value and value.lower() != "auto": + return f"{container_key}.{key}:{value}" + + instructions = body.get("instructions") + if isinstance(instructions, str) and instructions: + return f"instructions:{instructions[:512]}" + return "" + + def conversation_key_from_body(body: dict[str, Any]) -> str: """Derive a conversation-stable key for holdout assignment. @@ -105,6 +149,7 @@ def conversation_key_from_body(body: dict[str, Any]) -> str: message's text. The first user turn is immutable for a conversation's lifetime, which is exactly the stability we need. """ + body = _unwrap_response_create_body(body) model = str(body.get("model", "")) seed = model for msg in body.get("messages", []): @@ -118,6 +163,12 @@ def conversation_key_from_body(body: dict[str, Any]) -> str: seed += "\x00" + str(block.get("text", ""))[:512] break break + if "input" in body: + stable_response_key = _stable_response_identifier(body) + if stable_response_key: + seed += "\x00" + stable_response_key + elif not body.get("messages"): + seed += "\x00responses" return hashlib.sha256(seed.encode("utf-8", "ignore")).hexdigest() diff --git a/headroom/proxy/output_shaper.py b/headroom/proxy/output_shaper.py index abf61e9c9..e286d4204 100644 --- a/headroom/proxy/output_shaper.py +++ b/headroom/proxy/output_shaper.py @@ -49,6 +49,17 @@ LEGACY_THINKING_FLOOR = 1024 # Ordering for output_config.effort values. Unknown values are left alone. _EFFORT_RANK = {"low": 0, "medium": 1, "high": 2, "xhigh": 3, "max": 4} +_TEXT_VERBOSITY_RANK = {"low": 0, "medium": 1, "high": 2} + +_OPENAI_RESPONSES_OUTPUT_ITEM_TYPES = frozenset( + { + "custom_tool_call_output", + "function_call_output", + "local_shell_call_output", + "apply_patch_call_output", + } +) + # Sentinel prefix marks the steering block so application is idempotent and # the block is recognizable in logs/diffs. _STEERING_SENTINEL = "" @@ -250,6 +261,22 @@ def steering_text(level: int) -> str | None: return f"{_STEERING_SENTINEL}\n{text}\n{_STEERING_SUFFIX}" +def _replace_or_append_steering_block(existing: str, block: str) -> tuple[str, bool]: + """Replace an existing steering block in text, or append one at the tail.""" + start = existing.find(_STEERING_SENTINEL) + if start >= 0: + end = existing.find(_STEERING_SUFFIX, start) + end = len(existing) if end < 0 else end + len(_STEERING_SUFFIX) + prefix = existing[:start].rstrip() + suffix = existing[end:].lstrip("\n") + parts = [part for part in (prefix, block, suffix) if part] + updated = "\n\n".join(parts) + return updated, updated != existing + + updated = f"{existing.rstrip()}\n\n{block}" if existing.strip() else block + return updated, updated != existing + + def apply_verbosity_steering(body: dict[str, Any], level: int) -> bool: """Append the steering block to the tail of the system prompt. @@ -325,6 +352,181 @@ def route_effort( return labels +def _responses_part_text(value: Any) -> str: + if isinstance(value, str): + return value + if isinstance(value, list): + texts: list[str] = [] + for part in value: + if isinstance(part, str): + texts.append(part) + elif isinstance(part, dict) and isinstance(part.get("text"), str): + texts.append(part["text"]) + return "\n".join(text for text in texts if text) + return "" + + +def _responses_user_signal(item: dict[str, Any]) -> bool: + item_type = item.get("type") + role = item.get("role") + if role == "user": + content = item.get("content") + if isinstance(content, list): + for part in content: + if isinstance(part, dict) and part.get("type") in { + "input_file", + "input_image", + }: + return True + text = _responses_part_text(content) + return bool(text.strip()) + if item_type == "input_text": + text = _responses_part_text(item.get("text")) + return bool(text.strip()) + if item_type == "input_image": + return True + return False + + +def classify_openai_responses_input(input_data: Any) -> TurnKind: + """Classify OpenAI Responses ``input`` without content heuristics.""" + if isinstance(input_data, str): + return TurnKind.NEW_USER_ASK if input_data.strip() else TurnKind.UNKNOWN + if not isinstance(input_data, list) or not input_data: + return TurnKind.UNKNOWN + + saw_tool_output = False + saw_unknown = False + for item in input_data: + if not isinstance(item, dict): + saw_unknown = True + continue + item_type = item.get("type") + if item_type in _OPENAI_RESPONSES_OUTPUT_ITEM_TYPES: + saw_tool_output = True + continue + if _responses_user_signal(item): + return TurnKind.NEW_USER_ASK + if item_type in {"message", "function_call", "reasoning"}: + continue + saw_unknown = True + + if saw_tool_output and not saw_unknown: + return TurnKind.MECHANICAL_CONTINUATION + return TurnKind.UNKNOWN + + +def apply_openai_responses_verbosity_steering( + body: dict[str, Any], + level: int, +) -> bool: + """Append or replace steering in OpenAI Responses ``instructions``.""" + text = steering_text(level) + if text is None: + return False + + instructions = body.get("instructions") + if instructions is None: + body["instructions"] = text + return True + if not isinstance(instructions, str): + return False + + updated, changed = _replace_or_append_steering_block(instructions, text) + if changed: + body["instructions"] = updated + return changed + + +def route_openai_reasoning_effort( + body: dict[str, Any], + kind: TurnKind, + settings: OutputShaperSettings, +) -> list[str]: + """Lower explicitly-present OpenAI reasoning effort on mechanical turns.""" + if kind is not TurnKind.MECHANICAL_CONTINUATION: + return [] + + reasoning = body.get("reasoning") + if not isinstance(reasoning, dict): + return [] + effort = reasoning.get("effort") + target = settings.mechanical_effort + if ( + isinstance(effort, str) + and effort in _EFFORT_RANK + and target in _EFFORT_RANK + and _EFFORT_RANK[effort] > _EFFORT_RANK[target] + ): + reasoning["effort"] = target + return [f"output_shaper:reasoning_effort:{effort}->{target}"] + return [] + + +def route_openai_text_verbosity(body: dict[str, Any]) -> list[str]: + """Set or lower OpenAI ``text.verbosity`` conservatively.""" + model = str(body.get("model") or "").lower() + text_config = body.get("text") + can_create = model.startswith("gpt-5") + if text_config is None: + if not can_create: + return [] + body["text"] = {"verbosity": "low"} + return ["output_shaper:text_verbosity:unset->low"] + if not isinstance(text_config, dict): + return [] + + verbosity = text_config.get("verbosity") + if verbosity is None: + if not can_create: + return [] + text_config["verbosity"] = "low" + return ["output_shaper:text_verbosity:unset->low"] + if ( + isinstance(verbosity, str) + and verbosity in _TEXT_VERBOSITY_RANK + and _TEXT_VERBOSITY_RANK[verbosity] > _TEXT_VERBOSITY_RANK["low"] + ): + text_config["verbosity"] = "low" + return [f"output_shaper:text_verbosity:{verbosity}->low"] + return [] + + +def shape_openai_responses_request( + body: dict[str, Any], + settings: OutputShaperSettings | None = None, + level_override: int | None = None, +) -> ShapeResult: + """Apply OpenAI Responses output-shaping levers in place.""" + if settings is None: + settings = OutputShaperSettings.from_env() + result = ShapeResult() + if not settings.enabled: + return result + + assert result.labels is not None # __post_init__ guarantees + + level = settings.verbosity_level if level_override is None else level_override + if level > 0 and apply_openai_responses_verbosity_steering(body, level): + result.changed = True + result.labels.append(f"output_shaper:verbosity:L{level}") + + kind = classify_openai_responses_input(body.get("input")) + if settings.effort_router_enabled: + labels = route_openai_reasoning_effort(body, kind, settings) + if labels: + result.changed = True + result.labels.extend(labels) + logger.debug("OpenAIOutputShaper: turn=%s mutations=%s", kind.value, labels) + + labels = route_openai_text_verbosity(body) + if labels: + result.changed = True + result.labels.extend(labels) + + return result + + def shape_request( body: dict[str, Any], settings: OutputShaperSettings | None = None, diff --git a/tests/test_codex_responses_passthrough_bytes.py b/tests/test_codex_responses_passthrough_bytes.py new file mode 100644 index 000000000..af2158999 --- /dev/null +++ b/tests/test_codex_responses_passthrough_bytes.py @@ -0,0 +1,117 @@ +"""Byte-faithful passthrough for Codex Desktop /v1/responses posts (issue #1542). + +Codex Desktop sends ``POST /v1/responses`` with ``content-encoding: zstd``. The +handler decodes the body to parse it, but when nothing mutates the request it +must forward the *original decoded bytes* verbatim and must not re-advertise the +stale ``content-encoding`` header. Otherwise the upstream ChatGPT Codex endpoint +either re-canonicalizes a body it rejects, or tries to zstd-decode already-decoded +JSON — both surface to the client as ``400 {"detail":"Bad Request"}``. +""" + +from __future__ import annotations + +import json + +import pytest + +pytest.importorskip("fastapi") +pytest.importorskip("httpx") + +import httpx +from fastapi.testclient import TestClient + +from headroom.proxy.loopback_guard import require_loopback +from headroom.proxy.server import ProxyConfig, create_app + + +def _make_client(optimize: bool = False): + config = ProxyConfig( + optimize=optimize, + cache_enabled=False, + rate_limit_enabled=False, + cost_tracking_enabled=False, + ) + app = create_app(config) + app.dependency_overrides[require_loopback] = lambda: None + return app + + +def _fake_upstream_response(url: str) -> httpx.Response: + return httpx.Response( + 200, + json={ + "id": "resp_test", + "object": "response", + "output": [], + "usage": {"input_tokens": 12, "output_tokens": 3}, + }, + request=httpx.Request("POST", url), + ) + + +def _patch_capture(app): + """Replace the server's upstream forwarder with a capturing stub.""" + captured: dict = {} + server = app.state.proxy + + async def fake_retry(method, url, headers, body, stream=False, **kwargs): + captured["method"] = method + captured["url"] = url + captured["headers"] = dict(headers) + captured["body"] = body + captured["kwargs"] = kwargs + return _fake_upstream_response(url) + + server._retry_request = fake_retry + return captured + + +def test_unmutated_zstd_post_forwards_decoded_bytes_and_strips_content_encoding(): + zstandard = pytest.importorskip("zstandard") + app = _make_client(optimize=False) + + payload = { + "model": "gpt-5-codex", + "input": "list the files in this repo", + "instructions": "be terse", + } + raw = json.dumps(payload).encode("utf-8") + compressed = zstandard.ZstdCompressor().compress(raw) + + with TestClient(app) as client: + captured = _patch_capture(app) + resp = client.post( + "/v1/responses", + headers={ + "Authorization": "Bearer sk-test", + "Content-Type": "application/json", + "Content-Encoding": "zstd", + "originator": "codex_desktop", + }, + content=compressed, + ) + + assert resp.status_code == 200 + # Nothing mutated the request -> byte-faithful passthrough engages. + assert captured["kwargs"].get("body_mutated") is False + assert captured["kwargs"].get("original_body_bytes") == raw + # The stale content-encoding must not ride along with already-decoded bytes. + fwd_headers = {k.lower(): v for k, v in captured["headers"].items()} + assert "content-encoding" not in fwd_headers + + +def test_unmutated_plain_post_passes_original_bytes_through(): + app = _make_client(optimize=False) + raw = json.dumps({"model": "gpt-5-codex", "input": "hi"}).encode("utf-8") + + with TestClient(app) as client: + captured = _patch_capture(app) + resp = client.post( + "/v1/responses", + headers={"Authorization": "Bearer sk-test", "Content-Type": "application/json"}, + content=raw, + ) + + assert resp.status_code == 200 + assert captured["kwargs"].get("body_mutated") is False + assert captured["kwargs"].get("original_body_bytes") == raw diff --git a/tests/test_openai_codex_ws_lifecycle.py b/tests/test_openai_codex_ws_lifecycle.py index 059811903..10642bd23 100644 --- a/tests/test_openai_codex_ws_lifecycle.py +++ b/tests/test_openai_codex_ws_lifecycle.py @@ -24,6 +24,11 @@ from headroom.proxy.ws_session_registry import WebSocketSessionRegistry # --------------------------------------------------------------------------- +class _TokenCounter: + def count_text(self, text: str) -> int: + return len(text.split()) + + class _DummyMetrics: def __init__(self) -> None: self.active_ws_sessions = 0 @@ -73,7 +78,10 @@ class _DummyOpenAIHandler(OpenAIHandlerMixin): connect_timeout_seconds=10, ) self.usage_reporter = None - self.openai_provider = SimpleNamespace(get_context_limit=lambda model: 128_000) + self.openai_provider = SimpleNamespace( + get_context_limit=lambda model: 128_000, + get_token_counter=lambda model: _TokenCounter(), + ) self.openai_pipeline = SimpleNamespace(apply=MagicMock()) self.anthropic_backend = None self.cost_tracker = None @@ -300,6 +308,148 @@ def _codex_lite_headers(*, chatgpt: bool) -> dict[str, str]: return headers +@pytest.mark.asyncio +async def test_ws_first_frame_output_shaper_rewrites_without_compression(monkeypatch): + monkeypatch.setenv("HEADROOM_OUTPUT_SHAPER", "1") + monkeypatch.setenv("HEADROOM_VERBOSITY_LEVEL", "2") + monkeypatch.delenv("HEADROOM_OUTPUT_HOLDOUT", raising=False) + upstream_events = [ + json.dumps({"type": "response.created", "response": {"id": "r_1"}}), + json.dumps( + { + "type": "response.completed", + "response": { + "id": "r_1", + "usage": {"input_tokens": 10, "output_tokens": 1}, + }, + } + ), + ] + upstream = _FakeUpstream(upstream_events) + fake_ws_mod = _make_fake_websockets_module(upstream) + client_ws = _FakeWebSocket(frames=[_first_frame()]) + handler = _DummyOpenAIHandler() + handler.config.optimize = False + outcomes = [] + + async def _record_request_outcome(outcome): + outcomes.append(outcome) + + handler._record_request_outcome = _record_request_outcome + + with patch.dict(sys.modules, {"websockets": fake_ws_mod}): + await handler.handle_openai_responses_ws(client_ws) + + sent = json.loads(upstream.sent[0]) + payload = sent["response"] + assert "" in payload["instructions"] + assert payload["text"]["verbosity"] == "low" + assert any( + t == "output_shaper:verbosity:L2" + for t in outcomes[-1].transforms_applied + ) + + +@pytest.mark.asyncio +async def test_ws_output_shaper_stratum_uses_frame_input_tokens(monkeypatch): + monkeypatch.setenv("HEADROOM_OUTPUT_SHAPER", "1") + monkeypatch.setenv("HEADROOM_VERBOSITY_LEVEL", "2") + long_input = " ".join(f"word{i}" for i in range(2500)) + first_frame = json.dumps( + { + "type": "response.create", + "response": {"model": "gpt-5.4", "input": long_input}, + } + ) + upstream_events = [ + json.dumps({"type": "response.created", "response": {"id": "r_1"}}), + json.dumps( + { + "type": "response.completed", + "response": { + "id": "r_1", + "usage": {"input_tokens": 3000, "output_tokens": 1}, + }, + } + ), + ] + upstream = _FakeUpstream(upstream_events) + fake_ws_mod = _make_fake_websockets_module(upstream) + client_ws = _FakeWebSocket(frames=[first_frame]) + handler = _DummyOpenAIHandler() + outcomes = [] + + async def _record_request_outcome(outcome): + outcomes.append(outcome) + + handler._record_request_outcome = _record_request_outcome + + with patch.dict(sys.modules, {"websockets": fake_ws_mod}): + await handler.handle_openai_responses_ws(client_ws) + + transforms = outcomes[-1].transforms_applied + assert any(t.startswith("output_shaper:stratum:gpt|new_user_ask|s|") for t in transforms) + assert not any(t.startswith("output_shaper:stratum:gpt|new_user_ask|xs|") for t in transforms) + + +@pytest.mark.asyncio +async def test_ws_output_shaper_respects_bypass(monkeypatch): + monkeypatch.setenv("HEADROOM_OUTPUT_SHAPER", "1") + upstream_events = [ + json.dumps({"type": "response.created", "response": {"id": "r_1"}}), + json.dumps({"type": "response.completed", "response": {"id": "r_1"}}), + ] + upstream = _FakeUpstream(upstream_events) + fake_ws_mod = _make_fake_websockets_module(upstream) + first = _first_frame() + client_ws = _FakeWebSocket(frames=[first]) + client_ws.headers = { + "authorization": "Bearer test", + "x-headroom-bypass": "true", + } + handler = _DummyOpenAIHandler() + + with patch.dict(sys.modules, {"websockets": fake_ws_mod}): + await handler.handle_openai_responses_ws(client_ws) + + assert upstream.sent[0] == first + + +@pytest.mark.asyncio +async def test_ws_output_shaper_holdout_labels_without_rewrite(monkeypatch): + monkeypatch.setenv("HEADROOM_OUTPUT_SHAPER", "1") + monkeypatch.setenv("HEADROOM_OUTPUT_HOLDOUT", "1") + upstream_events = [ + json.dumps({"type": "response.created", "response": {"id": "r_1"}}), + json.dumps( + { + "type": "response.completed", + "response": { + "id": "r_1", + "usage": {"input_tokens": 10, "output_tokens": 1}, + }, + } + ), + ] + upstream = _FakeUpstream(upstream_events) + fake_ws_mod = _make_fake_websockets_module(upstream) + first = _first_frame() + client_ws = _FakeWebSocket(frames=[first]) + handler = _DummyOpenAIHandler() + outcomes = [] + + async def _record_request_outcome(outcome): + outcomes.append(outcome) + + handler._record_request_outcome = _record_request_outcome + + with patch.dict(sys.modules, {"websockets": fake_ws_mod}): + await handler.handle_openai_responses_ws(client_ws) + + assert upstream.sent[0] == first + transforms = outcomes[-1].transforms_applied + assert any(t.startswith("output_shaper:control:") for t in transforms) + assert not any(t == "output_shaper:verbosity:L2" for t in transforms) # --------------------------------------------------------------------------- # Tests # --------------------------------------------------------------------------- diff --git a/tests/test_openai_responses_output_shaper.py b/tests/test_openai_responses_output_shaper.py new file mode 100644 index 000000000..31104ecca --- /dev/null +++ b/tests/test_openai_responses_output_shaper.py @@ -0,0 +1,162 @@ +from __future__ import annotations + +import copy +from typing import Any + +import httpx +import pytest + +pytest.importorskip("fastapi") + +from fastapi.testclient import TestClient # noqa: E402 + +from headroom.proxy.loopback_guard import require_loopback # noqa: E402 +from headroom.proxy.server import ProxyConfig, create_app # noqa: E402 + + +def _make_client() -> TestClient: + app = create_app( + ProxyConfig( + optimize=False, + cache_enabled=False, + rate_limit_enabled=False, + cost_tracking_enabled=False, + log_requests=False, + http2=False, + ) + ) + app.dependency_overrides[require_loopback] = lambda: None + return TestClient(app) + + +async def _ok_response( + method: str, + url: str, + headers: dict[str, str], + body: dict[str, Any], + stream: bool = False, + **kwargs: Any, +) -> httpx.Response: + return httpx.Response( + 200, + json={ + "id": "resp_1", + "output": [], + "usage": {"input_tokens": 10, "output_tokens": 1}, + }, + ) + + +def test_http_responses_output_shaper_rewrites_and_labels(monkeypatch): + monkeypatch.setenv("HEADROOM_OUTPUT_SHAPER", "1") + monkeypatch.setenv("HEADROOM_VERBOSITY_LEVEL", "2") + monkeypatch.delenv("HEADROOM_OUTPUT_HOLDOUT", raising=False) + captured: dict[str, Any] = {} + outcomes: list[Any] = [] + + payload = { + "model": "gpt-5", + "input": [ + { + "type": "function_call_output", + "call_id": "call_1", + "output": "ok", + } + ], + "reasoning": {"effort": "xhigh"}, + "text": {"verbosity": "medium"}, + } + + with _make_client() as client: + proxy = client.app.state.proxy + + async def _fake_retry(*args: Any, **kwargs: Any) -> httpx.Response: + body = args[3] + captured["body"] = copy.deepcopy(body) + captured["retry_kwargs"] = dict(kwargs) + return await _ok_response(*args, **kwargs) + + async def _record_request_outcome(outcome: Any) -> None: + outcomes.append(outcome) + + proxy._retry_request = _fake_retry + proxy._record_request_outcome = _record_request_outcome + + response = client.post( + "/v1/responses", + headers={"authorization": "Bearer test-key"}, + json=payload, + ) + + assert response.status_code == 200 + sent = captured["body"] + assert "" in sent["instructions"] + assert sent["reasoning"]["effort"] == "low" + assert sent["text"]["verbosity"] == "low" + assert captured["retry_kwargs"]["body_mutated"] is True + assert captured["retry_kwargs"]["original_body_bytes"] is not None + transforms = outcomes[-1].transforms_applied + assert any(t.startswith("output_shaper:stratum:") for t in transforms) + assert "output_shaper:verbosity:L2" in transforms + assert "output_shaper:reasoning_effort:xhigh->low" in transforms + assert "output_shaper:text_verbosity:medium->low" in transforms + + +def test_http_responses_output_shaper_respects_bypass(monkeypatch): + monkeypatch.setenv("HEADROOM_OUTPUT_SHAPER", "1") + captured: dict[str, Any] = {} + payload = {"model": "gpt-5", "input": "hi"} + + with _make_client() as client: + proxy = client.app.state.proxy + + async def _fake_retry(*args: Any, **kwargs: Any) -> httpx.Response: + captured["body"] = copy.deepcopy(args[3]) + return await _ok_response(*args, **kwargs) + + proxy._retry_request = _fake_retry + + response = client.post( + "/v1/responses", + headers={ + "authorization": "Bearer test-key", + "x-headroom-bypass": "true", + }, + json=payload, + ) + + assert response.status_code == 200 + assert captured["body"] == payload + + +def test_http_responses_output_shaper_holdout_labels_without_rewrite(monkeypatch): + monkeypatch.setenv("HEADROOM_OUTPUT_SHAPER", "1") + monkeypatch.setenv("HEADROOM_OUTPUT_HOLDOUT", "1") + captured: dict[str, Any] = {} + outcomes: list[Any] = [] + payload = {"model": "gpt-5", "input": "hi"} + + with _make_client() as client: + proxy = client.app.state.proxy + + async def _fake_retry(*args: Any, **kwargs: Any) -> httpx.Response: + captured["body"] = copy.deepcopy(args[3]) + return await _ok_response(*args, **kwargs) + + async def _record_request_outcome(outcome: Any) -> None: + outcomes.append(outcome) + + proxy._retry_request = _fake_retry + proxy._record_request_outcome = _record_request_outcome + + response = client.post( + "/v1/responses", + headers={"authorization": "Bearer test-key"}, + json=payload, + ) + + assert response.status_code == 200 + assert captured["body"] == payload + transforms = outcomes[-1].transforms_applied + assert any(t.startswith("output_shaper:control:") for t in transforms) + assert "output_shaper:verbosity:L2" not in transforms diff --git a/tests/test_output_savings.py b/tests/test_output_savings.py index fef965687..e9d67a826 100644 --- a/tests/test_output_savings.py +++ b/tests/test_output_savings.py @@ -89,6 +89,71 @@ class TestArmAssignment: b = {"model": "m", "messages": [{"role": "user", "content": "task B"}]} assert conversation_key_from_body(a) != conversation_key_from_body(b) + def test_conversation_key_uses_responses_stable_metadata(self): + a = { + "model": "gpt-5", + "client_metadata": {"session_id": "session-1"}, + "input": "task A", + } + b = { + "model": "gpt-5", + "client_metadata": {"session_id": "session-2"}, + "input": "task A", + } + assert conversation_key_from_body(a) != conversation_key_from_body(b) + + def test_conversation_key_does_not_use_responses_delta_text(self): + user_turn = { + "model": "gpt-5", + "instructions": "same session instructions", + "input": "task A", + } + tool_turn = { + "model": "gpt-5", + "instructions": "same session instructions", + "input": [ + { + "type": "function_call_output", + "call_id": "call_1", + "output": "ok", + } + ], + } + assert conversation_key_from_body(user_turn) == conversation_key_from_body(tool_turn) + + def test_conversation_key_unwraps_ws_response_create(self): + http_body = {"model": "gpt-5", "input": "build a cache"} + ws_body = { + "type": "response.create", + "response": {"model": "gpt-5", "input": "build a cache"}, + } + assert conversation_key_from_body(http_body) == conversation_key_from_body(ws_body) + + def test_conversation_key_uses_responses_conversation_id(self): + a = { + "model": "gpt-5", + "conversation": "conv_1", + "input": [ + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "task A"}], + } + ], + } + b = { + "model": "gpt-5", + "conversation": "conv_2", + "input": [ + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "task B"}], + } + ], + } + assert conversation_key_from_body(a) != conversation_key_from_body(b) + # --------------------------------------------------------------------------- # baseline model diff --git a/tests/test_output_shaper.py b/tests/test_output_shaper.py index 4472ff848..bff332c67 100644 --- a/tests/test_output_shaper.py +++ b/tests/test_output_shaper.py @@ -13,9 +13,14 @@ from headroom.proxy.output_shaper import ( LEGACY_THINKING_FLOOR, OutputShaperSettings, TurnKind, + apply_openai_responses_verbosity_steering, apply_verbosity_steering, + classify_openai_responses_input, classify_turn, route_effort, + route_openai_reasoning_effort, + route_openai_text_verbosity, + shape_openai_responses_request, shape_request, steering_text, ) @@ -282,3 +287,121 @@ class TestShapeRequest: settings = OutputShaperSettings.from_env() assert settings.verbosity_level == 4 assert settings.mechanical_effort == "low" + + +class TestOpenAIResponsesClassify: + def test_string_input_is_new_ask(self): + assert classify_openai_responses_input("explain this") == TurnKind.NEW_USER_ASK + + def test_function_call_output_only_is_mechanical(self): + input_data = [ + { + "type": "function_call_output", + "call_id": "call_1", + "output": "ok", + } + ] + assert ( + classify_openai_responses_input(input_data) + == TurnKind.MECHANICAL_CONTINUATION + ) + + def test_mixed_user_message_and_tool_output_is_new_ask(self): + input_data = [ + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "also check foo.py"}], + }, + { + "type": "function_call_output", + "call_id": "call_1", + "output": "ok", + }, + ] + assert classify_openai_responses_input(input_data) == TurnKind.NEW_USER_ASK + + +class TestOpenAIResponsesSteering: + def test_instructions_steering_is_idempotent_and_replaced(self): + body = {"instructions": f"System.\n\n{steering_text(1)}"} + + assert apply_openai_responses_verbosity_steering(body, 2) is True + assert body["instructions"].count("") == 1 + assert steering_text(1) not in body["instructions"] + assert steering_text(2) in body["instructions"] + + snapshot = copy.deepcopy(body) + assert apply_openai_responses_verbosity_steering(body, 2) is False + assert body == snapshot + + +class TestOpenAIResponsesReasoning: + def test_reasoning_effort_lowers_only_for_mechanical_continuations(self): + body = {"reasoning": {"effort": "xhigh"}} + labels = route_openai_reasoning_effort( + body, + TurnKind.MECHANICAL_CONTINUATION, + ENABLED, + ) + assert labels == ["output_shaper:reasoning_effort:xhigh->low"] + assert body["reasoning"]["effort"] == "low" + + new_ask = {"reasoning": {"effort": "xhigh"}} + assert route_openai_reasoning_effort(new_ask, TurnKind.NEW_USER_ASK, ENABLED) == [] + assert new_ask["reasoning"]["effort"] == "xhigh" + + def test_reasoning_effort_is_not_injected_when_absent(self): + body: dict[str, Any] = {} + labels = route_openai_reasoning_effort( + body, + TurnKind.MECHANICAL_CONTINUATION, + ENABLED, + ) + assert labels == [] + assert "reasoning" not in body + + +class TestOpenAIResponsesTextVerbosity: + def test_text_verbosity_set_for_gpt5_family(self): + body = {"model": "gpt-5.1"} + labels = route_openai_text_verbosity(body) + assert labels == ["output_shaper:text_verbosity:unset->low"] + assert body["text"] == {"verbosity": "low"} + + def test_text_verbosity_not_injected_for_non_gpt5(self): + body = {"model": "gpt-4o"} + assert route_openai_text_verbosity(body) == [] + assert "text" not in body + + def test_existing_text_verbosity_is_lowered_for_any_model(self): + body = {"model": "gpt-4o", "text": {"verbosity": "medium"}} + labels = route_openai_text_verbosity(body) + assert labels == ["output_shaper:text_verbosity:medium->low"] + assert body["text"]["verbosity"] == "low" + + def test_shape_openai_responses_combines_steering_native_knobs(self): + body = { + "model": "gpt-5", + "input": [ + { + "type": "function_call_output", + "call_id": "call_1", + "output": "ok", + } + ], + "instructions": "System.", + "reasoning": {"effort": "xhigh"}, + "text": {"verbosity": "medium"}, + } + result = shape_openai_responses_request(body, ENABLED) + + assert result.changed is True + assert result.labels == [ + "output_shaper:verbosity:L2", + "output_shaper:reasoning_effort:xhigh->low", + "output_shaper:text_verbosity:medium->low", + ] + assert steering_text(2) in body["instructions"] + assert body["reasoning"]["effort"] == "low" + assert body["text"]["verbosity"] == "low"