From 54cfa361d308dec567615c346af7c77d52ebb676 Mon Sep 17 00:00:00 2001 From: Manmit Singh Date: Thu, 2 Jul 2026 07:32:15 +0530 Subject: [PATCH] fix(bedrock): fail fast when session-token auth lacks botocore (#1553) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## 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 --- CHANGELOG.md | 7 +++ headroom/backends/litellm.py | 15 ++++++ .../test_bedrock_botocore_preflight.py | 49 +++++++++++++++++++ 3 files changed, 71 insertions(+) create mode 100644 tests/test_backends/test_bedrock_botocore_preflight.py diff --git a/CHANGELOG.md b/CHANGELOG.md index dfba370d7..6b0690211 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## Unreleased ### Fixed +- `--backend bedrock` now fails fast with an actionable error when temporary + AWS credentials (`AWS_SESSION_TOKEN`) are used but botocore is not installed + (e.g. the slim default Docker image). litellm's session-token auth path + imports botocore, so the missing dependency previously surfaced only at + request time as a misleading `authentication_error: No module named + 'botocore'`. The proxy now tells the user to install the `bedrock` extra up + front ([#1551](https://github.com/headroomlabs-ai/headroom/issues/1551)). - Content detection no longer crashes the proxy on text containing an orphaned `+++ ` target line with no preceding `--- ` source line (common in `set -x` xtrace output and partial diffs). The bundled `unidiff` 0.4.0 parser diff --git a/headroom/backends/litellm.py b/headroom/backends/litellm.py index ff7cb5e27..7a410ed79 100644 --- a/headroom/backends/litellm.py +++ b/headroom/backends/litellm.py @@ -12,8 +12,10 @@ LiteLLM handles all the auth and format translation internally. from __future__ import annotations +import importlib.util import json import logging +import os import uuid from collections.abc import AsyncIterator from dataclasses import dataclass, field @@ -423,6 +425,19 @@ class LiteLLMBackend(Backend): # For Bedrock, fetch model map dynamically from AWS API if provider == "bedrock": + # litellm takes the botocore-backed `_auth_with_aws_session_token` + # path as soon as temporary credentials (AWS_SESSION_TOKEN) are + # present. botocore is an optional dependency (the `bedrock` + # extra); when it is absent — as in the slim default Docker image — + # the failure only surfaces at request time as a misleading + # `authentication_error: No module named 'botocore'` (#1551). Fail + # fast at startup with an actionable message instead. + if os.environ.get("AWS_SESSION_TOKEN") and importlib.util.find_spec("botocore") is None: + raise ImportError( + "Bedrock with temporary credentials (AWS_SESSION_TOKEN) requires " + "botocore, which is not installed. Install the bedrock extra: " + "pip install 'headroom-ai[bedrock]' (or pip install botocore)." + ) self._model_map = _fetch_bedrock_inference_profiles(region) litellm.set_verbose = False # Reduce noise else: diff --git a/tests/test_backends/test_bedrock_botocore_preflight.py b/tests/test_backends/test_bedrock_botocore_preflight.py new file mode 100644 index 000000000..acfbb5c8c --- /dev/null +++ b/tests/test_backends/test_bedrock_botocore_preflight.py @@ -0,0 +1,49 @@ +"""Preflight guard for Bedrock + temporary credentials without botocore (#1551). + +litellm takes the botocore-backed ``_auth_with_aws_session_token`` path as +soon as ``AWS_SESSION_TOKEN`` is present. botocore ships only with the +``bedrock`` extra, so on a slim install the failure used to surface at request +time as a misleading ``authentication_error: No module named 'botocore'``. +``LiteLLMBackend`` now fails fast at startup with an actionable message. +""" + +from __future__ import annotations + +from unittest.mock import patch + +import pytest + +from headroom.backends import litellm as litellm_mod +from headroom.backends.litellm import LiteLLMBackend + + +def test_bedrock_session_token_without_botocore_raises_actionable( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm_mod, "LITELLM_AVAILABLE", True) + monkeypatch.setenv("AWS_SESSION_TOKEN", "tmp-session-token") + + # Simulate botocore not installed (the slim default-image case). + with patch("importlib.util.find_spec", return_value=None): + with pytest.raises(ImportError, match="botocore") as exc: + LiteLLMBackend(provider="bedrock", region="us-west-2") + + # The message must point at the fix, not just name the missing module. + assert "headroom-ai[bedrock]" in str(exc.value) + + +def test_bedrock_without_session_token_does_not_trip_botocore_guard( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Static-credential Bedrock users don't need botocore — guard stays quiet.""" + monkeypatch.setattr(litellm_mod, "LITELLM_AVAILABLE", True) + monkeypatch.delenv("AWS_SESSION_TOKEN", raising=False) + + # botocore missing, but no session token → guard must NOT raise. Stub the + # downstream model-map fetch and litellm handle so construction completes. + with patch("importlib.util.find_spec", return_value=None): + with patch.object(litellm_mod, "_fetch_bedrock_inference_profiles", return_value={}): + with patch.object(litellm_mod, "litellm", create=True): + backend = LiteLLMBackend(provider="bedrock", region="us-west-2") + + assert backend.provider == "bedrock"