akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions
2696 commits 410 branches 169 tags 142 MiB
Commit graph

6 commits

Author SHA1 Message Date
Matt Haitana
64cb46e24b
fix(proxy): pass through cross-region prefixed Bedrock model IDs directly (#2330) 
## Description

When a model ID with a cross-region prefix (`au.`, `us.`, `eu.`,
`apac.`, `global.`) is sent to the Bedrock backend, `map_model_id` was
normalising it (e.g. `au.anthropic.claude-opus-4-8` → `claude-opus-4-8`)
then re-looking it up in the discovery map. If an APPLICATION inference
profile wrapping the same foundation model existed in the account, it
would be returned — routing the request to a profile the caller is not
authorised to invoke, resulting in a 403 from Bedrock even though the
system-defined profile is reachable directly.

Cross-region prefixed IDs are already fully-qualified system-defined
profile IDs; they must pass through unchanged.

Closes #

## Type of Change

- [x] Bug fix (non-breaking change that fixes an issue)

## Changes Made

- `headroom/backends/litellm.py`: added early-exit in `map_model_id` —
model IDs starting with `au.`, `us.`, `eu.`, `apac.`, or `global.` are
returned as `bedrock/<model_id>` without any discovery lookup
- `tests/test_bedrock_region.py`: two new regression tests covering the
exact failure mode and all five prefix families

## Testing

- [x] Unit tests pass (`pytest`)
- [x] New tests added for new functionality
- [x] Manual testing performed

### Test Output

```text
tests/test_bedrock_region.py ........................................... [ 95%]
..                                                                       [100%]

45 passed in 4.45s
```

## Real Behavior Proof

- Environment: headroom 0.32.0-dev, `--backend bedrock`, `--region
ap-southeast-2`, `--bedrock-profile <BEDROCK_PROFILE>`, proxy on port
8788
- Exact command / steps: Output
  ```
# Before fix — old map_model_id logic with a contaminated discovery map:
  # Input:      au.anthropic.claude-opus-4-8
  # Normalized: claude-opus-4-8
  # Resolved:   bedrock/<application-inference-profile-arn>  <-- 403

  # After fix — cross-region prefix detected, passed through directly:
  curl -s -X POST http://localhost:8788/v1/messages \
    -H "Content-Type: application/json" \
    -H "x-api-key: sk-ant-dummy" \
    -H "anthropic-version: 2023-06-01" \
-d
'{"model":"au.anthropic.claude-opus-4-8","max_tokens":64,"messages":[{"role":"user","content":"Reply
with just: fix works"}]}'
  ```
- Observed result:
`{"type":"message","role":"assistant","content":[{"type":"text","text":"fix
works"}],"model":"au.anthropic.claude-opus-4-8","stop_reason":"end_turn",...}`
— HTTP 200, routed to `bedrock/au.anthropic.claude-opus-4-8`
(system-defined profile) rather than the APPLICATION profile ARN
- Not tested: `apac.` and `global.` prefixes against a live AWS account
(covered by unit tests only)

## 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] 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

## Additional Notes

The existing `_fetch_bedrock_inference_profiles` already filters to
`typeEquals="SYSTEM_DEFINED"` so APPLICATION profiles are not added to
the discovery map during normal startup. This fix closes the remaining
gap where a caller passes a cross-region prefixed ID directly —
previously that ID was normalised before lookup, which could
accidentally match a stale or externally-injected map entry.

---------

Co-authored-by: JerrettDavis <mxjerrett@gmail.com>
Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: JD Davis <jd@jds-macbook-air.tail2a279.ts.net>
 2026-08-11 23:54:03 -05:00
Kenneth Wong
33c7f6cd3a
fix(bedrock): resolve global.* inference profiles + pin per-user app-profile ARNs (#1795) 
## What & why

Bedrock model resolution failed on accounts whose inference profiles use
the newer `global.` cross-region prefix and undated version suffixes. On
such an account, `list_inference_profiles` returns current-gen models as
`global.anthropic.claude-opus-4-8`, `global.anthropic.claude-sonnet-5`,
`global.anthropic.claude-opus-4-6-v1`,
`global.anthropic.claude-fable-5`, etc.

`_normalize_bedrock_profile_id` only stripped `us.`/`eu.`/`apac.`/`au.`
and only matched a dated `-vN:M` suffix, so every `global.`-prefixed
profile was silently dropped from the discovered model map. Requests
then fell through to the fabricated fallback id and Bedrock rejected
them:

```
litellm.BadRequestError: BedrockException - {"message":"The provided model identifier is invalid."}
```

On the affected account the discovered-profile count went from 5 → 14
after the fix.

## Changes

1. **`_normalize_bedrock_profile_id`** — strip the `global.` prefix in
addition to the region prefixes, and match undated version suffixes
(`-v1`, or none at all) alongside the legacy dated `-vN:M`.

2. **`HEADROOM_BEDROCK_MODEL_MAP` operator override** (read from the
process environment). AWS discovery keys the model map by normalized
model name, so it cannot disambiguate application inference profiles
that share one underlying model — e.g. a team where
`claude-sonnet-5-alice` and `claude-sonnet-5-bob` both resolve to
`claude-sonnet-5`. When you need requests billed to a *specific*
application profile (per-user cost attribution), pin it explicitly:

   ```

HEADROOM_BEDROCK_MODEL_MAP="claude-sonnet-5=arn:aws:bedrock:REGION:ACCT:application-inference-profile/abc123,claude-opus-4-8=arn:...:application-inference-profile/def456"
   ```

The plain model name (kept plain so a client's tool-search deferral
stays on) resolves to the pinned ARN, routed via the converse endpoint.
The override wins over discovery; when unset, discovery-only behaviour
is unchanged.

## Tests

`tests/test_bedrock_region.py` (43 passing): `global.`-prefixed
normalization across dated / bare-`-v1` / no-suffix shapes; override-map
parsing (empty, single, multi, whitespace, malformed-skip); and
`map_model_id` override routing (pinned name → app-profile ARN via
converse, wins over discovery; unpinned name falls through).

No behavioural change for accounts already on system-defined
region-prefixed profiles.

---------

Co-authored-by: JerrettDavis <mxjerrett@gmail.com>
 2026-07-09 16:56:23 -05:00
Matt Haitana
7d87aa2f1c
fix(bedrock): route ARNs via converse, named AWS profiles, and au. re… (#1456) 
## Description

Fix three related gaps in Bedrock support that prevented headroom from
working with Claude Code when `CLAUDE_CODE_USE_BEDROCK=0` and
`ANTHROPIC_BASE_URL` is pointed at the proxy:

1. **ARN passthrough used the wrong LiteLLM route** — application
inference profile ARNs (e.g.
`arn:aws:bedrock:ap-southeast-2:<account>:application-inference-profile/<id>`)
were forwarded as `bedrock/<arn>`, which LiteLLM rejects with HTTP 400
"Try calling via converse route". Fixed to `bedrock/converse/<arn>`.

2. **Named AWS profile not forwarded to completion calls** —
`--bedrock-profile` was wired through the CLI → config →
`LiteLLMBackend.__init__` and used to fetch the model map at startup,
but never stored on `self`. All four `acompletion()` call sites
(`send_message`, `stream_message`, `send_openai_message`,
`stream_openai_message`) passed only `aws_region_name` — the
actual Bedrock calls used ambient credentials regardless of the flag.
Fixed by storing `self.profile_name` and passing `aws_profile_name=` to
every `acompletion()` call.

3. **`ap-southeast-2` used the wrong region prefix** — Australia should
use `au.` for cross-region inference profile IDs, not `apac.`. Added
`ap-southeast-2 → "au"` to `_BEDROCK_REGION_PREFIXES` and `"au."` to the
strip list in `_normalize_bedrock_profile_id`.

Closes #

## Type of Change

- [x] Bug fix (non-breaking change that fixes an issue)

## Changes Made

- `backends/litellm.py`: route `arn:aws:` model IDs via
`bedrock/converse/<arn>` in `map_model_id`
- `backends/litellm.py`: store `profile_name` as `self.profile_name` in
`LiteLLMBackend.__init__`; pass `aws_profile_name=` to `acompletion()`
in all four call sites; use
`boto3.Session(profile_name=...)` for startup discovery; cache key is
`region:profile_name` to prevent cross-profile collisions
- `backends/litellm.py`: add `ap-southeast-2 → "au"` to
`_BEDROCK_REGION_PREFIXES`; add `"au."` to prefix strip list in
`_normalize_bedrock_profile_id`
- `providers/registry.py`: pass `profile_name=bedrock_profile` to
`LiteLLMBackend`
- `proxy/server.py`: pass `config.bedrock_profile` to
`create_proxy_backend`
- `docs/claude-code-bedrock-headroom.md`: remove false claim that ARNs
in `ANTHROPIC_DEFAULT_*_MODEL` bypass the proxy; fix troubleshooting
table
- `tests/test_bedrock_region.py`: update `test_arn_passthrough` to
expect `bedrock/converse/<arn>`; update cache key format; add
`test_profile_cache_isolation`,
`test_ap_southeast_2_uses_au_prefix`, and
`TestBedrockProfileForwardedToCompletion` (3 async tests asserting
`aws_profile_name` appears in `acompletion()` kwargs for named profiles
and is
absent for the no-profile case)
- `tests/test_provider_registry*.py`,
`test_vertex_claude_compression.py`: update `litellm_backend_cls` stubs
to accept `profile_name=None`

## Testing

- [x] Unit tests pass (`pytest`)
- [x] New tests added for new functionality
- [x] Manual testing performed

### Test Output

```text
$ pytest tests/test_bedrock_region.py tests/test_provider_registry.py tests/test_provider_registry_extended.py \
    -k "not test_fallback_when_boto3_import_fails and not test_fallback_when_api_call_fails and not test_successful_fetch" -q
collected 51 items / 3 deselected / 48 selected

tests/test_bedrock_region.py ...........................
tests/test_provider_registry.py ...........
tests/test_provider_registry_extended.py .......

48 passed, 3 deselected in 2.00s
```

Note: 3 deselected tests use patch("builtins.__import__") which hangs
under Python 3.13 — pre-existing issue unrelated to these changes.

## Real Behavior Proof

- Environment: macOS, Python 3.13, Claude Code with
`CLAUDE_CODE_USE_BEDROCK=0`, `ANTHROPIC_BASE_URL=http://127.0.0.1:8787`,
AWS ap-southeast-2, application inference profile ARNs in
`ANTHROPIC_DEFAULT_*_MODEL`
- Exact command / steps: `headroom proxy --port 8787 --backend bedrock
--region ap-southeast-2 --bedrock-profile "my-sso-profile"`
- Observed result: Requests routed correctly to
`bedrock/converse/arn:aws:bedrock:ap-southeast-2:...:application-inference-profile/<id>`
as confirmed in LiteLLM logs
- Not tested: EU/APAC region ARN passthrough (logic is identical);
non-SSO credential flows

```text
15:29:44 - LiteLLM:INFO: utils.py:4090 - 
LiteLLM completion() model= converse/arn:aws:bedrock:ap-southeast-2:<account>:application-inference-profile/<id>; provider = bedrock
2026-06-26 15:29:44,322 - LiteLLM - INFO - 
LiteLLM completion() model= converse/arn:aws:bedrock:ap-southeast-2:<account>:application-inference-profile/<id>; provider = bedrock
15:31:09 - LiteLLM:INFO: utils.py:4090 - 
LiteLLM completion() model= converse/arn:aws:bedrock:ap-southeast-2:<account>:application-inference-profile/<id>; provider = bedrock
2026-06-26 15:31:09,928 - LiteLLM - INFO - 
LiteLLM completion() model= converse/arn:aws:bedrock:ap-southeast-2:<account>:application-inference-profile/<id>; provider = bedrock
15:34:26 - LiteLLM:INFO: utils.py:4090 - 
LiteLLM completion() model= converse/arn:aws:bedrock:ap-southeast-2:<account>:application-inference-profile/<id>; provider = bedrock
2026-06-26 15:34:26,811 - LiteLLM - INFO - 
LiteLLM completion() model= converse/arn:aws:bedrock:ap-southeast-2:<account>:application-inference-profile/<id>; provider = bedrock
```

## 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
- [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

## Additional Notes

The 3 skipped tests (`test_fallback_when_boto3_import_fails`,
`test_fallback_when_api_call_fails`, `test_successful_fetch`) pre-exist
in the repo and use `patch("builtins.__import__")` which hangs under
Python 3.13. Not affected by these changes.

---------

Co-authored-by: Matt Haitana <mhaitana@costar.com>
Co-authored-by: JerrettDavis <mxjerrett@gmail.com>
 2026-07-02 22:51:05 -05:00
chopratejas
d5ca50cd03 fix(tests): stop module-level dotenv loaders from polluting os.environ during pytest collection 
# The bug

Several test modules and two production modules loaded the project `.env`
at *import time*. During pytest collection (where every test module is
imported once), this populated `os.environ` with API keys from `.env`.

The skipif guards in `test_proxy_passthrough_integration.py` (and
others) evaluate at collection time:

    @pytest.mark.skipif(not os.environ.get("OPENAI_API_KEY"), reason="...")

If the polluter module was collected *before* the guard, the guard saw
the leaked key, decided not to skip, and the integration tests ran
live against a fake key and failed. In a fresh local-dev venv with
`.env` + full `[dev]` extras, this manifested as ~16 spurious test
failures plus a misleading test runtime of 6+ minutes (live HTTP).

# Why now

CI does not see this (no `.env`). It only manifests when:
1. `litellm` (and friends) are installed — they run `dotenv.load_dotenv()`
   on import, populating `os.environ` from `.env`.
2. A `.env` file with real API keys exists locally.

Until the venv was provisioned with the full `[dev]` extras during
recent test work, `pytest.importorskip("litellm")` and
`from headroom.pricing import litellm_pricing` both silently no-op'd
(via try/except ImportError → `LITELLM_AVAILABLE=False`), so the leak
never triggered. With litellm now installed, the latent bug surfaced.

# The fix — three patterns

1. **Production modules** (`headroom/pricing/litellm_pricing.py`,
   `headroom/backends/litellm.py`): wrap the eager `import litellm` with
   a snapshot/restore of `os.environ`. Any keys litellm's bundled
   `python-dotenv` adds during import are deleted immediately. The
   module is fully imported and cached in `sys.modules` so subsequent
   imports hit the cache without re-running the side effect.

2. **Test modules using `pytest.importorskip("litellm")`**
   (`test_backend_bugs.py`, `test_bedrock_region.py`,
   `test_cost_tracker_counterfactual.py`): replace with
   `tests._dotenv.importorskip_no_env_leak("litellm")`, which does the
   same snapshot/restore around `importlib.import_module`.

3. **Test modules that intentionally need `.env` values for skipif
   guards** (`test_compression_summary_*.py`, `test_query_echo.py`,
   `test_cost_tracker_counterfactual.py`, `test_memory_usage_integration.py`,
   `test_bundled_tools_savings.py`): replace module-level
   `os.environ.setdefault(...)` / `dotenv.load_dotenv()` with
   `tests._dotenv.load_env_overrides()` (returns a local dict — does
   NOT mutate `os.environ`) plus `autouse_apply_env(...)` (function-
   scoped fixture that applies via `monkeypatch.setenv`, auto-cleaned
   at teardown). The skipif still works because
   `ANTHROPIC_KEY = os.environ.get(...) or _env_overrides.get(...)`
   reads from the local dict as fallback.

# Helper module

New `tests/_dotenv.py` exposes:
- `load_env_overrides() -> dict[str, str]` — read `.env` into a dict.
- `autouse_apply_env(overrides) -> fixture` — function-scoped autouse
  fixture that applies via `monkeypatch.setenv`.
- `importorskip_no_env_leak(module) -> module` — drop-in
  `pytest.importorskip` substitute that quarantines env mutations.

# Results

Local full-suite (excluding live-LLM and live-feed tests):
- Before: 46 failed, 4830 passed, 387s
- After:   2 failed, 4672 passed, 134s

The remaining 2 failures are unrelated environment-dependent tests
(missing `PIL` / Docker daemon).
 2026-04-26 09:15:37 -07:00
chopratejas
d9cc4f3991 Fix ruff lint errors in test files 2026-03-24 15:54:12 -07:00
Yitong Li
b0ece04b9b fix(bedrock): add EU/AP region support with graceful fallback 
Three issues fixed:

1. _fetch_bedrock_inference_profiles crashed the proxy on startup when
   boto3 was missing or the AWS API call failed (wrong credentials,
   permissions, network). Now catches exceptions and falls back to a
   static model map.

2. map_model_id produced invalid Bedrock model IDs for unmapped models.
   Bare names like 'claude-sonnet-4-20250514' became
   'bedrock/claude-sonnet-4-20250514' which is not a valid Bedrock
   identifier. Now constructs region-prefixed IDs like
   'bedrock/eu.anthropic.claude-sonnet-4-20250514-v1:0'.

3. No static fallback existed (_BEDROCK_MODEL_MAP was empty). Added
   _build_bedrock_fallback_map() that generates region-aware model IDs
   for all GA Claude models (us./eu./apac. prefixes).

Closes #28

Tests: 27 new tests covering region prefix mapping, static fallback map,
graceful degradation, and model ID mapping for EU/AP/US regions.
 2026-03-24 00:13:00 +08:00