Fix Codex WS proxy (Issue #86), LiteLLM metrics, and backend forwarding

- Fix WS /v1/responses: forward Sec-WebSocket-Protocol (subprotocol) to
upstream instead of stripping it — root cause of Codex HTTP 500 errors
- Fix WS relay: handle binary messages properly instead of crashing on
.decode(), add debug logging instead of silent except:pass
- Add Authorization header fallback from OPENAI_API_KEY env var for WS
- Extract response body from websockets InvalidStatus for error debugging
- Fix streaming /v1/responses: pass optimized_tokens (not original_tokens
twice) so compression savings appear in streaming metrics
- Fix hardcoded provider="bedrock" in 4 metrics/log locations — now uses
self.anthropic_backend.name so LiteLLM backends report correctly
- Forward --backend, --anyllm-provider, --region flags from wrap commands
(codex, aider) to the proxy subprocess via _start_proxy()
- Forward API key from request headers to LiteLLM acompletion() calls
- Forward region to Vertex AI (vertex_location) not just Bedrock
- Redesign proxy startup banner: show routing table instead of misleading
"Backend: Anthropic" label

Closes #86

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

headroom/backends/litellm.py

@@ -662,9 +662,19 @@ class LiteLLMBackend(Backend):
                     )
                     kwargs["messages"].insert(0, {"role": "system", "content": system_text})
 
-            # AWS region for Bedrock
-            if self.provider == "bedrock" and self.region:
-                kwargs["aws_region_name"] = self.region
+            # Provider-specific region config
+            if self.region:
+                if self.provider == "bedrock":
+                    kwargs["aws_region_name"] = self.region
+                elif self.provider in ("vertex_ai", "vertex_ai_beta"):
+                    kwargs["vertex_location"] = self.region
+
+            # Forward API key from request headers if present
+            auth_header = headers.get("authorization", headers.get("Authorization", ""))
+            if auth_header.startswith("Bearer "):
+                kwargs["api_key"] = auth_header[7:]
+            elif headers.get("x-api-key"):
+                kwargs["api_key"] = headers["x-api-key"]
 
             logger.debug(f"LiteLLM request: model={litellm_model}")
 
@@ -752,8 +762,19 @@ class LiteLLMBackend(Backend):
                     )
                     kwargs["messages"].insert(0, {"role": "system", "content": system_text})
 
-            if self.provider == "bedrock" and self.region:
-                kwargs["aws_region_name"] = self.region
+            # Provider-specific region config
+            if self.region:
+                if self.provider == "bedrock":
+                    kwargs["aws_region_name"] = self.region
+                elif self.provider in ("vertex_ai", "vertex_ai_beta"):
+                    kwargs["vertex_location"] = self.region
+
+            # Forward API key from request headers if present
+            auth_header = headers.get("authorization", headers.get("Authorization", ""))
+            if auth_header.startswith("Bearer "):
+                kwargs["api_key"] = auth_header[7:]
+            elif headers.get("x-api-key"):
+                kwargs["api_key"] = headers["x-api-key"]
 
             msg_id = f"msg_{uuid.uuid4().hex[:24]}"
 
@@ -961,12 +982,19 @@ class LiteLLMBackend(Backend):
                 if param in body:
                     kwargs[param] = body[param]
 
-            # Provider-specific config
-            if self.provider == "bedrock" and self.region:
-                kwargs["aws_region_name"] = self.region
-            elif self.provider == "databricks":
-                # Databricks uses env vars for auth
-                pass
+            # Provider-specific region config
+            if self.region:
+                if self.provider == "bedrock":
+                    kwargs["aws_region_name"] = self.region
+                elif self.provider in ("vertex_ai", "vertex_ai_beta"):
+                    kwargs["vertex_location"] = self.region
+
+            # Forward API key from request headers if present
+            auth_header = headers.get("authorization", headers.get("Authorization", ""))
+            if auth_header.startswith("Bearer "):
+                kwargs["api_key"] = auth_header[7:]
+            elif headers.get("x-api-key"):
+                kwargs["api_key"] = headers["x-api-key"]
 
             logger.debug(f"LiteLLM OpenAI request: model={litellm_model}")
 
