aws / aws/bedrock-agentcore-sdk-python

BedrockAgentCoreContext.get_workload_access_token() returns None for plain FastAPI runtimes (3 silent gotchas)

Ouverte
#484 0 commentaires 1 réaction 0 personnes assignées Voir sur GitHub
Langage dominant
Python
Étoiles
761
Forks
147
Merge moyen
1 j 23 h
PR mergées (30 j)
7

Description

## Summary

Inside an AgentCore Runtime container, `BedrockAgentCoreContext.get_workload_access_token()` returns `None` whenever the agent is served by a plain `FastAPI`/`Starlette` ASGI app rather than `BedrockAgentCoreApp`. This breaks `@requires_api_key` and `@requires_access_token`, plus any imperative `IdentityClient.get_api_key()` callers.

While debugging, I found **three** sequential silent failures that all need fixing for the documented identity flow to actually work end-to-end. Filing as one issue because they layer — masking each other if you only fix the top one.

## Reproduction

1. Build an AgentCore Runtime image whose ASGI entrypoint is a plain `fastapi.FastAPI()` app exposing `POST /invocations` (rather than `BedrockAgentCoreApp` + `@app.entrypoint`).
2. The agent does work in a background thread spawned via `threading.Thread(target=run_pipeline, ...)` (common pattern for return-200-fast async work).
3. Inside the pipeline, call `BedrockAgentCoreContext.get_workload_access_token()`. Observe `None`.
4. Caller-side: orchestrator Lambda calls `InvokeAgentRuntimeCommand({ ..., runtimeUserId: 'cognito-sub' })` with all required IAM (`InvokeAgentRuntime` + `InvokeAgentRuntimeForUser`).
5. Runtime role has `bedrock-agentcore:GetResourceApiKey` + `GetWorkloadAccessToken*` + `bedrock-agentcore-identity!*` SLR exists.

## Three silent failures, in the order you hit them

### 1. Plain FastAPI bypasses `_build_request_context`

The SDK populates the ContextVar via `_build_request_context` in `bedrock_agentcore/runtime/app.py` (called from `BedrockAgentCoreApp._handle_invocation` lines 537+). Plain FastAPI/Starlette apps bypass this entirely — the middleware never installs, so `BedrockAgentCoreContext.set_workload_access_token(...)` is never called even though the platform delivered the token on the wire.

**Header diagnostic** (logged from inside our handler against a SigV4-authorized runtime in us-east-1, 2026-05-18) — the platform sends the token under TWO header spellings on the same request:

```
HEADER (token-shaped) workloadaccesstoken: AgV4diWP... # matches ACCESS_TOKEN_HEADER constant
HEADER (token-shaped) x-amzn-bedrock-agentcore-runtime-workload-accesstoken: AgV4diWP... # undocumented
```

The undocumented long-form should either be removed or aliased explicitly. Today, code that reads only `WorkloadAccessToken` per the SDK constant is correct; code that reads the long form is relying on undocumented behavior.

### 2. ContextVar is per-thread

After bridging the header in user middleware, `get_workload_access_token()` *still* returned None in our pipeline. Cause: Python `ContextVar` storage is per-thread, not propagated across `threading.Thread` boundaries. Setting the token in the request handler thread doesn't reach a pipeline thread spawned from the handler.

This isn't called out in the `BedrockAgentCoreContext` docstring or in the runtime-oauth.html documentation. Issue #219 alludes to it (`ContextVar` is "safe under ASGI") but doesn't surface the threading caveat that bites users running long pipelines outside the request task.

### 3. Runtime role needs `secretsmanager:GetSecretValue` on `bedrock-agentcore-identity!*`

After (1) and (2) were fixed, `IdentityClient.get_api_key()` actually fired and got:

```
AccessDeniedException: ... not authorized to perform: secretsmanager:GetSecretValue
```

AgentCore Identity stores api-key credentials in Secrets Manager under reserved prefix `bedrock-agentcore-identity!*`. The `GetResourceApiKey` API surfaces the underlying secret to the caller, and AWS verifies the *caller* role (the runtime execution role) has `GetSecretValue` on the actual secret resource — not the SLR. This is *not* documented in https://docs.aws.amazon.com/bedrock-agentcore/latest/devguide/runtime-oauth.html or in the `GetResourceApiKey` API reference.

Tightly-scoped fix:

