fix(bedrock): forward userContext in Knowledge Base Retrieve requests

[devin-ai-integration]

TLDR

Problem this solves:

• Bedrock Knowledge Base search dropped the caller's userContext
• ACL-enabled data sources then answered with zero documents through the proxy
• Both the nested extra_body key and the OpenAI SDK's top-level shape were lost

How it solves it:

• Forwards userContext (or user_context) into the Bedrock Retrieve request body
• Reads extra_body first, then the top-level key, as retrievalConfiguration already does
• The value is forwarded as sent and Bedrock validates it, like retrievalConfiguration
• No logging added for it, since userId is personal data
• Regression tests cover both shapes plus the proxy-shaped call

User Flow

Before: a developer searching an ACL-enabled Bedrock Knowledge Base through the proxy gets no documents back, because the identity they attached never reaches Bedrock

1. Their app calls client.vector_stores.search(vector_store_id="T37J8R4WTM", query="What is LiteLLM?", max_num_results=2, extra_body={"userContext": {"userId": "alice@example.com"}}) with the OpenAI Python SDK pointed at the proxy
2. That sends POST https://litellm-domain/v1/vector_stores/T37J8R4WTM/search with {"query": "What is LiteLLM?", "max_num_results": 2, "userContext": {"userId": "alice@example.com"}}
3. The response is 200 with "object": "vector_store.search_results.page", but it was produced without alice's identity: on an ACL-enabled data source such as SharePoint data is empty, and on a store without ACLs every userId gets the same results
4. The same query sent straight to Bedrock with aws bedrock-agent-runtime retrieve ... --user-context 'userId=alice@example.com' returns the documents alice is allowed to see
5. Nobody holding a proxy key can read anything out of an ACL-enabled store, whatever userId they send

After: the same request carries alice's identity to Bedrock, so she gets the documents her ACL grants

1. Their app calls client.vector_stores.search(vector_store_id="T37J8R4WTM", query="What is LiteLLM?", max_num_results=2, extra_body={"userContext": {"userId": "alice@example.com"}}) with the OpenAI Python SDK pointed at the proxy
2. That sends POST https://litellm-domain/v1/vector_stores/T37J8R4WTM/search with {"query": "What is LiteLLM?", "max_num_results": 2, "userContext": {"userId": "alice@example.com"}}
3. The response is 200 with "object": "vector_store.search_results.page" and the documents alice is permitted to see, the same set the direct Bedrock call returns
4. A raw HTTP caller that nests it instead, {"query": ..., "extra_body": {"userContext": {"userId": "alice@example.com"}}}, gets the same result
5. A proxy key that can search the store can now read as any userId it sends, the same trust Bedrock gives a direct caller; the proxy does not derive the userId from the key's own user

Relevant issues

None on GitHub; reported by a customer through support

Affected release

Linear ticket

Resolves LIT-4415

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

• I have added meaningful tests
• The handful of test files covering my change pass locally, e.g. uv run pytest tests/test_litellm/<your_test_file>.py -v. Leave the suites (make test-unit-*, make test-unit) to CI: it finishes in ~15 minutes where a laptop takes an hour or more
• My PR passes all required CI/CD checks (e.g., lint, schema.d.ts sync check, etc.)
• My PR's scope is as isolated as possible; it only solves 1 specific problem
• I have received a Greptile Confidence Score of at least 4/5 before requesting a maintainer review (Greptile reviews automatically once the PR is opened; only comment @greptileai to re-request a review after pushing changes)

Delays in PR merge?

If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).

Screenshots / Proof of Fix

Both legs boot the proxy from a detached worktree at the named commit with 2 uvicorn workers and no database, against the real Bedrock Knowledge Base T37J8R4WTM in us-west-2. That store has no ACLs, so a valid identity returns the same page whether or not it reaches Bedrock. What separates the legs is Bedrock's own validation of the field: Bedrock rejects an empty userContext (and a non-string userId) with a 400, so a 400 through the proxy proves the field arrived and a 200 on the same request proves the proxy dropped it

