headroom/tests/test_proxy/test_openai_upstream_header.py
Kirill 6c48ac81f2
fix(proxy): honor x-headroom-base-url in dedicated OpenAI handlers (#1502)
## Description

The dedicated OpenAI handlers (`/v1/chat/completions`, `/v1/responses`)
ignore the `x-headroom-base-url` request header that the opencode/CLI
transports already send on every routed request
(`plugins/opencode/src/transport.ts`) and that the generic passthrough
route already honors (`providers/proxy_routes.py:953`).

As a result, OpenAI-compatible gateways (LiteLLM, CPA, self-hosted vLLM,
Azure OpenAI) route correctly for passthrough traffic, but the dedicated
chat/responses handlers fall back to the default `OPENAI_API_URL` and
send the request — and the user's provider key — to the wrong upstream.
This forces OpenCode users behind a custom gateway to run a hand-rolled
plugin that re-spawns the proxy with `OPENAI_TARGET_API_URL` instead of
the supported `HeadroomPlugin`.

Refs #1503 (feature-request issue with full spec — API surface, failure
modes, security considerations).

## Type of Change

- [x] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update
- [ ] Performance improvement
- [ ] Code refactoring (no functional changes)

Non-breaking: when the header is absent (the common case), behavior is
identical to before — `_resolve_openai_upstream` falls back to
`self.OPENAI_API_URL`.

## Changes Made

- Added `OpenAIHandlerMixin._resolve_openai_upstream(request)` — returns
`request.headers.get("x-headroom-base-url") or self.OPENAI_API_URL`.
Prefers the header, falls back to the configured URL.
- Used it at the two direct-path HTTP upstream sites:
- `handle_openai_chat` →
`build_copilot_upstream_url(self._resolve_openai_upstream(request),
"/v1/chat/completions")`
- `handle_openai_responses` →
`build_copilot_upstream_url(self._resolve_openai_upstream(request),
"/v1/responses")`
- This makes the dedicated handlers behave identically to the catch-all
passthrough and the Azure path (`_select_passthrough_base_url`,
`providers/proxy_routes.py:66,:953`), which already read the same
header.
- The header is already stripped before forwarding by
`helpers._strip_internal_headers`, so no upstream leakage /
fingerprinting is introduced.
- CHANGELOG entry under `### Bug Fixes`.

## Testing

<!-- Check what you actually ran, then paste the real command output
below. -->

- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [ ] Type checking passes (`mypy headroom`) — not run locally (maturin
native build not available in my env; covered by CI)
- [x] New tests added for new functionality
- [x] Manual testing performed

### Test Output

New `tests/test_proxy/test_openai_upstream_header.py` pins the
resolution contract (3 cases):

```text
$ pytest tests/test_proxy/test_openai_upstream_header.py -q
...
collected 3 items
tests/test_proxy/test_openai_upstream_header.py ...                      [100%]
========================= 3 passed, 1 warning in 0.25s =========================
```

Fail-before confirmed (unpatched handler raises `AttributeError:
_resolve_openai_upstream`):

```text
FAILED tests/test_proxy/test_openai_upstream_header.py::test_header_overrides_configured_url
FAILED tests/test_proxy/test_openai_upstream_header.py::test_missing_header_falls_back_to_configured_url
FAILED tests/test_proxy/test_openai_upstream_header.py::test_empty_header_falls_back_to_configured_url
========================= 3 failed, 1 warning in 0.29s =========================
```

Lint/format:

```text
$ ruff check headroom/proxy/handlers/openai.py tests/test_proxy/test_openai_upstream_header.py
Ruff: No issues found
$ ruff format --check headroom/proxy/handlers/openai.py tests/test_proxy/test_openai_upstream_header.py
2 files already formatted
```

## Real Behavior Proof

- Environment: macOS (darwin), Python 3.12 (pipx install of
`headroom-ai`), Headroom proxy `headroom proxy --port 8787` with
`OPENAI_TARGET_API_URL=https://cpa.funxyz.fun` (an OpenAI-compatible
gateway — "CLI Proxy API"). OpenCode with a custom `cpa` provider
(`@ai-sdk/openai-compatible`, `baseURL: https://cpa.funxyz.fun/v1`)
using the official `HeadroomPlugin`.
- Exact command / steps: traced the bug in the installed package source
— confirmed `handle_openai_chat` builds its upstream URL from
`self.OPENAI_API_URL` only (`proxy/handlers/openai.py:2487`), never
reading `x-headroom-base-url`, while `providers/proxy_routes.py:953`
reads it for passthrough. Then applied this patch and re-imported the
handler from the repo source via `PYTHONPATH`.
- Observed result: before the patch, `/v1/chat/completions` requests
ignored the `x-headroom-base-url: https://cpa.funxyz.fun` header (set by
the opencode transport) and routed to the default upstream, failing
against a non-OpenAI gateway — requiring a custom respawn-plugin
workaround. After the patch, `_resolve_openai_upstream` returns the
header value and the request forwards to the configured gateway; the
official `HeadroomPlugin` works without the env-var workaround. Unit
tests pass (3/3) and fail on the unpatched handler (3/3).
- Not tested: full `uv sync` CI matrix (native `headroom._core` maturin
build unavailable locally, so `headroom.proxy.server` import chain that
pulls `transforms/content_router` can't be exercised here — the edited
handler module imports fine and the focused unit tests exercise the new
method directly). WebSocket/Codex paths (`handle_openai_responses_ws`,
`_ws_http_fallback`) — intentionally out of scope (see Additional
Notes). `mypy headroom` — deferred to CI.