```python
runtime_role.addToPrincipalPolicy(new iam.PolicyStatement({
actions: ['secretsmanager:GetSecretValue'],
resources: [
f'arn:aws:secretsmanager:{region}:{account}:secret:bedrock-agentcore-identity!*',
],
}))
```

This is an SDK-adjacent issue rather than a strict SDK bug — but the right place to surface it is alongside the other Identity-flow gotchas, and the SDK's `IdentityClient.get_api_key` could include the canonical IAM grant in its docstring.

## Workaround

```python
# 1. Read the workload token off the request and thread it explicitly into
# the background pipeline so the new thread can re-set the ContextVar.
def _extract_workload_access_token(request: Request) -> str:
return (
request.headers.get("WorkloadAccessToken")
or request.headers.get("x-amzn-bedrock-agentcore-runtime-workload-accesstoken")
or ""
)

@app.post("/invocations")
async def invoke(request: Request, body: ...):
token = _extract_workload_access_token(request)
# ... start background thread, pass token as kwarg ...
spawn(target=run_pipeline, kwargs={..., 'workload_access_token': token})
return {"accepted": True}

def run_pipeline(*, workload_access_token: str = "", **kwargs):
if workload_access_token:
from bedrock_agentcore.runtime.context import BedrockAgentCoreContext
BedrockAgentCoreContext.set_workload_access_token(workload_access_token)
# ... rest of pipeline runs in this thread; @requires_api_key works now ...
```

Plus IAM grant on the runtime execution role (above).

## Expected behavior

At least one of:

1. **An exported, opt-in middleware**: `bedrock_agentcore.runtime.middleware.BedrockAgentCoreContextMiddleware` users can `app.add_middleware(...)` to get parity with `BedrockAgentCoreApp` for the request-handler path. Combined with a clearly-documented per-thread propagation pattern in the docstring.

2. **OR a `BedrockAgentCoreContext` API that doesn't depend on ContextVar**: e.g., `BedrockAgentCoreContext.from_request(request)` that reads the headers directly and returns a context object the user can pass into `@requires_api_key(context=...)`. Sidesteps both the middleware gap and the threading gap.

3. **AND**: docstring updates on `BedrockAgentCoreContext`, `IdentityClient`, and the Identity getting-started doc that explicitly call out:
- "ContextVar is per-thread; if you spawn background threads, re-set the token on entry"
- "Caller IAM principal must have `secretsmanager:GetSecretValue` on `bedrock-agentcore-identity!*`"
- The undocumented long-form header — either remove or formally alias

## Environment

- `bedrock-agentcore` `1.9.1` (PyPI)
- AgentCore Runtime: ARM64 container, Python 3.13, FastAPI 0.136.1, uvicorn 0.47.0
- Region: `us-east-1`
- Inbound auth: SigV4 (default IAM, no `AuthorizerConfiguration` set)

## PR proposal

I'm willing to send a PR for the middleware + docstring updates. Two questions before I draft:

1. **Shape preference?** Opt-in `BedrockAgentCoreContextMiddleware` that users explicitly add (option 1 above) vs. helper API on `BedrockAgentCoreContext` (option 2)? The middleware is a smaller diff but only fixes the request-handler path; the helper is more invasive but addresses the threading gotcha cleanly.
2. **Should the threading caveat live in code (e.g., a `BedrockAgentCoreContext.copy_for_thread()` helper) or just docs?**

Happy to send whichever shape maintainers prefer.

Guide de contribution

Ouvrir le guide de contribution

Piste de recherche

Commencez par _build_request_context dans bedrock_agentcore/runtime/app.py, puis lisez BedrockAgentCoreContext et IdentityClient ainsi que la documentation runtime-oauth.html. Suivez la manière dont les requêtes FastAPI simples et les threads d’arrière-plan reçoivent le token, et examinez les exigences IAM documentées. Le travail est terminé lorsque les maintainers s’accordent sur la forme de l’API et que les indications concernant le contexte, le threading et Secrets Manager sont mises à jour en conséquence.

Rédigé par le modèle d'indexation à partir du texte de l'issue.

Évaluation

Stack technique
aws, fastapi, python
Domaine
authentication, backend, security
Type d'issue
Bug
Difficulté
5/5
Temps estimé
Plus d'une semaine
Activité
Calme
Clarté
Plutôt claire
Accessibilité débutants
42/100

Recevez les nouvelles issues par e-mail

Un résumé court des issues GitHub adaptées aux débutants.