@@ -1086,8 +1114,19 @@ class LiteLLMBackend(Backend):
             if "stream_options" in body:
                 kwargs["stream_options"] = body["stream_options"]
 
-            if self.provider == "bedrock" and self.region:
-                kwargs["aws_region_name"] = self.region
+            # Provider-specific region config
+            if self.region:
+                if self.provider == "bedrock":
+                    kwargs["aws_region_name"] = self.region
+                elif self.provider in ("vertex_ai", "vertex_ai_beta"):
+                    kwargs["vertex_location"] = self.region
+
+            # Forward API key from request headers if present
+            auth_header = headers.get("authorization", headers.get("Authorization", ""))
+            if auth_header.startswith("Bearer "):
+                kwargs["api_key"] = auth_header[7:]
+            elif headers.get("x-api-key"):
+                kwargs["api_key"] = headers["x-api-key"]
 
             response = await acompletion(**kwargs)
 

headroom/cli/proxy.py

@@ -250,13 +250,12 @@ def proxy(
     if license_key:
         license_status = f"MANAGED (key={license_key[:8]}...)"
 
-    effective_region = bedrock_region or region
-    backend_status = "Anthropic (direct API)"
+    anthropic_url = config.anthropic_api_url or "https://api.anthropic.com"
+    openai_url = config.openai_api_url or "https://api.openai.com"
     backend_section = ""
 
     if config.backend == "anyllm" or config.backend.startswith("anyllm-"):
         # any-llm backend
-        backend_status = f"{effective_anyllm_provider.title()} via any-llm"
         backend_section = """
   Set credentials for your provider (e.g., OPENAI_API_KEY, MISTRAL_API_KEY)
   Providers: https://mozilla-ai.github.io/any-llm/providers/
@@ -268,14 +267,6 @@ def proxy(
         provider = config.backend.replace("litellm-", "")
         provider_config = get_provider_config(provider)
 
-        # Build backend status
-        if provider_config.uses_region:
-            backend_status = (
-                f"{provider_config.display_name} via LiteLLM (region={effective_region})"
-            )
-        else:
-            backend_status = f"{provider_config.display_name} via LiteLLM"
-
         # Build usage instructions from provider config
         env_vars_str = (
             ", ".join(provider_config.env_vars) if provider_config.env_vars else "See docs"
@@ -314,7 +305,6 @@ Memory (Multi-Provider):
 Starting proxy server...
 
   URL:          http://{config.host}:{config.port}
-  Backend:      {backend_status}
   Mode:         {config.mode}
   Optimization: {"ENABLED" if config.optimize else "DISABLED"}
   Caching:      {"ENABLED" if config.cache_enabled else "DISABLED"}
@@ -322,18 +312,19 @@ Starting proxy server...
   Memory:       {memory_status}
   License:      {license_status}
 {backend_section}
-Usage with Claude Code:
-  ANTHROPIC_BASE_URL=http://{config.host}:{config.port} claude
+Routing:
+  /v1/messages         → {anthropic_url}
+  /v1/chat/completions → {openai_url}
+  /v1/responses        → {openai_url}  (HTTP + WebSocket)
 
-Usage with OpenAI-compatible clients:
-  OPENAI_BASE_URL=http://{config.host}:{config.port}/v1 your-app
+Usage:
+  Claude Code:   ANTHROPIC_BASE_URL=http://{config.host}:{config.port} claude
+  Codex / OpenAI: OPENAI_BASE_URL=http://{config.host}:{config.port}/v1 your-app
 {memory_section}
 Endpoints:
   GET  /health     Health check
   GET  /stats      Detailed statistics
   GET  /metrics    Prometheus metrics
-  POST /v1/messages           Anthropic API
-  POST /v1/chat/completions   OpenAI API
 
 Press Ctrl+C to stop.
 """)

headroom/cli/wrap.py

@@ -54,7 +54,14 @@ def _get_log_path() -> Path:
     return log_dir / "proxy.log"
 
 
