mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
54cfa361d3
## Description With `--backend bedrock` and **temporary** AWS credentials (`AWS_SESSION_TOKEN`, as produced by SSO / STS assume-role / `credential_process`), every request fails. litellm self-signs Bedrock requests without botocore for *static* IAM keys, but as soon as a session token is present it takes the `_auth_with_aws_session_token` path in `litellm/llms/bedrock/base_aws_llm.py`, which imports `botocore`. botocore is an optional dependency — it ships only with headroom's `bedrock` extra, and the default Docker image is built with `HEADROOM_EXTRAS=proxy,code`, so botocore is absent. The failure surfaces only at request time as a misleading `authentication_error: No module named 'botocore'` (and as a bare `Invalid API key` in Claude Code). This PR makes the Bedrock backend **fail fast at startup** with an actionable message when a session token is set but botocore is missing — directly addressing the "clearer error message" the reporter asked for. It mirrors the existing optional-dependency guard pattern already used for boto3 in `backends/litellm.py`. Scope note: this does not change what the published image ships — whether to add botocore/`bedrock` to the default image extras is a separate sizing decision I left to maintainers. Static-credential Bedrock users (who never hit the botocore path) are unaffected. Refs #1551 ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) ## Changes Made - `headroom/backends/litellm.py`: when initializing the Bedrock backend with `AWS_SESSION_TOKEN` set and `botocore` not importable, raise an `ImportError` pointing at `pip install 'headroom-ai[bedrock]'` instead of letting the request fail later with a misleading auth error. - `tests/test_backends/test_bedrock_botocore_preflight.py`: regression tests — the guard raises an actionable error for the session-token-without-botocore case, and stays quiet for the static-credential case. - `CHANGELOG.md`: note under Unreleased → Fixed. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check`, `ruff format --check`) - [x] New tests added for new functionality - [x] Manual testing performed ### Test Output Regression test fails before the fix (no guard → no error raised), passes after: ```text # before fix (guard removed) FAILED tests/test_backends/test_bedrock_botocore_preflight.py::test_bedrock_session_token_without_botocore_raises_actionable # after fix tests/test_backends/test_bedrock_botocore_preflight.py .. [100%] 2 passed, 1 warning in 0.13s ``` `ruff check` / `ruff format --check` on the changed files: clean. ## Real Behavior Proof - Environment: macOS (arm64), Python venv, editable install (`pip install -e .`, no `bedrock` extra → botocore absent, matching the reported slim-image condition), `pytest`. - Exact command / steps: `python -m pytest tests/test_backends/test_bedrock_botocore_preflight.py`. (1) Removed the guard and ran the test → it failed because `LiteLLMBackend(provider="bedrock")` with `AWS_SESSION_TOKEN` set and botocore absent did NOT raise (reproducing the original "no early signal" behavior). (2) Applied the guard. (3) Re-ran → both tests pass, and the raised `ImportError` contains the `headroom-ai[bedrock]` install hint. - Observed result: with `AWS_SESSION_TOKEN` set and botocore not importable, the backend now raises a clear, actionable `ImportError` at construction time instead of deferring to litellm's later `No module named 'botocore'` auth error. Without a session token the guard does not fire, so static-credential users are unaffected. - Not tested: I did not run a live Bedrock request against AWS with real temporary credentials (no AWS account/STS access in this environment); the reporter already confirmed that installing botocore makes the identical request succeed, and this change surfaces that requirement at startup. ## 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 - [x] New and existing unit tests pass locally with my changes - [x] I have updated the CHANGELOG.md Co-authored-by: JD Davis <mxjerrett@gmail.com> |
||
|---|---|---|
| .. | ||
| __init__.py | ||
| test_bedrock_botocore_preflight.py | ||
| test_litellm_cache_stats.py | ||