## Review Readiness

- [x] I have performed a self-review
- [x] This PR is ready for human review

## Checklist

- [x] My code follows the project's style guidelines
- [x] I have performed a self-review of my code
- [x] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation — no public
API/docs surface; the header is already documented as an internal
control flag in `helpers.py:1489-1495`
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] New and existing unit tests pass locally with my changes
- [x] I have updated the CHANGELOG.md if applicable

## Additional Notes

**Scope boundary — WebSocket paths intentionally unchanged.** The two WS
sites (`handle_openai_responses_ws`, `_ws_http_fallback`) are
Codex-specific and left as-is:

1. They short-circuit to `chatgpt.com` under ChatGPT-session auth (not
arbitrary gateways).
2. The WS path strips `x-headroom-base-url` from `upstream_headers`
(`_strip_internal`, ~line 3756) before the upstream URL is built, and
`_ws_http_fallback` receives already-stripped headers as a parameter.
Honoring the header there would require threading it through the WS
internals and changing a signature, for a path a custom
OpenAI-compatible WebSocket gateway is unlikely to use. The HTTP paths
cover the realistic gateway case. Happy to do it as a follow-up if
maintainers want it.

**Issue-first.** This is a behaviour change, so per CONTRIBUTING a
feature-request issue (#1503) is open for triage with the full spec (API
surface, user stories, failure modes, security). This PR implements it;
holding for maintainer 👍 before treating as ready to merge.

---------

Co-authored-by: ShutovKS <shutovks@example.local>
Co-authored-by: JerrettDavis <mxjerrett@gmail.com>
2026-07-01 22:24:39 -05:00

80 lines
2.9 KiB
Python

"""Tests for ``OpenAIHandlerMixin._resolve_openai_upstream``.
The dedicated OpenAI handlers (``/v1/chat/completions``,
``/v1/responses``) must honor the ``x-headroom-base-url`` request header
so OpenAI-compatible gateways (LiteLLM, CPA, self-hosted vLLM, Azure
OpenAI) route correctly — consistent with the generic passthrough route
that already honors it (see ``providers/proxy_routes.py``).
These tests pin the resolution contract:
- header present → its value wins
- header absent → configured ``OPENAI_API_URL`` fallback
- header empty or whitespace-only → fallback (no blanking)
"""
from __future__ import annotations
import pytest
fastapi = pytest.importorskip("fastapi")
httpx = pytest.importorskip("httpx")
from starlette.datastructures import Headers # noqa: E402
from headroom.proxy.handlers.openai import OpenAIHandlerMixin # noqa: E402
class _FakeRequest:
"""Minimal stand-in exposing ``headers`` like a real Starlette request.
Uses ``starlette.datastructures.Headers`` so header lookup is
case-insensitive, matching the production ``request.headers`` — a
plain ``dict`` would let case-folding regressions pass silently.
"""
def __init__(self, headers: dict[str, str]) -> None:
self.headers = Headers(headers=headers)
def _stub_proxy(fallback_url: str) -> OpenAIHandlerMixin:
"""A bare mixin instance with only ``OPENAI_API_URL`` configured."""
return type( # type: ignore[return-value]
"_S",
(OpenAIHandlerMixin,),
{"OPENAI_API_URL": fallback_url},
)()
def test_header_overrides_configured_url() -> None:
proxy = _stub_proxy("https://api.openai.test")
# The transport sends the upstream origin (no /v1 path).
request = _FakeRequest({"x-headroom-base-url": "https://gateway.example"})
assert proxy._resolve_openai_upstream(request) == "https://gateway.example"
def test_missing_header_falls_back_to_configured_url() -> None:
proxy = _stub_proxy("https://api.openai.test")
request = _FakeRequest({})
assert proxy._resolve_openai_upstream(request) == "https://api.openai.test"
def test_empty_header_falls_back_to_configured_url() -> None:
"""An explicitly empty or whitespace-only header must not blank the upstream."""
proxy = _stub_proxy("https://api.openai.test")
empty = _FakeRequest({"x-headroom-base-url": ""})
assert proxy._resolve_openai_upstream(empty) == "https://api.openai.test"
whitespace = _FakeRequest({"x-headroom-base-url": " "})
assert proxy._resolve_openai_upstream(whitespace) == "https://api.openai.test"
def test_header_lookup_is_case_insensitive() -> None:
"""Transports may send mixed-case header names; lookup must still resolve."""
proxy = _stub_proxy("https://api.openai.test")
# Real transports routinely send Title-Case header names.
request = _FakeRequest({"X-Headroom-Base-Url": "https://gateway.example"})
assert proxy._resolve_openai_upstream(request) == "https://gateway.example"