UN-3152 [FIX] Bedrock bearer token lost on LLM.complete() re-validation#1957
Merged
Conversation
`LLM.complete()` re-runs `adapter.validate({**self.kwargs, **kwargs})`
on every call (llm.py:286). The first pass translates `aws_bearer_token`
into LiteLLM's `api_key` kwarg and pops both the bearer field and
`auth_type`, so the second pass sees:
- `api_key` present (output of first pass)
- no `aws_bearer_token` (popped)
- no `auth_type` (popped) → resolver falls into legacy mode
Pydantic's `model_dump()` only emits declared fields, and `api_key` was
not declared on `AWSBedrockLLMParameters` — so the second pass silently
dropped the token. LiteLLM then fell through to SigV4 signing, found no
AWS credentials, and 401'd with `Unable to locate credentials` (the
exact symptom reported when testing a bearer-token Bedrock adapter from
the UI).
Fix:
- Declare `api_key: str | None = None` on both bedrock parameter classes
so it survives Pydantic's round-trip. Embedding doesn't re-validate
today, but added for symmetry to defend against future regressions.
- Update the resolver to pop `api_key` in `iam_role` and `access_keys`
modes (drop stale tokens from previous bearer-mode kwargs) and to
pop a blank `api_key` (Pydantic's `None` default) in legacy mode while
preserving a non-blank one (the round-trip case the LLM path needs).
Two regression tests cover the round-trip path: the second
`validate()` call must keep `api_key` intact for both LLM and embedding
parameter classes. 37 tests pass.
Contributor
WalkthroughThis PR refines Bedrock credential resolution to properly manage LiteLLM's bearer token parameter ( ChangesBedrock Bearer Token Credential Management
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Contributor
Test ResultsSummary
Runner Tests - Full Report
SDK1 Tests - Full Report
|
|
Contributor
|
| Filename | Overview |
|---|---|
| unstract/sdk1/src/unstract/sdk1/adapters/base1.py | Adds `api_key: str |
| unstract/sdk1/tests/test_bedrock_adapter.py | Adds two regression tests (LLM and embedding bearer-token round-trip) that directly exercise the fixed behaviour; both tests are well-structured and correctly simulate the second validate() call. |
Reviews (1): Last reviewed commit: "[FIX] Bedrock bearer token lost on `LLM...." | Re-trigger Greptile
athul-rs
approved these changes
May 12, 2026
jaseemjaskp
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Fix a bug introduced by #1952 (bearer-token Bedrock auth) where the bearer token is silently dropped on
LLM.complete()'s secondvalidate()call, causing the UI "Test Connection" action to fail withBedrockException Invalid Authentication - Unable to locate credentialsfor any Bedrock LLM adapter configured with Bedrock API Key (Bearer Token) auth.Why
LLM.complete()re-runsadapter.validate({**self.kwargs, **kwargs})on every call (unstract/sdk1/src/unstract/sdk1/llm.py:286). The first pass translatesaws_bearer_token→ LiteLLM'sapi_keykwarg and pops both the bearer field andauth_type, so the second pass sees:api_keypresent (output of the first pass)aws_bearer_token(popped)auth_type(popped) → resolver falls into legacy modePydantic's
model_dump()only emits declared fields, andapi_keywas not declared onAWSBedrockLLMParameters— so the second pass silently dropped the token. LiteLLM then fell through to SigV4 signing, found no AWS credentials, and 401'd withUnable to locate credentials(the exact UI-visible symptom).This was 100% reproducible in the UI: pick Bedrock API Key (Bearer Token) → paste a valid key → click Test Connection → fail.
How
api_key: str | None = Noneon bothAWSBedrockLLMParametersandAWSBedrockEmbeddingParametersso the field survives Pydantic'smodel_dump(). Embedding doesn't re-validate today, but the symmetric declaration defends against a future regression._resolve_bedrock_aws_credentialsto popapi_keyiniam_roleandaccess_keysmodes (drops Pydantic'sNonedefault, plus any staleapi_keyfrom a previous bearer-mode kwargs dict). In legacy mode, pop a blankapi_keywhile preserving a non-blank one — the round-trip case the LLM path needs.Can this PR break any existing features. If yes, please list possible items. If no, please explain why.
No. Validated against every non-bearer flow. Output kwargs are byte-identical to pre-fix:
access_keysiam_roleauth_type)bearer_tokenLLM.complete()re-validateAll 19 pre-existing access_keys / iam_role / legacy / stale-key tests still pass — they explicitly assert
"api_key" not in outfor non-bearer modes, so the resolver pop guarantees no regression.Database Migrations
None.
Env Config
None.
Relevant Docs
Related Issues or PRs
Dependencies Versions
None changed.
Notes on Testing
Run locally:
37 passed (35 pre-existing + 2 new round-trip regression tests):
test_llm_bearer_token_survives_revalidation— pins the LLM bug fix; secondvalidate()call must keepapi_keyintact.test_embedding_bearer_token_survives_revalidation— defensive parity for embedding.End-to-end smoke check (manual):
unstract/sdk1.access_keysandiam_roleadapters to confirm no regression.Screenshots
The UI failure screenshots were attached to the bug report — re-testing them with this branch should show Test Connection passing instead of the
Unable to locate credentialstoast.Checklist
I have read and understood the Contribution Guidelines.