-def _start_proxy(port: int, *, learn: bool = False) -> subprocess.Popen:
+def _start_proxy(
+    port: int,
+    *,
+    learn: bool = False,
+    backend: str | None = None,
+    anyllm_provider: str | None = None,
+    region: str | None = None,
+) -> subprocess.Popen:
     """Start Headroom proxy as a background subprocess.
 
     Logs are written to ~/.headroom/logs/proxy.log to avoid pipe buffer
@@ -72,6 +79,19 @@ def _start_proxy(port: int, *, learn: bool = False) -> subprocess.Popen:
     if learn:
         cmd.append("--learn")
 
+    # Forward backend configuration to proxy subprocess
+    _backend = backend or os.environ.get("HEADROOM_BACKEND")
+    if _backend:
+        cmd.extend(["--backend", _backend])
+
+    _anyllm = anyllm_provider or os.environ.get("HEADROOM_ANYLLM_PROVIDER")
+    if _anyllm:
+        cmd.extend(["--anyllm-provider", _anyllm])
+
+    _region = region or os.environ.get("HEADROOM_REGION")
+    if _region:
+        cmd.extend(["--region", _region])
+
     log_path = _get_log_path()
     log_file = open(log_path, "a")  # noqa: SIM115
 
@@ -233,7 +253,15 @@ def _inject_rtk_instructions(file_path: Path, verbose: bool = False) -> bool:
     return True
 
 
-def _ensure_proxy(port: int, no_proxy: bool, *, learn: bool = False) -> subprocess.Popen | None:
+def _ensure_proxy(
+    port: int,
+    no_proxy: bool,
+    *,
+    learn: bool = False,
+    backend: str | None = None,
+    anyllm_provider: str | None = None,
+    region: str | None = None,
+) -> subprocess.Popen | None:
     """Start or verify proxy. Returns process handle if we started it."""
     if not no_proxy:
         if _check_proxy(port):
@@ -242,7 +270,13 @@ def _ensure_proxy(port: int, no_proxy: bool, *, learn: bool = False) -> subproce
         else:
             click.echo(f"  Starting Headroom proxy on port {port}...")
             try:
-                proc = _start_proxy(port, learn=learn)
+                proc = _start_proxy(
+                    port,
+                    learn=learn,
+                    backend=backend,
+                    anyllm_provider=anyllm_provider,
+                    region=region,
+                )
                 click.echo(f"  Proxy ready on http://127.0.0.1:{port}")
                 return proc
             except RuntimeError as e:
@@ -303,6 +337,9 @@ def _launch_tool(
     env_vars_display: list[str],
     *,
     learn: bool = False,
+    backend: str | None = None,
+    anyllm_provider: str | None = None,
+    region: str | None = None,
 ) -> None:
     """Common logic: start proxy, launch tool, clean up."""
     proxy_holder: list[subprocess.Popen | None] = [None]
@@ -318,7 +355,14 @@ def _launch_tool(
         click.echo("  ╚═══════════════════════════════════════════════╝")
         click.echo()
 
-        proxy_holder[0] = _ensure_proxy(port, no_proxy, learn=learn)
+        proxy_holder[0] = _ensure_proxy(
+            port,
+            no_proxy,
+            learn=learn,
+            backend=backend,
+            anyllm_provider=anyllm_provider,
+            region=region,
+        )
 
         click.echo()
         click.echo(f"  Launching {tool_label} (API routed through Headroom)...")
@@ -449,10 +493,31 @@ def claude(
 @click.option(
     "--learn", is_flag=True, help="Enable live traffic learning (patterns saved to AGENTS.md)"
 )
+@click.option(
+    "--backend",
+    default=None,
+    help="API backend for the proxy: 'anthropic', 'anyllm', 'litellm-vertex', etc. (env: HEADROOM_BACKEND)",
+)
+@click.option(
+    "--anyllm-provider",
+    default=None,
+    help="Provider for any-llm backend: openai, mistral, groq, etc. (env: HEADROOM_ANYLLM_PROVIDER)",
+)
+@click.option(
+    "--region", default=None, help="Cloud region for Bedrock/Vertex (env: HEADROOM_REGION)"
+)
 @click.option("--verbose", "-v", is_flag=True, help="Verbose output")
 @click.argument("codex_args", nargs=-1, type=click.UNPROCESSED)
 def codex(
-    port: int, no_rtk: bool, no_proxy: bool, learn: bool, verbose: bool, codex_args: tuple
+    port: int,
+    no_rtk: bool,
+    no_proxy: bool,
+    learn: bool,
+    backend: str | None,
+    anyllm_provider: str | None,
+    region: str | None,
+    verbose: bool,
+    codex_args: tuple,
 ) -> None:
     """Launch OpenAI Codex CLI through Headroom proxy.
 
