BAI-299 Implement native AWS Bedrock Converse transport - #107
Open
imranq2 wants to merge 1023 commits into
Open
Conversation
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
…s for improved consistency and maintainability
…ls-framework, language-model-common, and langgraph-prebuilt
…update uv.lock for dependency version upgrades
…ess ID placeholder
Switch to uv
…or improved configuration
…or marketplace plugins
… refactoring plugin loading logic
…, and types-boto3
…der for enhanced plugin management
…der for enhanced plugin management
…shots and refactor related components
…ore and update related components
…ved readability and maintainability
…del-common to versions 1.0.54 and 2.0.25 respectively
…matic fallback to in-memory storage
…pes and improve configuration handling
…n and dependency version
…ough usage, and upstream errors Adds start_time/end_time/duration_ms to usage records; upserts a per-session rollup (model-router-sessions) bucketed by tier with independent MODEL_ROUTING_USAGE_SESSION_TRACKING_ENABLED toggle; extends usage tracking to the api_type:anthropic passthrough route (previously uncovered) by buffering non-streaming responses and sniffing token usage out of streaming SSE bytes without altering them; and adds an ErrorTracker recording upstream failures to model-router-errors.
Fix usage-tracking IDOR (spoofable auth headers) and switch JSON logging to structlog
…r-errors
openai.APIError (raised by the SDK for SSE data events with no HTTP status
code, e.g. Bedrock Mantle sending {"error": {...}} mid-stream) was falling
through to the generic exception handler, which only recorded str(exc) - a
generic "server had an error" message with no diagnostic value. Add a
dedicated except branch that captures the SDK's .body/.code/.type/.param
in model-router-errors, and surface the same detail to the client instead
of just the generic passthrough message.
- Retry transient mid-stream Bedrock errors (ModelStreamErrorException, ThrottlingException, etc.) instead of skipping the retry loop entirely just because they carry no HTTP status code. - Record errors on the Anthropic-passthrough dispatch path, which previously re-raised connection/timeout failures with zero trace in model-router-errors. - Retry transport-level failures (httpx.TransportError) in _send_with_bedrock_retry with the same backoff+jitter as HTTP throttling, instead of failing outright on a momentary network blip. - Distinguish AWS credential failures (no profile configured, STS/SSO ClientError) from TokenRetrievalError instead of lumping them into a generic, unhelpful error_type with no actionable message. - Record mid-stream SDK failures (after streaming has already started) to model-router-errors via a new on_stream_error hook, instead of only ever showing them inline to the client.
…on, extract credential-error helper APIConnectionError/APITimeoutError/APIResponseValidationError are subclasses of openai.APIError but carry no code/type/param/body, so they were being misrecorded as bedrock_stream_error with a null-filled body instead of their real exception type. Fixed by adding an exact-type guard and extracting the shared classify/record/re-raise logic into _handle_unexpected_upstream_error, since a bare raise from the APIError except clause does not fall through to the sibling except Exception clause. Also moves _bedrock_credential_error_detail into aws_auth.py, alongside the BedrockCredentialsUnavailableError it already depends on, addressing the router.py file-size concern raised on the PR without a full-file refactor.
Documents the plan to let BEDROCK_TRANSPORT=native route haiku/sonnet traffic through Bedrock's Converse API instead of Bedrock Mantle, as a manual operator toggle during Mantle incidents.
Route the new toggle through LanguageModelGatewayEnvironmentVariables (model_routing_bedrock_transport) and thread it into CodingModelRouter via constructor injection from api.py, matching how every other model_routing_* setting already flows in — instead of a bare module-level os.environ.get() inside the model_routing package.
…errors ErrorTracker.record_error and CodingModelRouter._record_error now accept response_headers, populated from exc.response.headers for HTTP-status Bedrock Mantle errors and from the initial stream handshake response for mid-stream SSE-embedded errors (which carry no response of their own). Bedrock Mantle's error bodies are often too generic to act on alone (e.g. a bare "internal_server_error" with no further detail); the AWS request ID and other response headers give something to correlate against AWS-side logs.
…tream error wrapping
Rework the >=400 error-response handling in the Anthropic-passthrough
branch: parse the upstream error body once, after fallback-route
annotation, instead of twice with two different criteria. Previously
clean_error was computed from the pre-annotation body, so whenever the
annotated body failed the second JSON check the response silently fell
back to the stale, un-annotated message — dropping the fallback-route
note in exactly the case it exists to surface. Also tighten the
passthrough check to require a usable error shape (not just an "error"
key), so malformed bodies like {"error": null} get normalized instead
of forwarded verbatim, and stop leaking Python dict-repr syntax into
client-facing error text.
Adds test coverage for all three fixed scenarios, which previously had
none.
Task-by-task plan for the approved design spec: the model_routing_bedrock_transport toggle, bedrock_converse_client.py (request/response conversion, boto3 client cache, streaming adapter), and router.py dispatch wiring with error handling. Notes one deviation from the spec found while grounding the plan in the actual code: the Converse request is built from the already-OpenAI-shaped, already-budget-enforced body, not the raw Anthropic request, since context-budget enforcement runs upstream of the dispatch point for every api_type=="openai" route.
Adds model_routing_bedrock_transport (default "mantle") and threads it from api.py through CodingModelRouter's constructor as self._bedrock_transport. No dispatch behavior changes yet — the flag does nothing until later tasks add the native-Bedrock code path that reads it.
_get_bedrock_runtime_client caches one client per (AWS_PROFILE, region) pair rather than building one per request. _is_transient_bedrock_error_code reuses bedrock_client.py's existing Bedrock exception-name taxonomy for retry decisions on the native path.
_get_bedrock_runtime_client's check-then-set on _CLIENT_CACHE was unguarded — concurrent calls (this function runs via asyncio.to_thread from an async request path) racing on the same new (profile, region) key could each construct their own client, one silently overwriting the other. Double-checked locking with a threading.Lock ensures only one client is ever built per key.
The prior version of test_concurrent_calls_for_same_new_key_construct_only_one_client synchronized thread entry via threading.Barrier but mocked client construction as near-instant, so the GIL rarely switched between the two threads' check-then-set sequences — the test passed even against the pre-fix, unguarded code (verified: reverting the lock and re-running showed the old test wasn't a reliable regression guard). Adding a brief sleep inside the mocked construction call widens the critical section enough that the test now reliably fails without the lock and passes with it.
_openai_to_converse_request builds converse()/converse_stream() kwargs from the already-OpenAI-shaped, already-budget-enforced request body router.py produces today — not from the original Anthropic request, since context-budget enforcement runs upstream of the native-transport dispatch point and only understands the OpenAI shape. Covers text, tool_use/ tool_result round-trips, system prompts, and tool_choice mapping; image content blocks are dropped with a warning rather than supported.
… requests _openai_to_converse_request had no branch for tool_choice=="none" (a value message_translator.py's _anthropic_to_openai_request legitimately produces), so toolConfig.toolChoice was left unset while tools was still sent — an absent toolChoice behaves like "auto" in Bedrock's Converse API, silently inverting the caller's explicit "don't use tools" instruction. Converse has no direct "disable tools" toolChoice value, so the fix omits toolConfig entirely when tool_choice is "none" — if the model isn't told about any tools, it can't call one.
_converse_response_to_anthropic mirrors _openai_to_anthropic_response's contract for the native transport. stopReason values without a direct Anthropic equivalent (guardrail_intervened, content_filtered) default to end_turn rather than raising on an unrecognized value.
_iter_converse_stream_events pulls boto3's synchronous EventStream onto the event loop via asyncio.to_thread, one event per thread hop. _stream_bedrock_converse_to_anthropic translates Converse stream events into the same Anthropic SSE sequence _stream_oai_sdk_to_anthropic already emits for the Mantle path, with the same sinks/on_stream_error contract.
_converse_stream_with_usage_tracking records usage via the generic UsageTracker.record_usage (not the OpenAI-specific record_usage_from_openai_response, since Converse's usage keys differ from OpenAI's) after the stream completes, matching the skip-when-no- tokens behavior of its OAI counterpart.
Adds the auth=="aws" && bedrock_transport=="native" branch ahead of the existing openai-SDK Mantle path, and _dispatch_bedrock_native_nonstreaming to handle it: converse() via asyncio.to_thread, throttle-retry on transient ClientError codes, credential-error and non-transient-error recording via the existing _record_error/_error_response helpers with error_type="bedrock_native_error". Streaming requests still fall through to a not-yet-implemented method, completed in the next task.
_dispatch_bedrock_native_streaming mirrors the non-streaming dispatch's retry/error-recording logic, then hands off to _converse_stream_with_usage_tracking (or the bare translator when usage tracking is disabled). Mid-stream failures — after the SSE response has already started — are recorded the same way router.py's own _record_mid_stream_error does for the Mantle path, so they reach model-router-errors instead of only being shown inline to the client.
Sets MODEL_ROUTING_BEDROCK_TRANSPORT=native so haiku/sonnet traffic now routes through Bedrock's Converse API instead of Bedrock Mantle by default in this compose file. Flip back to "mantle" (or delete the line) once no longer needed.
…ng credential-error test Final whole-branch review found the native Converse streaming path could never show inline error text to the client on an early mid-stream failure (the Mantle path does, via a lazy message_start gate this path lacked, sending message_start eagerly instead) — the exact failure mode this fallback exists to serve produced a silently empty assistant message. Added the same inline "[bedrock-native-proxy error]" text injection Mantle already has, tracking whether any content block has ever opened (not just currently-open ones). Also added a streaming-dispatch credential/ClientError test (previously only the non-streaming path had one), corrected a docker-compose comment that inaccurately implied native Bedrock access might be provisioned separately from Mantle access (both share the same AWS credentials, so the "native" default stays), and updated the native-transport design spec's non-streaming data-flow section to reference the actual shipped _openai_to_converse_request function instead of the superseded _anthropic_to_converse_request name.
…verse transport MCP-prefixed tool names routinely exceed toolSpec.name's 64-character limit, failing the whole request with a ValidationException. Deterministically shorten names sent to Bedrock and restore the original name in tool_use blocks returned to the client.
Addresses aikido-pr-checks' review comments on PR #200: router.py and the new bedrock_converse_client.py had grown into files mixing several distinct responsibilities. Extracted _dispatch_bedrock_native_nonstreaming/ _streaming out of CodingModelRouter into a new BedrockNativeDispatcher class (bedrock_native_dispatcher.py), constructor-injected with the client provider, a usage-tracker getter, and error/latency-recording callbacks — mirrors the BaseChatCompletionsProvider adapter pattern already used for other backends. Split bedrock_converse_client.py into converse_request_translator.py (OpenAI<->Converse format conversion) and converse_stream_adapter.py (event-stream adaptation + usage-tracking wrapper); the remaining file owns only client lifecycle and native error-code classification. Replaced the module-level _CLIENT_CACHE dict/lock with an injectable BedrockRuntimeClientProvider instance, removing the hidden global state. Also fixes two usage-tracking gaps found while reviewing the native transport's usage handling: - Native Bedrock Converse responses (streaming and non-streaming) were dropping Bedrock's cacheReadInputTokens/cacheWriteInputTokens instead of mapping them onto Anthropic's cache_read_input_tokens/ cache_creation_input_tokens. Clients that derive context-window usage from these fields (e.g. Claude Code, see anthropics/claude-code#13385) can't distinguish a missing field from "no caching happened," so this silently broke context accounting for any client routed through the native transport. - UsageTracker.record_usage now accepts raw_usage and stores the upstream response's usage object verbatim across all three response paths (Anthropic passthrough, OpenAI/Mantle, native Bedrock Converse), so fields not yet normalized into a top-level column aren't silently dropped from model-router-usage.
…ve Bedrock Extends the anthropics/claude-code#13385 fix to every route, not just the native Bedrock Converse transport fixed previously: - OpenAI/Bedrock Mantle (non-streaming): _openai_to_anthropic_response now maps usage.prompt_tokens_details.cached_tokens (bedrock-mantle is a conformant OpenAI Chat Completions endpoint, per docs.aws.amazon.com/bedrock/latest/userguide/inference-chat-completions-mantle.html) onto cache_read_input_tokens. cache_creation_input_tokens has no OpenAI equivalent (that API doesn't charge separately for cache writes) but is still always included as 0 rather than omitted, for the same reason as the native-transport fix: a missing field looks identical to a broken one to a client computing context usage from it. - OpenAI/Bedrock Mantle (streaming): _stream_oai_sdk_to_anthropic now captures cached_tokens from each chunk's usage and surfaces it in the final message_delta event, plus stores the raw CompletionUsage dump for Mongo logging. Also fixes a pre-existing bug found while touching this code: message_delta's usage.output_tokens read a stale local variable that was never updated from usage_sink, so it always reported 0 instead of the real completion token count. - Anthropic passthrough: the client-facing SSE/JSON bytes are already relayed verbatim from Anthropic's own API, which already includes these fields correctly — no client-facing change needed. Only _parse_anthropic_sse_usage's raw_usage capture (for Mongo logging) was extended to carry message_start's cache fields forward instead of discarding everything but input_tokens/output_tokens. raw_usage (the full upstream usage object, verbatim) is now recorded to model-router-usage from every streaming path, matching what non-streaming already had.
…anslators
Both the native Converse and Mantle/OpenAI streaming translators could
produce a client-facing SSE stream with zero content blocks ever opened,
or with a content_block_delta referencing a block that was never
started — both are invalid Anthropic protocol and break Claude Code's
SSE parser ("Content block not found"), forcing a fallback to
non-streaming that starves the client's context-usage tracking (the
"0% until auto-compact" symptom).
- converse_stream_adapter.py: Bedrock's Converse API skips
contentBlockStart entirely for qwen.qwen3-coder-next over the native
transport and goes straight to contentBlockDelta — lazily open the
block on first delta instead of assuming a prior start. Also
guarantee at least one (empty) content block when a completion ends
with none opened and no error (e.g. the model exhausts max_tokens
entirely within hidden reasoning).
- stream_converter.py (Mantle/OpenAI path): same "guarantee at least
one content block" fix, plus fix the inline error-visibility guard to
check `not open_blocks` instead of `not sent_message_start` — the
latter is already True as soon as the first chunk arrives, so a
mid-stream error before any visible text was silently dropped instead
of surfacing as error text.
Added regression tests for both translators.
Records the 2026-07-14 investigation into the "0% until auto-compact" context-usage bug: symptom, root cause in both streaming translators, why the router's own logging didn't reveal which transport was actually live, the fix, and which other response paths were audited and confirmed already safe.
…ransport ran
MODEL_ROUTING_BEDROCK_TRANSPORT now defaults to "native" instead of
"mantle" — "mantle" becomes the manual opt-out instead of the default.
Mantle's own reliability (bare, undiagnosable internal_server_error
500s with no further diagnostic detail, as seen in a services-env
incident) makes native the better default now that both transports
have been proven out in production. Only the literal value "mantle"
(case-insensitive) opts back in; anything else, including a typo,
resolves to "native" rather than silently landing on the less
reliable transport. CodingModelRouter's own constructor default is
left at "mantle" since api.py always passes bedrock_transport
explicitly from this env var — flipping the constructor default too
would have silently changed which path ~20 existing tests exercise
for no behavioral benefit.
Every model-router-usage/model-router-errors record now also carries
a bedrock_transport field ("native"/"mantle") whenever
backend == "aws_bedrock". auth/api_type alone can't disambiguate this
today — native Converse dispatch is also reached via auth="aws",
api_type="openai" routes — which is exactly what made diagnosing the
services-env incident harder than it needed to be. record_error fills
this in automatically from CodingModelRouter._bedrock_transport (no
call-site changes needed); record_usage's callers pass it explicitly,
hardcoded per translator since each usage-tracking wrapper is
transport-specific by construction (_oai_stream_with_usage_tracking is
always Mantle, _converse_stream_with_usage_tracking and
BedrockNativeDispatcher are always native).
Removed the now-redundant MODEL_ROUTING_BEDROCK_TRANSPORT: native
override from docker-compose.yml (native is the default everywhere
now) and updated docs/coding_model_router.md and the original
native-transport design spec to reflect the new default and field.
…Bedrock dispatch paths - record_usage/session rollup now capture retry_count/total_retries across all six dispatch paths (Mantle + native Converse, streaming + non-streaming, plus the Anthropic-format Bedrock/passthrough path). - Native Bedrock Converse streaming now reports sse_event_count, matching Mantle's existing behavior - usage records for native-transport streams were silently missing this sanity field. - Dispatch pacing (_pace_bedrock_dispatch) was only wired into the unused Anthropic-format Bedrock route; now also applied to the Mantle and native Converse paths that actually serve the haiku/sonnet tiers in production, so concurrent requests no longer burst against Bedrock unthrottled.
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.
Summary
Implements native AWS Bedrock Converse API transport as a fallback mechanism when the standard
bedrock-runtimeAPI is unavailable. This allows the gateway to route requests directly to Bedrock's Converse API endpoint, bypassing the Mantle layer.Changes
New Features
bedrock_native_dispatcher.pywith full support for both streaming and non-streaming Converse API callsMODEL_ROUTING_BEDROCK_TRANSPORTenvironment variable enables native transport (defaults to off; switches to on in bd49c78)usage_tracker.pyto properly track cache and input/output token usage for all Converse API routesCode Structure
bedrock_native_dispatcher.py- Core native transport dispatcher (408 lines)bedrock_converse_client.py- Async wrapper for Bedrock Converse API (77 lines)converse_request_translator.py- OpenAI-to-Converse request conversion (249 lines)converse_stream_adapter.py- SSE stream adapter for streaming responses (383 lines)bedrock_client.py- Extracted shared client caching logicBug Fixes
Tests
test_coding_model_router.pywith native transport scenariosTesting Plan
Related
docs/superpowers/plans/2026-07-13-native-bedrock-transport.mddocs/superpowers/specs/2026-07-13-native-bedrock-transport-design.md