Files
coder/aibridge
Yevhenii Shcherbina 8bf6f43016 feat: support cross-account Bedrock AssumeRole in AI Bridge (#26527)
# Support IAM role assumption for AWS Bedrock in AI Bridge

## Summary

Implements
https://linear.app/codercom/issue/AIGOV-371/support-dynamic-bedrock-assumerole-across-aws-accounts-for-ai-gateway

A Bedrock provider can now be configured with an IAM role to assume.
Before calling Bedrock, the gateway assumes that role via STS and signs
requests with the resulting temporary credentials. Whether the role
lives in the same account or another one is entirely a matter of the
role's trust policy.

## Problem

Many organizations prohibit long-lived AWS access keys and expect
workloads to authenticate through assumed IAM roles instead. A common
case is an organization that runs Bedrock across several AWS accounts,
one per business unit, and needs each unit's usage billed to its own
account by assuming a role there. AI Bridge previously authenticated a
Bedrock provider only with static keys or the gateway's own ambient AWS
identity, which is shared by every provider, with no way to assume a
role. These deployments had no clean path.

## How it works

When a provider is configured with a role ARN, the gateway uses its base
identity to assume that role via STS and signs Bedrock requests with the
temporary credentials it returns. The base identity is whatever the AWS
default credential chain resolves, IRSA, EKS Pod Identity, EC2 Instance
Profile, or static keys.

Credentials are resolved once when the provider is set up and are then
cached and rotated, so individual requests are served from the cache
rather than triggering a new STS call. A deployment that needs several
roles configures several providers, each pointing at its own role.

## Configuration

The role ARN is part of the Bedrock provider settings and is set through
the AI provider API. It is optional: a provider with no role ARN behaves
exactly as before.

## Scope and trade-offs

- This PR is backend only. The settings UI for the role ARN ships in a
follow-up.
- Configuration is not exposed through environment variables.
Environment-based provider configuration is being phased out in favor of
database-managed providers, so the role ARN is intentionally database
and API only.

Follow-up PR: https://github.com/coder/coder/pull/26578
2026-06-24 12:03:27 -04:00
..

aibridge

aibridge provides an HTTP handler that intercepts AI client requests bound for upstream AI providers (Anthropic, OpenAI, Copilot). It records token usage, prompts, and tool invocations per user. Optionally supports centralized MCP tool injection with allowlist/denylist filtering.

The handler is mounted by a host process. Today that host is coderd, which mounts the handler at /api/v2/aibridge/<provider>/*. Running aibridge as a separate process is planned for the future.

Architecture

┌─────────────────┐     ┌───────────────────────────────────────────┐
│    AI Client    │     │                    aibridge               │
│  (Claude Code,  │────▶│  ┌─────────────────┐    ┌─────────────┐   │
│   Cursor, etc.) │     │  │  RequestBridge  │───▶│  Providers  │   │
└─────────────────┘     │  │  (http.Handler) │    │  (Anthropic │   │
                        │  └─────────────────┘    │   OpenAI)   │   │
                        │                         └──────┬──────┘   │
                        │                                │          │
                        │                                ▼          │    ┌─────────────┐
                        │  ┌─────────────────┐    ┌─────────────┐   │    │  Upstream   │
                        │  │    Recorder     │◀───│ Interceptor │─── ───▶│    API      │
                        │  │ (tokens, tools, │    │ (streaming/ │   │    │ (Anthropic  │
                        │  │  prompts)       │    │  blocking)  │   │    │   OpenAI)   │
                        │  └────────┬────────┘    └──────┬──────┘   │    └─────────────┘
                        │           │                    │          │
                        │           ▼             ┌──────▼──────┐   │
                        │  ┌ ─ ─ ─ ─ ─ ─ ─ ┐      │  MCP Proxy  │   │
                        │  │    Database   │      │   (tools)   │   │
                        │  └ ─ ─ ─ ─ ─ ─ ─ ┘      └─────────────┘   │
                        └───────────────────────────────────────────┘

Components

  • RequestBridge: The main http.Handler that routes requests to providers
  • Provider: Defines bridged routes (intercepted) and passthrough routes (proxied)
  • Interceptor: Handles request/response processing and streaming
  • Recorder: Interface for capturing usage data (tokens, prompts, tools)
  • MCP Proxy (optional): Connects to MCP servers to list tool, inject them into requests, and invoke them in an inner agentic loop

Request Flow

  1. Client sends request to /anthropic/v1/messages or /openai/v1/chat/completions
  2. Actor extraction: Request must have an actor in context (via AsActor()). The host is responsible for authenticating the caller before invoking the handler.
  3. Upstream call: Request forwarded to the AI provider
  4. Response relay: Response streamed/sent to client
  5. Recording: Token usage, prompts, and tool invocations recorded

With MCP enabled: Tools from configured MCP servers are centrally defined and injected into requests (prefixed bmcp_). Allowlist/denylist regex patterns control which tools are available. When the model selects an injected tool, the gateway invokes it in an inner agentic loop, and continues the conversation loop until complete.

Passthrough routes (/v1/models, /v1/messages/count_tokens) are reverse-proxied directly.

Observability

Prometheus Metrics

Create metrics with NewMetrics(prometheus.Registerer):

Metric Type Description
interceptions_total Counter Intercepted request count
interceptions_inflight Gauge Currently processing requests
interceptions_duration_seconds Histogram Request duration
passthrough_total Counter Non-intercepted requests forwarded to the upstream
prompts_total Counter User prompt count
tokens_total Counter Token usage (input, output, cache read/write, provider extras)
injected_tool_invocations_total Counter Injected MCP tool invocations performed by the handler
non_injected_tool_selections_total Counter Client-defined tool selections returned by the model
circuit_breaker_state Gauge Circuit breaker state per provider/endpoint (0=closed, 0.5=half, 1=open)
circuit_breaker_trips_total Counter Times the circuit breaker transitioned to open
circuit_breaker_rejects_total Counter Requests rejected due to an open circuit breaker

Recorder Interface

Implement Recorder to persist usage data to your database:

  • aibridge_interceptions - request metadata (provider, model, initiator, timestamps)
  • aibridge_token_usages - input/output and cache read/write token counts per response
  • aibridge_user_prompts - user prompts
  • aibridge_tool_usages - tool invocations (injected and client-defined)
  • aibridge_model_thoughts - model reasoning content (thinking, reasoning summaries, commentary)
type Recorder interface {
    RecordInterception(ctx context.Context, req *InterceptionRecord) error
    RecordInterceptionEnded(ctx context.Context, req *InterceptionRecordEnded) error
    RecordTokenUsage(ctx context.Context, req *TokenUsageRecord) error
    RecordPromptUsage(ctx context.Context, req *PromptUsageRecord) error
    RecordToolUsage(ctx context.Context, req *ToolUsageRecord) error
    RecordModelThought(ctx context.Context, req *ModelThoughtRecord) error
}

Supported Routes

Each provider instance is mounted under /api/v2/aibridge/<name>, where <name> is the provider's configured name. For example, with an Anthropic provider named my-anthropic, its /messages endpoint would be reachable at /api/v2/aibridge/my-anthropic/v1/messages.

If a name is not set, the route path defaults to the provider's type: anthropic, openai, or copilot. The table below uses the default names.

(/*) denotes a route that handles both the exact path and any subpaths. A trailing /* denotes subpaths only.

Provider Route Type
Anthropic /anthropic/v1/messages Bridged (intercepted)
Anthropic /anthropic/v1/messages/count_tokens Passthrough
Anthropic /anthropic/v1/models(/*) Passthrough
Anthropic /anthropic/api/event_logging/* Passthrough
OpenAI /openai/v1/chat/completions Bridged (intercepted)
OpenAI /openai/v1/responses Bridged (intercepted)
OpenAI /openai/v1/responses/* Passthrough
OpenAI /openai/v1/conversations(/*) Passthrough
OpenAI /openai/v1/models(/*) Passthrough
Copilot /copilot/chat/completions Bridged (intercepted)
Copilot /copilot/responses Bridged (intercepted)
Copilot /copilot/models(/*) Passthrough
Copilot /copilot/agents/* Passthrough
Copilot /copilot/mcp/* Passthrough
Copilot /copilot/.well-known/* Passthrough