@@ -467,6 +532,7 @@ def codex(
         headroom wrap codex -- "fix the bug"        # Pass prompt to codex
         headroom wrap codex --no-rtk                # Skip rtk setup
         headroom wrap codex --port 9999             # Custom proxy port
+        headroom wrap codex --backend anyllm --anyllm-provider groq
     """
     codex_bin = shutil.which("codex")
     if not codex_bin:
@@ -499,6 +565,9 @@ def codex(
         tool_label="CODEX",
         env_vars_display=[f"OPENAI_BASE_URL=http://127.0.0.1:{port}/v1"],
         learn=learn,
+        backend=backend,
+        anyllm_provider=anyllm_provider,
+        region=region,
     )
 
 
@@ -512,10 +581,23 @@ def codex(
 @click.option("--no-rtk", is_flag=True, help="Skip rtk installation and conventions injection")
 @click.option("--no-proxy", is_flag=True, help="Skip proxy startup (use existing proxy)")
 @click.option("--learn", is_flag=True, help="Enable live traffic learning")
+@click.option(
+    "--backend", default=None, help="API backend: 'anthropic', 'anyllm', 'litellm-vertex', etc."
+)
+@click.option("--anyllm-provider", default=None, help="Provider for any-llm backend")
+@click.option("--region", default=None, help="Cloud region for Bedrock/Vertex")
 @click.option("--verbose", "-v", is_flag=True, help="Verbose output")
 @click.argument("aider_args", nargs=-1, type=click.UNPROCESSED)
 def aider(
-    port: int, no_rtk: bool, no_proxy: bool, learn: bool, verbose: bool, aider_args: tuple
+    port: int,
+    no_rtk: bool,
+    no_proxy: bool,
+    learn: bool,
+    backend: str | None,
+    anyllm_provider: str | None,
+    region: str | None,
+    verbose: bool,
+    aider_args: tuple,
 ) -> None:
     """Launch aider through Headroom proxy.
 
@@ -530,6 +612,7 @@ def aider(
         headroom wrap aider -- --model gpt-4o            # Use GPT-4o
         headroom wrap aider -- --model claude-sonnet-4   # Use Claude
         headroom wrap aider --no-rtk                     # Skip rtk setup
+        headroom wrap aider --backend litellm-vertex --region us-central1
     """
     aider_bin = shutil.which("aider")
     if not aider_bin:
@@ -562,6 +645,9 @@ def aider(
             f"ANTHROPIC_BASE_URL=http://127.0.0.1:{port}",
         ],
         learn=learn,
+        backend=backend,
+        anyllm_provider=anyllm_provider,
+        region=region,
     )
 
 

headroom/proxy/server.py

@@ -2653,8 +2653,11 @@ class HeadroomProxy:
                     usage = backend_response.body.get("usage", {})
                     output_tokens = usage.get("output_tokens", 0)
 
+                    _backend_name = (
+                        self.anthropic_backend.name if self.anthropic_backend else "anthropic"
+                    )
                     await self.metrics.record_request(
-                        provider="bedrock",
+                        provider=_backend_name,
                         model=model,
                         input_tokens=optimized_tokens,
                         output_tokens=output_tokens,
@@ -2674,7 +2677,7 @@ class HeadroomProxy:
                             RequestLog(
                                 request_id=request_id,
                                 timestamp=datetime.now().isoformat(),
-                                provider="bedrock",
+                                provider=_backend_name,
                                 model=model,
                                 input_tokens_original=original_tokens,
                                 input_tokens_optimized=optimized_tokens,
@@ -4869,8 +4872,11 @@ class HeadroomProxy:
                 total_latency = (time.time() - start_time) * 1000
                 output_tokens = stream_state["output_tokens"]
 
+                _backend_name = (
+                    self.anthropic_backend.name if self.anthropic_backend else "anthropic"
+                )
                 await self.metrics.record_request(
-                    provider="bedrock",
+                    provider=_backend_name,
                     model=model,
                     input_tokens=optimized_tokens,
                     output_tokens=output_tokens,
@@ -4891,7 +4897,7 @@ class HeadroomProxy:
                         RequestLog(
                             request_id=request_id,
                             timestamp=datetime.now().isoformat(),
-                            provider="bedrock",
+                            provider=_backend_name,
                             model=model,
                             input_tokens_original=original_tokens,
                             input_tokens_optimized=optimized_tokens,
@@ -6264,6 +6270,14 @@ class HeadroomProxy:
 
             body["input"] = messages_to_responses_items(opt_msgs, original_items, preserved_indices)
 
+        # /v1/responses is OpenAI-specific (Codex) — always routes direct.
+        # LiteLLM/AnyLLM backends use /v1/chat/completions or /v1/messages.
+        if self.anthropic_backend is not None:
+            logger.debug(
+                f"[{request_id}] /v1/responses always routes to OpenAI direct "
+                f"(backend '{self.anthropic_backend.name}' not used for Responses API)"
+            )
+
         url = f"{self.OPENAI_API_URL}/v1/responses"
 
         try:
@@ -6277,7 +6291,7 @@ class HeadroomProxy:
                     model,
                     request_id,
                     original_tokens,
-                    original_tokens,
+                    optimized_tokens,
                     tokens_saved,
                     transforms_applied,
                     tags,
@@ -6360,31 +6374,46 @@ class HeadroomProxy:
             )
             return
 
-        await websocket.accept()
         request_id = await self._next_request_id()
 
         # Forward client headers to upstream, adding required OpenAI-Beta header
         ws_headers = dict(websocket.headers)
 
+        # Extract subprotocol from client — this is an application-level negotiation
+        # that MUST be forwarded end-to-end (unlike sec-websocket-key which is per-connection).
+        # Codex and OpenAI negotiate a subprotocol; stripping it causes OpenAI to return 500.
+        client_subprotocols: list[str] = []
+        raw_protocol = ws_headers.get("sec-websocket-protocol", "")
+        if raw_protocol:
+            client_subprotocols = [p.strip() for p in raw_protocol.split(",") if p.strip()]
+
+        # Accept client connection with the requested subprotocol
+        if client_subprotocols:
+            await websocket.accept(subprotocol=client_subprotocols[0])
+        else:
+            await websocket.accept()
+
         # Build upstream WebSocket URL (http→ws, https→wss)
         base = self.OPENAI_API_URL
         ws_base = base.replace("https://", "wss://").replace("http://", "ws://")
         upstream_url = f"{ws_base}/v1/responses"
 
-        # Forward all relevant headers (auth, org, project, beta, etc.)
-        # Skip hop-by-hop headers that shouldn't be forwarded
+        # Forward all client headers except hop-by-hop / per-connection headers.
+        # These are WebSocket handshake mechanics that the `websockets` library
+        # generates fresh for the upstream connection — forwarding them would conflict.
+        # Everything else (auth, org, beta, user-agent, custom headers) is forwarded as-is.
         _skip_headers = frozenset(
             {
-                "host",
-                "connection",
-                "upgrade",
-                "sec-websocket-key",
-                "sec-websocket-version",
-                "sec-websocket-extensions",
-                "sec-websocket-accept",
-                "sec-websocket-protocol",
-                "content-length",
-                "transfer-encoding",
+                "host",  # must match upstream, not local proxy
+                "connection",  # hop-by-hop
+                "upgrade",  # hop-by-hop
+                "sec-websocket-key",  # per-connection cryptographic nonce
+                "sec-websocket-version",  # protocol version (websockets lib sets this)
+                "sec-websocket-extensions",  # per-connection negotiation
+                "sec-websocket-accept",  # server-side only
+                "sec-websocket-protocol",  # handled via subprotocols param below
+                "content-length",  # hop-by-hop
+                "transfer-encoding",  # hop-by-hop
             }
         )
         upstream_headers: dict[str, str] = {}
@@ -6392,10 +6421,30 @@ class HeadroomProxy:
             if k.lower() not in _skip_headers:
                 upstream_headers[k] = v
 
+        # Ensure Authorization header is present — fall back to OPENAI_API_KEY env var.
+        # Safety net for clients that don't forward auth headers via WebSocket upgrade.
+        _has_auth = "authorization" in {k.lower() for k in upstream_headers}
+        if not _has_auth:
+            api_key = os.environ.get("OPENAI_API_KEY")
+            if api_key:
+                upstream_headers["Authorization"] = f"Bearer {api_key}"
+                logger.debug(f"[{request_id}] WS: injected Authorization from OPENAI_API_KEY env")
+            else:
+                logger.warning(
+                    f"[{request_id}] WS: no Authorization header from client and "
+                    f"OPENAI_API_KEY not set — upstream will likely reject"
+                )
+
         # Ensure the required beta header is present — OpenAI returns 500 without it
         if "openai-beta" not in {k.lower() for k in upstream_headers}:
             upstream_headers["OpenAI-Beta"] = "responses-api=v1"
 
+        logger.debug(
+            f"[{request_id}] WS upstream headers: "
+            f"{[k for k in upstream_headers if k.lower() != 'authorization']}, "
+            f"subprotocols={client_subprotocols}"
+        )
+
         try:
             # Receive the first message from client (the response.create request)
             first_msg_raw = await websocket.receive_text()
@@ -6479,6 +6528,11 @@ class HeadroomProxy:
             async with websockets.connect(
                 upstream_url,
                 additional_headers=upstream_headers,
+                subprotocols=(
+                    [websockets.Subprotocol(p) for p in client_subprotocols]
+                    if client_subprotocols and hasattr(websockets, "Subprotocol")
+                    else client_subprotocols or None
+                ),
                 ssl=ssl_ctx if upstream_url.startswith("wss://") else None,
             ) as upstream:
                 # Send (potentially compressed) first message
@@ -6490,16 +6544,28 @@ class HeadroomProxy:
                         while True:
                             msg = await websocket.receive_text()
                             await upstream.send(msg)
-                    except Exception:
+                    except Exception as relay_err:
+                        if "WebSocketDisconnect" not in type(relay_err).__name__:
+                            logger.debug(
+                                f"[{request_id}] WS client→upstream relay ended: {relay_err}"
+                            )
                         with contextlib.suppress(Exception):
                             await upstream.close()
 
                 async def _upstream_to_client() -> None:
                     try:
                         async for msg in upstream:
-                            await websocket.send_text(msg if isinstance(msg, str) else msg.decode())
-                    except Exception:
-                        pass
+                            if isinstance(msg, str):
+                                await websocket.send_text(msg)
+                            elif isinstance(msg, bytes):
+                                await websocket.send_bytes(msg)
+                            else:
+                                await websocket.send_text(str(msg))
+                    except Exception as relay_err:
+                        if "WebSocketDisconnect" not in type(relay_err).__name__:
+                            logger.debug(
+                                f"[{request_id}] WS upstream→client relay ended: {relay_err}"
+                            )
                     finally:
                         with contextlib.suppress(Exception):
                             await websocket.close()
@@ -6524,7 +6590,19 @@ class HeadroomProxy:
 
         except Exception as e:
             if "WebSocketDisconnect" not in type(e).__name__:
-                logger.error(f"[{request_id}] WS proxy error: {e}")
+                # Extract response body from websockets InvalidStatus for better debugging
+                error_detail = str(e)
+                if hasattr(e, "response"):
+                    try:
+                        resp = e.response
+                        body_bytes = getattr(resp, "body", None) or b""
+                        if body_bytes:
+                            error_detail += (
+                                f" | body: {body_bytes[:500].decode('utf-8', errors='replace')}"
+                            )
+                    except Exception:
+                        pass
+                logger.error(f"[{request_id}] WS proxy error: {error_detail}")
             with contextlib.suppress(Exception):
                 await websocket.close(code=1011, reason=str(e)[:120])