Config (config.yaml, keys come from the environment):

model_list:
  - model_name: dummy
    litellm_params:
      model: openai/dummy
      api_key: not-used
vector_store_registry:
  - vector_store_name: "acl-knowledgebase"
    litellm_params:
      vector_store_id: "T37J8R4WTM"
      custom_llm_provider: "bedrock"
      aws_region_name: "us-west-2"
      aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID
      aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY
general_settings:
  master_key: os.environ/LITELLM_MASTER_KEY

Boot, same on both legs apart from the checkout (.env carries the two AWS keys and LITELLM_MASTER_KEY, no DATABASE_URL):

$ git -C <worktree> rev-parse HEAD
<hash below>
$ python litellm/proxy/proxy_cli.py --config config.yaml --port $PORT --num_workers 2 --detailed_debug
INFO:     Uvicorn running on http://0.0.0.0:$PORT
INFO:     Started parent process
INFO:     Started server process
INFO:     Started server process
prisma_client: None

What Bedrock itself says per shape, calling Retrieve directly with the same keys (the AWS CLI validates userContext client-side, so {} never leaves the laptop through it; the signed curl shows the server's verdict):

$ for uc in '{"userId":"alice@example.com"}' '{}' '{"userId":12345}'; do echo "userContext=$uc"; curl -s -w '\nHTTP %{http_code}\n' --aws-sigv4 'aws:amz:us-west-2:bedrock' --user "$AWS_ACCESS_KEY_ID:$AWS_SECRET_ACCESS_KEY" -H 'Content-Type: application/json' -X POST https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases/T37J8R4WTM/retrieve -d '{"retrievalQuery":{"text":"What is LiteLLM?"},"retrievalConfiguration":{"vectorSearchConfiguration":{"numberOfResults":2}},"userContext":'"$uc"'}'; done
userContext={"userId":"alice@example.com"}
{"retrievalResults":[{"content":{"text":"Try LiteLLM Enterprise ... ## What is LiteLLM? LiteLLM simplifies **model access** ...","type":"TEXT"},"location":{"type":"WEB","webLocation":{"url":"https://www.litellm.ai"}},"score":0.6238149},{"content":{"text":". Without LiteLLM this would be hours of work ...","type":"TEXT"},"location":{"type":"WEB","webLocation":{"url":"https://www.litellm.ai"}},"score":0.5865586}]}
HTTP 200
userContext={}
{"message":"1 validation error detected: Value at 'userContext.userId' failed to satisfy constraint: Member must not be null"}
HTTP 400
userContext={"userId":12345}
{"Message":"NUMBER_VALUE cannot be converted to String"}
HTTP 400

The page every 200 below returns is the same two chunks from https://www.litellm.ai with scores 0.6238149 and 0.5865586; its text fields are shortened with ... for length only. $KEY is the proxy master key. openai SDK 2.33.0

Before (4e99640)

OpenAI SDK, valid userContext

1. from openai import OpenAI
   client = OpenAI(base_url="http://localhost:39749/v1", api_key=KEY)
   r = client.vector_stores.with_raw_response.search(
       vector_store_id="T37J8R4WTM", query="What is LiteLLM?", max_num_results=2,
       extra_body={"userContext": {"userId": "alice@example.com"}})
   print(r.http_response.status_code, r.http_response.text)

2. 200 {"object":"vector_store.search_results.page","search_query":"What is LiteLLM?","data":[{"score":0.6238149,"content":[{"text":"Try LiteLLM Enterprise ...","type":"text"}],"file_id":"https://www.litellm.ai","filename":"www.litellm.ai","attributes":{"x-amz-bedrock-kb-source-uri":"https://www.litellm.ai","x-amz-bedrock-kb-chunk-id":"1%3A0%3AidYPg5YByRuP5PdK96co","x-amz-bedrock-kb-data-source-id":"CCEJIRXXFI"}},{"score":0.5865586,"content":[{"text":". Without LiteLLM this would be hours of work ...","type":"text"}],"file_id":"https://www.litellm.ai","filename":"www.litellm.ai","attributes":{"x-amz-bedrock-kb-source-uri":"https://www.litellm.ai","x-amz-bedrock-kb-chunk-id":"1%3A0%3AjdYPg5YByRuP5PdK96co","x-amz-bedrock-kb-data-source-id":"CCEJIRXXFI"}}]}
   

OpenAI SDK, empty userContext

1. Same call with extra_body={"userContext": {}}

2. 200 {"object":"vector_store.search_results.page","search_query":"What is LiteLLM?","data":[{"score":0.6238149,...},{"score":0.5865586,...}]}
   

3. Same call with extra_body={"userContext": {"userId": 12345}}

4. 200 {"object":"vector_store.search_results.page","search_query":"What is LiteLLM?","data":[{"score":0.6238149,...},{"score":0.5865586,...}]}
   

curl, userContext nested under extra_body, valid

1. curl -s -w '\nHTTP %{http_code}\n' http://localhost:39749/v1/vector_stores/T37J8R4WTM/search -H "Authorization: Bearer $KEY" -H 'Content-Type: application/json' -d '{"query":"What is LiteLLM?","max_num_results":2,"extra_body":{"userContext":{"userId":"alice@example.com"}}}'
   

2. {"object":"vector_store.search_results.page","search_query":"What is LiteLLM?","data":[{"score":0.6238149,...},{"score":0.5865586,...}]}
   HTTP 200
   

curl, userContext nested under extra_body, empty

1. curl -s -w '\nHTTP %{http_code}\n' http://localhost:39749/v1/vector_stores/T37J8R4WTM/search -H "Authorization: Bearer $KEY" -H 'Content-Type: application/json' -d '{"query":"What is LiteLLM?","max_num_results":2,"extra_body":{"userContext":{}}}'
   

2. {"object":"vector_store.search_results.page","search_query":"What is LiteLLM?","data":[{"score":0.6238149,...},{"score":0.5865586,...}]}
   HTTP 200
   

3. Same with "userContext":{"userId":12345}

4. {"object":"vector_store.search_results.page","search_query":"What is LiteLLM?","data":[{"score":0.6238149,...},{"score":0.5865586,...}]}
   HTTP 200
   

curl, top-level userContext, valid

1. curl -s -w '\nHTTP %{http_code}\n' http://localhost:39749/v1/vector_stores/T37J8R4WTM/search -H "Authorization: Bearer $KEY" -H 'Content-Type: application/json' -d '{"query":"What is LiteLLM?","max_num_results":2,"userContext":{"userId":"alice@example.com"}}'
   

2. {"object":"vector_store.search_results.page","search_query":"What is LiteLLM?","data":[{"score":0.6238149,...},{"score":0.5865586,...}]}
   HTTP 200
   

curl, top-level userContext, empty

1. curl -s -w '\nHTTP %{http_code}\n' http://localhost:39749/v1/vector_stores/T37J8R4WTM/search -H "Authorization: Bearer $KEY" -H 'Content-Type: application/json' -d '{"query":"What is LiteLLM?","max_num_results":2,"userContext":{}}'
   

2. {"object":"vector_store.search_results.page","search_query":"What is LiteLLM?","data":[{"score":0.6238149,...},{"score":0.5865586,...}]}
   HTTP 200
   

3. Same with "userContext":{"userId":12345}

4. {"object":"vector_store.search_results.page","search_query":"What is LiteLLM?","data":[{"score":0.6238149,...},{"score":0.5865586,...}]}
   HTTP 200
   

Every shape got the same 200 page, including the two Bedrock rejects outright, so nothing the caller put in userContext reached Bedrock

After (033aa8b)

The tip 884087f only annotates a test helper on top of this hash; no runtime file changed since, so this leg stands

OpenAI SDK, valid userContext

1. from openai import OpenAI
   client = OpenAI(base_url="http://localhost:42934/v1", api_key=KEY)
   r = client.vector_stores.with_raw_response.search(
       vector_store_id="T37J8R4WTM", query="What is LiteLLM?", max_num_results=2,
       extra_body={"userContext": {"userId": "alice@example.com"}})
   print(r.http_response.status_code, r.http_response.text)

2. 200 {"object":"vector_store.search_results.page","search_query":"What is LiteLLM?","data":[{"score":0.6238149,"content":[{"text":"Try LiteLLM Enterprise ...","type":"text"}],"file_id":"https://www.litellm.ai","filename":"www.litellm.ai","attributes":{"x-amz-bedrock-kb-source-uri":"https://www.litellm.ai","x-amz-bedrock-kb-chunk-id":"1%3A0%3AidYPg5YByRuP5PdK96co","x-amz-bedrock-kb-data-source-id":"CCEJIRXXFI"}},{"score":0.5865586,"content":[{"text":". Without LiteLLM this would be hours of work ...","type":"text"}],"file_id":"https://www.litellm.ai","filename":"www.litellm.ai","attributes":{"x-amz-bedrock-kb-source-uri":"https://www.litellm.ai","x-amz-bedrock-kb-chunk-id":"1%3A0%3AjdYPg5YByRuP5PdK96co","x-amz-bedrock-kb-data-source-id":"CCEJIRXXFI"}}]}
   

OpenAI SDK, empty userContext

1. Same call with extra_body={"userContext": {}}

2. 400 {"error":{"message":"litellm.BadRequestError: BedrockException - {\"message\":\"1 validation error detected: Value at 'userContext.userId' failed to satisfy constraint: Member must not be null\"}","type":"invalid_request_error","param":null,"code":"400"}}
   

3. Same call with extra_body={"userContext": {"userId": 12345}}

4. 400 {"error":{"message":"litellm.BadRequestError: BedrockException - {\"Message\":\"NUMBER_VALUE cannot be converted to String\"}","type":"invalid_request_error","param":null,"code":"400"}}
   

curl, userContext nested under extra_body, valid

1. curl -s -w '\nHTTP %{http_code}\n' http://localhost:42934/v1/vector_stores/T37J8R4WTM/search -H "Authorization: Bearer $KEY" -H 'Content-Type: application/json' -d '{"query":"What is LiteLLM?","max_num_results":2,"extra_body":{"userContext":{"userId":"alice@example.com"}}}'
   

2. {"object":"vector_store.search_results.page","search_query":"What is LiteLLM?","data":[{"score":0.6238149,...},{"score":0.5865586,...}]}
   HTTP 200
   

curl, userContext nested under extra_body, empty

1. curl -s -w '\nHTTP %{http_code}\n' http://localhost:42934/v1/vector_stores/T37J8R4WTM/search -H "Authorization: Bearer $KEY" -H 'Content-Type: application/json' -d '{"query":"What is LiteLLM?","max_num_results":2,"extra_body":{"userContext":{}}}'
   

2. {"error":{"message":"litellm.BadRequestError: BedrockException - {\"message\":\"1 validation error detected: Value at 'userContext.userId' failed to satisfy constraint: Member must not be null\"}","type":"invalid_request_error","param":null,"code":"400"}}
   HTTP 400
   

3. Same with "userContext":{"userId":12345}

4. {"error":{"message":"litellm.BadRequestError: BedrockException - {\"Message\":\"NUMBER_VALUE cannot be converted to String\"}","type":"invalid_request_error","param":null,"code":"400"}}
   HTTP 400
   

curl, top-level userContext, valid

1. curl -s -w '\nHTTP %{http_code}\n' http://localhost:42934/v1/vector_stores/T37J8R4WTM/search -H "Authorization: Bearer $KEY" -H 'Content-Type: application/json' -d '{"query":"What is LiteLLM?","max_num_results":2,"userContext":{"userId":"alice@example.com"}}'
   

2. {"object":"vector_store.search_results.page","search_query":"What is LiteLLM?","data":[{"score":0.6238149,...},{"score":0.5865586,...}]}
   HTTP 200
   

curl, top-level userContext, empty

1. curl -s -w '\nHTTP %{http_code}\n' http://localhost:42934/v1/vector_stores/T37J8R4WTM/search -H "Authorization: Bearer $KEY" -H 'Content-Type: application/json' -d '{"query":"What is LiteLLM?","max_num_results":2,"userContext":{}}'
   

2. {"error":{"message":"litellm.BadRequestError: BedrockException - {\"message\":\"1 validation error detected: Value at 'userContext.userId' failed to satisfy constraint: Member must not be null\"}","type":"invalid_request_error","param":null,"code":"400"}}
   HTTP 400
   

3. Same with "userContext":{"userId":12345}

4. {"error":{"message":"litellm.BadRequestError: BedrockException - {\"Message\":\"NUMBER_VALUE cannot be converted to String\"}","type":"invalid_request_error","param":null,"code":"400"}}
   HTTP 400
   

Every proxy status now equals Bedrock's own status for the same shape on all three client paths: the valid identity still returns the page, and the two shapes Bedrock rejects come back as its 400, which is only possible if userContext reached it

Observations from the legs:

• {"userId": ""} returns 200 on both legs, Bedrock accepts it
• The AWS CLI rejects {} before sending, Bedrock rejects it too
• No ACL-enabled store reachable, valid identity invisible in results

Type

🐛 Bug Fix

Caveats (if any)

Low

• A proxy key that can search the store can read as any userId it sends: the proxy forwards the identity as sent and does not derive it from the key's own user, the same trust Bedrock gives a direct caller and the same pass-through the retrievalConfiguration filters already get. Kept as is: mapping the key's user onto userContext would be a new opt-in feature with its own identity-format decisions (a SharePoint ACL wants the user principal name, not a proxy user id) and would break the shape the ticket describes, an app passing extra_body.userContext for end users who are not proxy users; the docs state the trust boundary
• userContext is forwarded as sent, no local validation, so a bad shape is Bedrock's 400 where it used to be silently ignored
  • Seen live at this tip: {}, a non-string userId, and a bare string all come back 400, while {"userId": ""} is accepted
  • Kept as is: Bedrock's validation is the source of truth and a local copy of it would drift, the same reason retrievalConfiguration is not validated locally
• A user_context set on a store in vector_store_registry is forwarded too, the same way its retrievalConfiguration already is
  • A caller's userContext overrides it, while a caller's user_context loses to it, because the endpoint lays the store's litellm_params over the request body before either spelling is read; kept as is since that merge predates this PR and applies to every key
  • On the chat completions file_search path the store-level value is the only way to set it, since the tool block has no per-request field; the ticket scopes the fix to the vector store search API
• With --detailed_debug the shared HTTP handler logs the outbound Retrieve body, userId included, as it does for every request field; this PR adds no logging of its own, and redacting one field there is a change to the shared handler outside this fix
• Top-level retrievalConfiguration on the proxy path is still dropped, pre-existing and out of scope
• The ACL-enabled leg is not exercised live: only a store without ACLs is reachable from CI, so the proof shows Bedrock validating the forwarded field rather than filtering results by it. Standing up a SharePoint or OneDrive backed store is a multi-hour data-source setup outside this fix, and the malformed-value legs already prove the field reaches Bedrock
• Neither key was in the docs; docs(bedrock): document userContext for Knowledge Base search litellm-docs#1500 documents both
• The CircleCI reds at the tip are fleet-wide or known flaky classes, and no fix on main since the merge base touches them: 22 Together AI tests fail with Unable to access model in 7 of the last 10 pipelines on other branches with no green run anywhere, and the sumologic NDJSON, Responses streaming iterator, router cooldown, and test_openai_endpoints::test_chat_completion tests fail in 1 or 2 of those 10 with 5 green jobs each; none touches a file this PR changes, and CircleCI is not a required check

Final Attestation

• The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR

• 033aa8b passes /live-pr-risk

[devin-ai-integration]

🤖 Devin AI Engineer

I'll be helping with this pull request! Here's what you should know:

✅ I will automatically:

• Address comments on this PR. Add '(aside)' to your comment to have me ignore it.
• Look at CI failures and help fix them

Note: I can only respond to comments from users who have write access to this repository.

⚙️ Control Options:

• Disable automatic comment, CI, and merge conflict monitoring

[codspeed]

Merging this PR will not alter performance

✅ 31 untouched benchmarks

---

_{Comparing litellm_bedrock_kb_user_context (884087f) with main (b6143b3)}

[greptile-apps]

Greptile Summary

This PR forwards caller-provided Bedrock Knowledge Base userContext data into Retrieve requests so ACL-aware searches receive the intended identity.

• Accepts both userContext and user_context.
• Gives nested extra_body values precedence over top-level LiteLLM parameters.
• Adds the corresponding Bedrock request types.
• Adds transformation-level and public search-path regression coverage.

Confidence Score: 5/5

The PR appears safe to merge, with no outstanding correctness, security, or repository-rule failures identified.

The forwarding logic handles both supported field spellings, preserves the intended precedence, omits absent context, and is covered through both transformation and public search paths. The previous typing finding was fully addressed and its thread is resolved.

Important Files Changed

Filename	Overview
litellm/llms/bedrock/vector_stores/transformation.py	Extracts either supported user-context spelling with documented precedence and forwards the value unchanged to Bedrock.
litellm/types/integrations/rag/bedrock_knowledgebase.py	Adds typed, read-only definitions for Bedrock Knowledge Base user context.
tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py	Covers forwarding from both input sources, precedence, omission, and the newly typed helper.
tests/test_litellm/vector_stores/test_main.py	Verifies that the public search path includes top-level user context in the serialized Bedrock request.

_{Reviews (2): Last reviewed commit: "test(bedrock): type the vector store sea..." | Re-trigger Greptile}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[mateo-berri]

@greptileai

[VANDRANKI]

Community review, not a merge-gate approval.

I traced _user_context() in litellm/llms/bedrock/vector_stores/transformation.py. It builds a tuple of (extra_body, litellm_params) filtered to actual Mapping instances (so a None extra_body is dropped, not an error), then does a lazy next() scan over sources -> ("userContext", "user_context") and returns the first key whose value is not None. That gives priority: extra_body["userContext"] > extra_body["user_context"] > litellm_params["userContext"] > litellm_params["user_context"], and it's a real short-circuit since next() stops at the first match rather than evaluating the whole generator.

I checked this against the three new tests in test_bedrock_vector_store_transformation.py:

• extra_body-only userContext forwards correctly, and retrievalConfiguration from the other optional param still comes through unaffected.
• litellm_params-only user_context (snake_case, no extra_body at all) forwards correctly, which confirms None extra_body doesn't blow up the isinstance(source, Mapping) filter.
• when both extra_body and litellm_params set a (conflicting) userContext, extra_body wins, matching the priority order I traced above.

The end-to-end test in test_main.py (test_search_forwards_top_level_user_context_to_bedrock_retrieve) goes one level up through vector_stores/main.py's search() and asserts the posted JSON body actually contains userContext, not just that the transformation function's return value has it, so this isn't just a unit-level check in isolation.

One thing I did not verify: whether userId is the only field Bedrock's real KB Retrieve API accepts under userContext, or whether BedrockKBUserContext (a TypedDict with just userId: ReadOnly[str]) is a deliberately narrow first cut. If AWS's Retrieve API supports other userContext fields, this type would need extending later, but that's a scope question, not a bug in what's here.

The rest of the diff (existing test now asserting "userContext" not in body when neither source sets it) is a reasonable regression guard for the negative case too.

This reads as a real, well-scoped, well-tested fix, and I didn't find a logic error tracing it.

[mateo-berri]

bugbot run

[cursor]

✅ Bugbot reviewed your changes and found no new issues!

Comment @cursor review or bugbot run to trigger another review on this PR

^{Reviewed by Cursor Bugbot for commit 884087f. Configure here.}

[mateo-berri]

LGTM