From b6706b50cff18923c8f161cd7ca820e9bc6614d2 Mon Sep 17 00:00:00 2001 From: mateo Date: Tue, 22 Sep 2026 22:46:28 +0000 Subject: [PATCH 1/7] fix(bedrock): stream /v1/messages Invoke bytes through instead of holding them in a 1024-byte chunker Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../anthropic_claude3_transformation.py | 12 +-- .../test_anthropic_claude3_transformation.py | 102 +++++++++++++++++- 2 files changed, 102 insertions(+), 12 deletions(-) diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index f46edc766c7e..35cb1abbe602 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -771,7 +771,7 @@ def get_async_streaming_response_iterator( model=model, ) completion_stream: Final = aws_decoder.aiter_bytes( - httpx_response.aiter_bytes(chunk_size=aws_decoder.DEFAULT_CHUNK_SIZE) + httpx_response.aiter_bytes() ) # Convert decoded Bedrock events to Server-Sent Events expected by Anthropic clients. return self.bedrock_sse_wrapper( @@ -919,16 +919,6 @@ async def _promote_message_stop_usage( class AmazonAnthropicClaudeMessagesStreamDecoder(AWSEventStreamDecoder): - def __init__( - self, - model: str, - ) -> None: - """ - Iterator to return Bedrock invoke response in anthropic /messages format - """ - super().__init__(model=model) - self.DEFAULT_CHUNK_SIZE = 1024 - def _chunk_parser(self, chunk_data: dict) -> GChunk | ModelResponseStream | dict: """ Parse the chunk data into anthropic /messages format diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index ea8b722b849e..c1db6e25af75 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -1,12 +1,16 @@ import asyncio +import base64 import copy import json import os +import struct +import zlib from datetime import datetime from types import SimpleNamespace -from typing import Final +from typing import AsyncIterator, Final from unittest.mock import Mock +import httpx import pytest # Ensure the project root is on the import path so `litellm` can be imported when @@ -3395,3 +3399,99 @@ def test_bedrock_invoke_eager_input_streaming_beta_not_duplicated_with_client_he ) assert result["anthropic_beta"] == [FINE_GRAINED_TOOL_STREAMING_BETA] + + +def _bedrock_event_frame(payload: dict) -> bytes: + def _header(name: str, value: str) -> bytes: + return ( + bytes([len(name)]) + + name.encode() + + bytes([7]) + + struct.pack(">H", len(value)) + + value.encode() + ) + + headers: Final = ( + _header(":message-type", "event") + + _header(":event-type", "chunk") + + _header(":content-type", "application/json") + ) + body: Final = json.dumps( + {"bytes": base64.b64encode(json.dumps(payload).encode()).decode()} + ).encode() + prelude: Final = struct.pack(">II", 12 + len(headers) + len(body) + 4, len(headers)) + prelude_crc: Final = struct.pack(">I", zlib.crc32(prelude)) + message_crc: Final = struct.pack(">I", zlib.crc32(prelude + prelude_crc + headers + body)) + return prelude + prelude_crc + headers + body + message_crc + + +class _GatedAsyncByteStream(httpx.AsyncByteStream): + def __init__(self, chunks: list[bytes], gate: asyncio.Event) -> None: + self._chunks = chunks + self._gate = gate + + async def __aiter__(self) -> AsyncIterator[bytes]: + yield self._chunks[0] + await self._gate.wait() + for chunk in self._chunks[1:]: + yield chunk + + async def aclose(self) -> None: + return None + + +@pytest.mark.asyncio +async def test_get_async_streaming_response_iterator_yields_small_frame_before_upstream_pauses(): + gate: Final = asyncio.Event() + response: Final = httpx.Response( + 200, + stream=_GatedAsyncByteStream( + chunks=[ + _bedrock_event_frame( + { + "type": "message_start", + "message": { + "id": "msg_test", + "type": "message", + "role": "assistant", + "content": [], + "model": "us.anthropic.claude-sonnet-4-6", + "usage": {"input_tokens": 3, "output_tokens": 1}, + }, + } + ), + _bedrock_event_frame( + { + "type": "message_stop", + "usage": {"input_tokens": 3, "output_tokens": 9}, + } + ), + ], + gate=gate, + ), + ) + + iterator: Final = AmazonAnthropicClaudeMessagesConfig().get_async_streaming_response_iterator( + model="us.anthropic.claude-sonnet-4-6", + httpx_response=response, + request_body={"model": "us.anthropic.claude-sonnet-4-6"}, + litellm_logging_obj=LiteLLMLoggingObj( + model="bedrock/us.anthropic.claude-sonnet-4-6", + messages=[{"role": "user", "content": "Hello"}], + stream=True, + call_type="chat", + start_time=datetime.now(), + litellm_call_id="test_small_frame_before_upstream_pauses", + function_id="test_small_frame_before_upstream_pauses", + ), + ) + + first: Final = await asyncio.wait_for(anext(iterator), timeout=1) + assert first.startswith(b"event: message_start\n"), first + + gate.set() + remaining: list[bytes] = [] + async for chunk in iterator: + remaining.append(chunk) + assert any(chunk.startswith(b"event: message_stop\n") for chunk in remaining), remaining + await iterator.aclose() From 2b4d0ca8992136aaa68bfaa2bb86ce08a910bb15 Mon Sep 17 00:00:00 2001 From: mateo Date: Tue, 22 Sep 2026 22:54:06 +0000 Subject: [PATCH 2/7] style(bedrock): apply ruff format to invoke messages stream passthrough Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../anthropic_claude3_transformation.py | 4 +- .../test_anthropic_claude3_transformation.py | 189 +++++------------- 2 files changed, 55 insertions(+), 138 deletions(-) diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index 35cb1abbe602..14bd2bee6cfd 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -770,9 +770,7 @@ def get_async_streaming_response_iterator( aws_decoder: Final = AmazonAnthropicClaudeMessagesStreamDecoder( model=model, ) - completion_stream: Final = aws_decoder.aiter_bytes( - httpx_response.aiter_bytes() - ) + completion_stream: Final = aws_decoder.aiter_bytes(httpx_response.aiter_bytes()) # Convert decoded Bedrock events to Server-Sent Events expected by Anthropic clients. return self.bedrock_sse_wrapper( completion_stream=completion_stream, diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index c1db6e25af75..76a50e5165f4 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -53,9 +53,7 @@ async def _dummy_stream(): # type: ignore[return-type] _dummy_stream(), litellm_logging_obj=LiteLLMLoggingObj( model="bedrock/invoke/anthropic.claude-3-sonnet-20240229-v1:0", - messages=[ - {"role": "user", "content": "Hello, can you tell me a short joke?"} - ], + messages=[{"role": "user", "content": "Hello, can you tell me a short joke?"}], stream=True, call_type="chat", start_time=datetime.now(), @@ -232,9 +230,7 @@ async def _dummy_stream(): # type: ignore[return-type] def test_chunk_parser_usage_transformation(): """Ensure Bedrock invocation metrics are transformed to Anthropic usage keys.""" - decoder = AmazonAnthropicClaudeMessagesStreamDecoder( - model="bedrock/invoke/anthropic.claude-3-sonnet-20240229-v1:0" - ) + decoder = AmazonAnthropicClaudeMessagesStreamDecoder(model="bedrock/invoke/anthropic.claude-3-sonnet-20240229-v1:0") chunk = { "type": "message_delta", @@ -263,9 +259,7 @@ def test_chunk_parser_preserves_cache_usage_fields_with_invocation_metrics(): fields and cache tokens end up billed at $0. """ - decoder = AmazonAnthropicClaudeMessagesStreamDecoder( - model="bedrock/invoke/anthropic.claude-sonnet-4-6" - ) + decoder = AmazonAnthropicClaudeMessagesStreamDecoder(model="bedrock/invoke/anthropic.claude-sonnet-4-6") chunk = { "type": "message_stop", @@ -291,9 +285,7 @@ def test_chunk_parser_preserves_cache_usage_fields_with_invocation_metrics(): def test_chunk_parser_maps_cache_token_counts_from_invocation_metrics(): """Cache itemization inside invocationMetrics maps to Anthropic usage keys.""" - decoder = AmazonAnthropicClaudeMessagesStreamDecoder( - model="bedrock/invoke/anthropic.claude-sonnet-4-6" - ) + decoder = AmazonAnthropicClaudeMessagesStreamDecoder(model="bedrock/invoke/anthropic.claude-sonnet-4-6") chunk = { "type": "message_stop", @@ -316,9 +308,7 @@ def test_chunk_parser_maps_cache_token_counts_from_invocation_metrics(): def test_chunk_parser_keeps_existing_token_counts_over_invocation_metrics(): """Token counts reported in the chunk's own usage block win over invocationMetrics.""" - decoder = AmazonAnthropicClaudeMessagesStreamDecoder( - model="bedrock/invoke/anthropic.claude-sonnet-4-6" - ) + decoder = AmazonAnthropicClaudeMessagesStreamDecoder(model="bedrock/invoke/anthropic.claude-sonnet-4-6") chunk = { "type": "message_stop", @@ -353,9 +343,7 @@ async def test_bedrock_sse_wrapper_preserves_cache_usage_with_invocation_metrics final usage billed cache reads and writes at $0. """ - decoder = AmazonAnthropicClaudeMessagesStreamDecoder( - model="bedrock/invoke/anthropic.claude-sonnet-4-6" - ) + decoder = AmazonAnthropicClaudeMessagesStreamDecoder(model="bedrock/invoke/anthropic.claude-sonnet-4-6") cfg = AmazonAnthropicClaudeMessagesConfig() raw_chunks = [ @@ -565,11 +553,7 @@ def test_normalize_custom_field_on_tools(): assert request4["tools"] is None # Case 5: an explicit top-level flag wins over a conflicting wrapped one - request5 = { - "tools": [ - {"name": "Read", "defer_loading": False, "custom": {"defer_loading": True}} - ] - } + request5 = {"tools": [{"name": "Read", "defer_loading": False, "custom": {"defer_loading": True}}]} normalize_custom_field_on_tools(request5) assert request5["tools"][0] == {"name": "Read", "defer_loading": False} @@ -590,9 +574,7 @@ def test_normalize_custom_field_on_tools(): assert request7["tools"] == [{"name": "Read"}, {"name": "Write"}] -@pytest.mark.parametrize( - "deferred_marker", [{"custom": {"defer_loading": True}}, {"defer_loading": True}] -) +@pytest.mark.parametrize("deferred_marker", [{"custom": {"defer_loading": True}}, {"defer_loading": True}]) def test_bedrock_invoke_messages_transform_emits_top_level_defer_loading( deferred_marker, ): @@ -725,9 +707,7 @@ def test_bedrock_invoke_messages_skips_thinking_injection_when_already_enabled( "max_tokens": 32000, "stream": False, "thinking": {"type": "enabled", "budget_tokens": 2048}, - "context_management": { - "edits": [{"type": "clear_thinking_20251015", "keep": "all"}] - }, + "context_management": {"edits": [{"type": "clear_thinking_20251015", "keep": "all"}]}, } result = cfg.transform_anthropic_messages_request( model="global.anthropic.claude-sonnet-4-6-v1:0", @@ -829,9 +809,7 @@ def test_remove_ttl_from_cache_control_processes_tools(local_model_cost_map): "messages": [], } - cfg._remove_ttl_from_cache_control( - request, model="anthropic.claude-3-5-sonnet-20241022-v2:0" - ) + cfg._remove_ttl_from_cache_control(request, model="anthropic.claude-3-5-sonnet-20241022-v2:0") # Tool ttl should be stripped assert "ttl" not in request["tools"][0]["cache_control"] @@ -867,9 +845,7 @@ def test_remove_ttl_from_cache_control_preserves_tools_ttl_for_claude_4_5(local_ ], } - cfg._remove_ttl_from_cache_control( - request, model="us.anthropic.claude-sonnet-4-5-20250929-v1:0" - ) + cfg._remove_ttl_from_cache_control(request, model="us.anthropic.claude-sonnet-4-5-20250929-v1:0") # Both tools and system should preserve ttl for Claude 4.5 assert request["tools"][0]["cache_control"]["ttl"] == "1h" @@ -953,9 +929,7 @@ def test_bedrock_messages_strips_output_config(): headers={}, ) - assert "output_config" not in result, ( - "output_config should be stripped for models that don't support it" - ) + assert "output_config" not in result, "output_config should be stripped for models that don't support it" assert result.get("max_tokens") == 4096 @@ -988,9 +962,7 @@ def test_bedrock_messages_preserves_output_config_for_claude_4_6(): headers={}, ) - assert "output_config" in result, ( - "output_config should be preserved for supported models" - ) + assert "output_config" in result, "output_config should be preserved for supported models" assert result["output_config"] == {"effort": "high"} assert result.get("max_tokens") == 4096 @@ -1142,9 +1114,7 @@ def test_bedrock_messages_converts_output_config_format_to_inline_schema(): ("anthropic.claude-opus-4-7", "xhigh"), ], ) -def test_bedrock_messages_normalizes_output_config_effort_for_opus( - model, expected_effort -): +def test_bedrock_messages_normalizes_output_config_effort_for_opus(model, expected_effort): """Bedrock /v1/messages accepts ``xhigh`` and forwards the provider-safe effort.""" from unittest.mock import patch @@ -1202,9 +1172,7 @@ def test_bedrock_messages_does_not_mutate_callers_messages_when_embedding_schema headers={}, ) - assert caller_messages == [ - {"role": "user", "content": [{"type": "text", "text": "Hello"}]} - ] + assert caller_messages == [{"role": "user", "content": [{"type": "text", "text": "Hello"}]}] assert caller_message == { "role": "user", "content": [{"type": "text", "text": "Hello"}], @@ -1520,9 +1488,7 @@ def test_bedrock_messages_strips_context_management(): messages = [{"role": "user", "content": [{"type": "text", "text": "Hello"}]}] optional_params = { "max_tokens": 4096, - "context_management": { - "edits": [{"type": "clear_thinking_20251015", "keep": "all"}] - }, + "context_management": {"edits": [{"type": "clear_thinking_20251015", "keep": "all"}]}, } result = cfg.transform_anthropic_messages_request( @@ -1533,9 +1499,7 @@ def test_bedrock_messages_strips_context_management(): headers={}, ) - assert "context_management" not in result, ( - "context_management should be stripped — Bedrock Invoke rejects it" - ) + assert "context_management" not in result, "context_management should be stripped — Bedrock Invoke rejects it" assert result.get("max_tokens") == 4096 @@ -1660,7 +1624,9 @@ def test_bedrock_messages_allowlist_filters_anthropic_only_fields(): ["dangerous-tool-use-2026-09-03,interleaved-thinking-2025-05-14", "interleaved-thinking-2025-05-14"], ids=["client_sends_beta", "client_omits_beta"], ) -def test_bedrock_messages_forwards_safeguards_with_dangerous_tool_use_beta(local_beta_headers_config, client_beta_header): +def test_bedrock_messages_forwards_safeguards_with_dangerous_tool_use_beta( + local_beta_headers_config, client_beta_header +): """ Claude Code's server-side auto-mode classifier sends `safeguards` alongside the dangerous-tool-use-2026-09-03 beta. Bedrock Invoke accepts the pair, answers @@ -1768,12 +1734,8 @@ def test_bedrock_messages_filters_user_provided_unsupported_beta_header(): ) betas = result.get("anthropic_beta") or [] - assert "advisor-tool-2026-03-01" not in betas, ( - "user-provided beta not in the Bedrock mapping must be dropped" - ) - assert "context-1m-2025-08-07" in betas, ( - "user-provided beta that IS in the Bedrock mapping should survive" - ) + assert "advisor-tool-2026-03-01" not in betas, "user-provided beta not in the Bedrock mapping must be dropped" + assert "context-1m-2025-08-07" in betas, "user-provided beta that IS in the Bedrock mapping should survive" def test_bedrock_messages_renames_user_provided_aliased_beta_header(): @@ -1801,9 +1763,7 @@ def test_bedrock_messages_renames_user_provided_aliased_beta_header(): assert "advanced-tool-use-2025-11-20" not in betas, ( "Anthropic-direct spelling should be rewritten, not forwarded verbatim" ) - assert "tool-search-tool-2025-10-19" in betas, ( - "user-provided beta should be renamed to the Bedrock-side spelling" - ) + assert "tool-search-tool-2025-10-19" in betas, "user-provided beta should be renamed to the Bedrock-side spelling" @pytest.mark.asyncio @@ -2065,9 +2025,7 @@ async def _stream(): # type: ignore[return-type] "global.anthropic.claude-fable-5", ], ) -def test_bedrock_clear_thinking_injects_adaptive_with_effort_for_adaptive_models( - local_model_cost_map, model -): +def test_bedrock_clear_thinking_injects_adaptive_with_effort_for_adaptive_models(local_model_cost_map, model): """clear_thinking_20251015 without a top-level ``thinking`` field must inject ``thinking.type=adaptive`` plus ``output_config.effort`` on adaptive-thinking models (Opus 4.7/4.8, Fable 5). The legacy ``thinking.type=enabled`` shape is @@ -2077,9 +2035,7 @@ def test_bedrock_clear_thinking_injects_adaptive_with_effort_for_adaptive_models cfg = AmazonAnthropicClaudeMessagesConfig() request = { "max_tokens": 32000, - "context_management": { - "edits": [{"type": "clear_thinking_20251015", "keep": "all"}] - }, + "context_management": {"edits": [{"type": "clear_thinking_20251015", "keep": "all"}]}, } changed = cfg._ensure_thinking_for_clear_thinking_context_management( @@ -2102,9 +2058,7 @@ def test_bedrock_clear_thinking_converts_legacy_enabled_budget_to_effort(): "type": "enabled", "budget_tokens": DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, }, - "context_management": { - "edits": [{"type": "clear_thinking_20251015", "keep": "all"}] - }, + "context_management": {"edits": [{"type": "clear_thinking_20251015", "keep": "all"}]}, } changed = cfg._ensure_thinking_for_clear_thinking_context_management( @@ -2122,10 +2076,7 @@ def test_resolve_clear_thinking_budget_tokens_honors_explicit_zero(): and only fall back to the minimum when the caller omits the budget.""" cfg = AmazonAnthropicClaudeMessagesConfig() assert cfg._resolve_clear_thinking_budget_tokens(0) == 0 - assert ( - cfg._resolve_clear_thinking_budget_tokens(None) - == BEDROCK_MIN_THINKING_BUDGET_TOKENS - ) + assert cfg._resolve_clear_thinking_budget_tokens(None) == BEDROCK_MIN_THINKING_BUDGET_TOKENS assert cfg._resolve_clear_thinking_budget_tokens(12000) == 12000 @@ -2135,9 +2086,7 @@ def test_bedrock_clear_thinking_keeps_enabled_for_non_adaptive_models(): cfg = AmazonAnthropicClaudeMessagesConfig() request = { "max_tokens": 32000, - "context_management": { - "edits": [{"type": "clear_thinking_20251015", "keep": "all"}] - }, + "context_management": {"edits": [{"type": "clear_thinking_20251015", "keep": "all"}]}, } changed = cfg._ensure_thinking_for_clear_thinking_context_management( @@ -2162,9 +2111,7 @@ def test_bedrock_invoke_transform_emits_adaptive_thinking_for_opus_4_8(): optional_params = { "max_tokens": 32000, "stream": False, - "context_management": { - "edits": [{"type": "clear_thinking_20251015", "keep": "all"}] - }, + "context_management": {"edits": [{"type": "clear_thinking_20251015", "keep": "all"}]}, } result = cfg.transform_anthropic_messages_request( @@ -2201,9 +2148,7 @@ def test_bedrock_invoke_transform_normalizes_system_role_message_into_system(): assert all(m.get("role") != "system" for m in result["messages"]) assert result["messages"] == [{"role": "user", "content": "hi"}] - assert result["system"] == [ - {"type": "text", "text": "You are a careful assistant."} - ] + assert result["system"] == [{"type": "text", "text": "You are a careful assistant."}] def test_bedrock_invoke_transform_merges_system_role_into_existing_system(): @@ -2318,9 +2263,7 @@ def test_bedrock_invoke_transform_keeps_mid_conversation_system_role_in_place(lo ) assert result["messages"] == messages - assert result["system"] == [ - {"type": "text", "text": "Base.", "cache_control": {"type": "ephemeral"}} - ] + assert result["system"] == [{"type": "text", "text": "Base.", "cache_control": {"type": "ephemeral"}}] def test_bedrock_invoke_transform_hoists_only_leading_system_run(local_model_cost_map): @@ -2503,13 +2446,13 @@ def test_bedrock_invoke_transform_converted_system_carries_only_its_content(loca assert result["messages"][2] == { "role": "user", "content": [ - { - "type": "text", - "text": ( - "Operator note (not from the user): the following was " - "originally a mid-conversation system-role reminder." - ), - }, + { + "type": "text", + "text": ( + "Operator note (not from the user): the following was " + "originally a mid-conversation system-role reminder." + ), + }, {"type": "text", "text": "[Truncated: PARTIAL view of big1.txt]"}, ], } @@ -2645,10 +2588,7 @@ def test_as_system_content_blocks_handles_each_shape(): def test_effort_from_thinking_budget_tiers(budget_tokens, expected_effort): """The budget -> effort mapping pins each tier boundary so a shifted threshold is caught.""" - assert ( - AmazonAnthropicClaudeMessagesConfig._effort_from_thinking_budget(budget_tokens) - == expected_effort - ) + assert AmazonAnthropicClaudeMessagesConfig._effort_from_thinking_budget(budget_tokens) == expected_effort def test_inject_adaptive_thinking_preserves_existing_effort(): @@ -2657,9 +2597,7 @@ def test_inject_adaptive_thinking_preserves_existing_effort(): cfg = AmazonAnthropicClaudeMessagesConfig() request = {"output_config": {"effort": "max", "other": "keep"}} - cfg._inject_adaptive_thinking_for_clear_thinking( - request, budget_tokens=24000, model="us.anthropic.claude-fable-5" - ) + cfg._inject_adaptive_thinking_for_clear_thinking(request, budget_tokens=24000, model="us.anthropic.claude-fable-5") assert request["thinking"] == {"type": "adaptive"} assert request["output_config"] == {"effort": "max", "other": "keep"} @@ -2672,9 +2610,7 @@ def test_bedrock_clear_thinking_noops_when_thinking_already_adaptive(): request = { "max_tokens": 32000, "thinking": {"type": "adaptive"}, - "context_management": { - "edits": [{"type": "clear_thinking_20251015", "keep": "all"}] - }, + "context_management": {"edits": [{"type": "clear_thinking_20251015", "keep": "all"}]}, } changed = cfg._ensure_thinking_for_clear_thinking_context_management( @@ -2694,9 +2630,7 @@ def test_bedrock_clear_thinking_replaces_disabled_thinking_on_adaptive_model(): request = { "max_tokens": 32000, "thinking": {"type": "disabled"}, - "context_management": { - "edits": [{"type": "clear_thinking_20251015", "keep": "all"}] - }, + "context_management": {"edits": [{"type": "clear_thinking_20251015", "keep": "all"}]}, } changed = cfg._ensure_thinking_for_clear_thinking_context_management( @@ -2716,9 +2650,7 @@ def test_bedrock_clear_thinking_leaves_enabled_thinking_on_non_adaptive_model(): request = { "max_tokens": 32000, "thinking": {"type": "enabled", "budget_tokens": 8000}, - "context_management": { - "edits": [{"type": "clear_thinking_20251015", "keep": "all"}] - }, + "context_management": {"edits": [{"type": "clear_thinking_20251015", "keep": "all"}]}, } changed = cfg._ensure_thinking_for_clear_thinking_context_management( @@ -2753,9 +2685,7 @@ def test_bedrock_messages_preserves_clear_tool_uses_context_management_and_adds_ messages = [{"role": "user", "content": [{"type": "text", "text": "Hi"}]}] optional_params = { "max_tokens": 4096, - "context_management": { - "edits": [{"type": "clear_tool_uses_20250919"}] - }, + "context_management": {"edits": [{"type": "clear_tool_uses_20250919"}]}, } result = cfg.transform_anthropic_messages_request( @@ -2766,12 +2696,11 @@ def test_bedrock_messages_preserves_clear_tool_uses_context_management_and_adds_ headers={}, ) - assert result.get("context_management") == { - "edits": [{"type": "clear_tool_uses_20250919"}] - }, "clear_tool_uses_20250919 edit must reach Bedrock InvokeModel body" + assert result.get("context_management") == {"edits": [{"type": "clear_tool_uses_20250919"}]}, ( + "clear_tool_uses_20250919 edit must reach Bedrock InvokeModel body" + ) assert "context-management-2025-06-27" in result.get("anthropic_beta", []), ( - "context-management-2025-06-27 beta must reach the InvokeModel body so " - "the tool-call-clearing edit is accepted" + "context-management-2025-06-27 beta must reach the InvokeModel body so the tool-call-clearing edit is accepted" ) @@ -2848,9 +2777,9 @@ def test_bedrock_messages_filters_clear_thinking_keeps_clear_tool_uses( cm = result.get("context_management") assert cm is not None - assert [e.get("type") for e in cm["edits"]] == [ - "clear_tool_uses_20250919" - ], "clear_thinking_20251015 must still be stripped (LiteLLM-internal)" + assert [e.get("type") for e in cm["edits"]] == ["clear_tool_uses_20250919"], ( + "clear_thinking_20251015 must still be stripped (LiteLLM-internal)" + ) betas = result.get("anthropic_beta", []) assert "context-management-2025-06-27" in betas @@ -2991,9 +2920,7 @@ def test_bedrock_messages_tool_search_follows_claude_tool_search_rule(local_mode assert cfg._supports_tool_search_on_bedrock(model) is expected -def test_bedrock_messages_thinking_shape_follows_exact_bedrock_entry_flag( - local_model_cost_map, monkeypatch -): +def test_bedrock_messages_thinking_shape_follows_exact_bedrock_entry_flag(local_model_cost_map, monkeypatch): """The outbound thinking payload must follow the exact Bedrock cost-map entry. Before threading the caller's provider through the capability probes, the probe was pinned to ``"anthropic"``: the exact ``global.anthropic.claude-opus-4-8`` @@ -3403,22 +3330,14 @@ def test_bedrock_invoke_eager_input_streaming_beta_not_duplicated_with_client_he def _bedrock_event_frame(payload: dict) -> bytes: def _header(name: str, value: str) -> bytes: - return ( - bytes([len(name)]) - + name.encode() - + bytes([7]) - + struct.pack(">H", len(value)) - + value.encode() - ) + return bytes([len(name)]) + name.encode() + bytes([7]) + struct.pack(">H", len(value)) + value.encode() headers: Final = ( _header(":message-type", "event") + _header(":event-type", "chunk") + _header(":content-type", "application/json") ) - body: Final = json.dumps( - {"bytes": base64.b64encode(json.dumps(payload).encode()).decode()} - ).encode() + body: Final = json.dumps({"bytes": base64.b64encode(json.dumps(payload).encode()).decode()}).encode() prelude: Final = struct.pack(">II", 12 + len(headers) + len(body) + 4, len(headers)) prelude_crc: Final = struct.pack(">I", zlib.crc32(prelude)) message_crc: Final = struct.pack(">I", zlib.crc32(prelude + prelude_crc + headers + body)) From d3dcd9a437d78d823bec7a31454bb9a6c2d9200d Mon Sep 17 00:00:00 2001 From: mateo Date: Tue, 22 Sep 2026 22:54:24 +0000 Subject: [PATCH 3/7] style(bedrock): drop drive-by reformat of existing invoke messages tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_anthropic_claude3_transformation.py | 189 +++++++++++++----- 1 file changed, 135 insertions(+), 54 deletions(-) diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index 76a50e5165f4..c1db6e25af75 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -53,7 +53,9 @@ async def _dummy_stream(): # type: ignore[return-type] _dummy_stream(), litellm_logging_obj=LiteLLMLoggingObj( model="bedrock/invoke/anthropic.claude-3-sonnet-20240229-v1:0", - messages=[{"role": "user", "content": "Hello, can you tell me a short joke?"}], + messages=[ + {"role": "user", "content": "Hello, can you tell me a short joke?"} + ], stream=True, call_type="chat", start_time=datetime.now(), @@ -230,7 +232,9 @@ async def _dummy_stream(): # type: ignore[return-type] def test_chunk_parser_usage_transformation(): """Ensure Bedrock invocation metrics are transformed to Anthropic usage keys.""" - decoder = AmazonAnthropicClaudeMessagesStreamDecoder(model="bedrock/invoke/anthropic.claude-3-sonnet-20240229-v1:0") + decoder = AmazonAnthropicClaudeMessagesStreamDecoder( + model="bedrock/invoke/anthropic.claude-3-sonnet-20240229-v1:0" + ) chunk = { "type": "message_delta", @@ -259,7 +263,9 @@ def test_chunk_parser_preserves_cache_usage_fields_with_invocation_metrics(): fields and cache tokens end up billed at $0. """ - decoder = AmazonAnthropicClaudeMessagesStreamDecoder(model="bedrock/invoke/anthropic.claude-sonnet-4-6") + decoder = AmazonAnthropicClaudeMessagesStreamDecoder( + model="bedrock/invoke/anthropic.claude-sonnet-4-6" + ) chunk = { "type": "message_stop", @@ -285,7 +291,9 @@ def test_chunk_parser_preserves_cache_usage_fields_with_invocation_metrics(): def test_chunk_parser_maps_cache_token_counts_from_invocation_metrics(): """Cache itemization inside invocationMetrics maps to Anthropic usage keys.""" - decoder = AmazonAnthropicClaudeMessagesStreamDecoder(model="bedrock/invoke/anthropic.claude-sonnet-4-6") + decoder = AmazonAnthropicClaudeMessagesStreamDecoder( + model="bedrock/invoke/anthropic.claude-sonnet-4-6" + ) chunk = { "type": "message_stop", @@ -308,7 +316,9 @@ def test_chunk_parser_maps_cache_token_counts_from_invocation_metrics(): def test_chunk_parser_keeps_existing_token_counts_over_invocation_metrics(): """Token counts reported in the chunk's own usage block win over invocationMetrics.""" - decoder = AmazonAnthropicClaudeMessagesStreamDecoder(model="bedrock/invoke/anthropic.claude-sonnet-4-6") + decoder = AmazonAnthropicClaudeMessagesStreamDecoder( + model="bedrock/invoke/anthropic.claude-sonnet-4-6" + ) chunk = { "type": "message_stop", @@ -343,7 +353,9 @@ async def test_bedrock_sse_wrapper_preserves_cache_usage_with_invocation_metrics final usage billed cache reads and writes at $0. """ - decoder = AmazonAnthropicClaudeMessagesStreamDecoder(model="bedrock/invoke/anthropic.claude-sonnet-4-6") + decoder = AmazonAnthropicClaudeMessagesStreamDecoder( + model="bedrock/invoke/anthropic.claude-sonnet-4-6" + ) cfg = AmazonAnthropicClaudeMessagesConfig() raw_chunks = [ @@ -553,7 +565,11 @@ def test_normalize_custom_field_on_tools(): assert request4["tools"] is None # Case 5: an explicit top-level flag wins over a conflicting wrapped one - request5 = {"tools": [{"name": "Read", "defer_loading": False, "custom": {"defer_loading": True}}]} + request5 = { + "tools": [ + {"name": "Read", "defer_loading": False, "custom": {"defer_loading": True}} + ] + } normalize_custom_field_on_tools(request5) assert request5["tools"][0] == {"name": "Read", "defer_loading": False} @@ -574,7 +590,9 @@ def test_normalize_custom_field_on_tools(): assert request7["tools"] == [{"name": "Read"}, {"name": "Write"}] -@pytest.mark.parametrize("deferred_marker", [{"custom": {"defer_loading": True}}, {"defer_loading": True}]) +@pytest.mark.parametrize( + "deferred_marker", [{"custom": {"defer_loading": True}}, {"defer_loading": True}] +) def test_bedrock_invoke_messages_transform_emits_top_level_defer_loading( deferred_marker, ): @@ -707,7 +725,9 @@ def test_bedrock_invoke_messages_skips_thinking_injection_when_already_enabled( "max_tokens": 32000, "stream": False, "thinking": {"type": "enabled", "budget_tokens": 2048}, - "context_management": {"edits": [{"type": "clear_thinking_20251015", "keep": "all"}]}, + "context_management": { + "edits": [{"type": "clear_thinking_20251015", "keep": "all"}] + }, } result = cfg.transform_anthropic_messages_request( model="global.anthropic.claude-sonnet-4-6-v1:0", @@ -809,7 +829,9 @@ def test_remove_ttl_from_cache_control_processes_tools(local_model_cost_map): "messages": [], } - cfg._remove_ttl_from_cache_control(request, model="anthropic.claude-3-5-sonnet-20241022-v2:0") + cfg._remove_ttl_from_cache_control( + request, model="anthropic.claude-3-5-sonnet-20241022-v2:0" + ) # Tool ttl should be stripped assert "ttl" not in request["tools"][0]["cache_control"] @@ -845,7 +867,9 @@ def test_remove_ttl_from_cache_control_preserves_tools_ttl_for_claude_4_5(local_ ], } - cfg._remove_ttl_from_cache_control(request, model="us.anthropic.claude-sonnet-4-5-20250929-v1:0") + cfg._remove_ttl_from_cache_control( + request, model="us.anthropic.claude-sonnet-4-5-20250929-v1:0" + ) # Both tools and system should preserve ttl for Claude 4.5 assert request["tools"][0]["cache_control"]["ttl"] == "1h" @@ -929,7 +953,9 @@ def test_bedrock_messages_strips_output_config(): headers={}, ) - assert "output_config" not in result, "output_config should be stripped for models that don't support it" + assert "output_config" not in result, ( + "output_config should be stripped for models that don't support it" + ) assert result.get("max_tokens") == 4096 @@ -962,7 +988,9 @@ def test_bedrock_messages_preserves_output_config_for_claude_4_6(): headers={}, ) - assert "output_config" in result, "output_config should be preserved for supported models" + assert "output_config" in result, ( + "output_config should be preserved for supported models" + ) assert result["output_config"] == {"effort": "high"} assert result.get("max_tokens") == 4096 @@ -1114,7 +1142,9 @@ def test_bedrock_messages_converts_output_config_format_to_inline_schema(): ("anthropic.claude-opus-4-7", "xhigh"), ], ) -def test_bedrock_messages_normalizes_output_config_effort_for_opus(model, expected_effort): +def test_bedrock_messages_normalizes_output_config_effort_for_opus( + model, expected_effort +): """Bedrock /v1/messages accepts ``xhigh`` and forwards the provider-safe effort.""" from unittest.mock import patch @@ -1172,7 +1202,9 @@ def test_bedrock_messages_does_not_mutate_callers_messages_when_embedding_schema headers={}, ) - assert caller_messages == [{"role": "user", "content": [{"type": "text", "text": "Hello"}]}] + assert caller_messages == [ + {"role": "user", "content": [{"type": "text", "text": "Hello"}]} + ] assert caller_message == { "role": "user", "content": [{"type": "text", "text": "Hello"}], @@ -1488,7 +1520,9 @@ def test_bedrock_messages_strips_context_management(): messages = [{"role": "user", "content": [{"type": "text", "text": "Hello"}]}] optional_params = { "max_tokens": 4096, - "context_management": {"edits": [{"type": "clear_thinking_20251015", "keep": "all"}]}, + "context_management": { + "edits": [{"type": "clear_thinking_20251015", "keep": "all"}] + }, } result = cfg.transform_anthropic_messages_request( @@ -1499,7 +1533,9 @@ def test_bedrock_messages_strips_context_management(): headers={}, ) - assert "context_management" not in result, "context_management should be stripped — Bedrock Invoke rejects it" + assert "context_management" not in result, ( + "context_management should be stripped — Bedrock Invoke rejects it" + ) assert result.get("max_tokens") == 4096 @@ -1624,9 +1660,7 @@ def test_bedrock_messages_allowlist_filters_anthropic_only_fields(): ["dangerous-tool-use-2026-09-03,interleaved-thinking-2025-05-14", "interleaved-thinking-2025-05-14"], ids=["client_sends_beta", "client_omits_beta"], ) -def test_bedrock_messages_forwards_safeguards_with_dangerous_tool_use_beta( - local_beta_headers_config, client_beta_header -): +def test_bedrock_messages_forwards_safeguards_with_dangerous_tool_use_beta(local_beta_headers_config, client_beta_header): """ Claude Code's server-side auto-mode classifier sends `safeguards` alongside the dangerous-tool-use-2026-09-03 beta. Bedrock Invoke accepts the pair, answers @@ -1734,8 +1768,12 @@ def test_bedrock_messages_filters_user_provided_unsupported_beta_header(): ) betas = result.get("anthropic_beta") or [] - assert "advisor-tool-2026-03-01" not in betas, "user-provided beta not in the Bedrock mapping must be dropped" - assert "context-1m-2025-08-07" in betas, "user-provided beta that IS in the Bedrock mapping should survive" + assert "advisor-tool-2026-03-01" not in betas, ( + "user-provided beta not in the Bedrock mapping must be dropped" + ) + assert "context-1m-2025-08-07" in betas, ( + "user-provided beta that IS in the Bedrock mapping should survive" + ) def test_bedrock_messages_renames_user_provided_aliased_beta_header(): @@ -1763,7 +1801,9 @@ def test_bedrock_messages_renames_user_provided_aliased_beta_header(): assert "advanced-tool-use-2025-11-20" not in betas, ( "Anthropic-direct spelling should be rewritten, not forwarded verbatim" ) - assert "tool-search-tool-2025-10-19" in betas, "user-provided beta should be renamed to the Bedrock-side spelling" + assert "tool-search-tool-2025-10-19" in betas, ( + "user-provided beta should be renamed to the Bedrock-side spelling" + ) @pytest.mark.asyncio @@ -2025,7 +2065,9 @@ async def _stream(): # type: ignore[return-type] "global.anthropic.claude-fable-5", ], ) -def test_bedrock_clear_thinking_injects_adaptive_with_effort_for_adaptive_models(local_model_cost_map, model): +def test_bedrock_clear_thinking_injects_adaptive_with_effort_for_adaptive_models( + local_model_cost_map, model +): """clear_thinking_20251015 without a top-level ``thinking`` field must inject ``thinking.type=adaptive`` plus ``output_config.effort`` on adaptive-thinking models (Opus 4.7/4.8, Fable 5). The legacy ``thinking.type=enabled`` shape is @@ -2035,7 +2077,9 @@ def test_bedrock_clear_thinking_injects_adaptive_with_effort_for_adaptive_models cfg = AmazonAnthropicClaudeMessagesConfig() request = { "max_tokens": 32000, - "context_management": {"edits": [{"type": "clear_thinking_20251015", "keep": "all"}]}, + "context_management": { + "edits": [{"type": "clear_thinking_20251015", "keep": "all"}] + }, } changed = cfg._ensure_thinking_for_clear_thinking_context_management( @@ -2058,7 +2102,9 @@ def test_bedrock_clear_thinking_converts_legacy_enabled_budget_to_effort(): "type": "enabled", "budget_tokens": DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, }, - "context_management": {"edits": [{"type": "clear_thinking_20251015", "keep": "all"}]}, + "context_management": { + "edits": [{"type": "clear_thinking_20251015", "keep": "all"}] + }, } changed = cfg._ensure_thinking_for_clear_thinking_context_management( @@ -2076,7 +2122,10 @@ def test_resolve_clear_thinking_budget_tokens_honors_explicit_zero(): and only fall back to the minimum when the caller omits the budget.""" cfg = AmazonAnthropicClaudeMessagesConfig() assert cfg._resolve_clear_thinking_budget_tokens(0) == 0 - assert cfg._resolve_clear_thinking_budget_tokens(None) == BEDROCK_MIN_THINKING_BUDGET_TOKENS + assert ( + cfg._resolve_clear_thinking_budget_tokens(None) + == BEDROCK_MIN_THINKING_BUDGET_TOKENS + ) assert cfg._resolve_clear_thinking_budget_tokens(12000) == 12000 @@ -2086,7 +2135,9 @@ def test_bedrock_clear_thinking_keeps_enabled_for_non_adaptive_models(): cfg = AmazonAnthropicClaudeMessagesConfig() request = { "max_tokens": 32000, - "context_management": {"edits": [{"type": "clear_thinking_20251015", "keep": "all"}]}, + "context_management": { + "edits": [{"type": "clear_thinking_20251015", "keep": "all"}] + }, } changed = cfg._ensure_thinking_for_clear_thinking_context_management( @@ -2111,7 +2162,9 @@ def test_bedrock_invoke_transform_emits_adaptive_thinking_for_opus_4_8(): optional_params = { "max_tokens": 32000, "stream": False, - "context_management": {"edits": [{"type": "clear_thinking_20251015", "keep": "all"}]}, + "context_management": { + "edits": [{"type": "clear_thinking_20251015", "keep": "all"}] + }, } result = cfg.transform_anthropic_messages_request( @@ -2148,7 +2201,9 @@ def test_bedrock_invoke_transform_normalizes_system_role_message_into_system(): assert all(m.get("role") != "system" for m in result["messages"]) assert result["messages"] == [{"role": "user", "content": "hi"}] - assert result["system"] == [{"type": "text", "text": "You are a careful assistant."}] + assert result["system"] == [ + {"type": "text", "text": "You are a careful assistant."} + ] def test_bedrock_invoke_transform_merges_system_role_into_existing_system(): @@ -2263,7 +2318,9 @@ def test_bedrock_invoke_transform_keeps_mid_conversation_system_role_in_place(lo ) assert result["messages"] == messages - assert result["system"] == [{"type": "text", "text": "Base.", "cache_control": {"type": "ephemeral"}}] + assert result["system"] == [ + {"type": "text", "text": "Base.", "cache_control": {"type": "ephemeral"}} + ] def test_bedrock_invoke_transform_hoists_only_leading_system_run(local_model_cost_map): @@ -2446,13 +2503,13 @@ def test_bedrock_invoke_transform_converted_system_carries_only_its_content(loca assert result["messages"][2] == { "role": "user", "content": [ - { - "type": "text", - "text": ( - "Operator note (not from the user): the following was " - "originally a mid-conversation system-role reminder." - ), - }, + { + "type": "text", + "text": ( + "Operator note (not from the user): the following was " + "originally a mid-conversation system-role reminder." + ), + }, {"type": "text", "text": "[Truncated: PARTIAL view of big1.txt]"}, ], } @@ -2588,7 +2645,10 @@ def test_as_system_content_blocks_handles_each_shape(): def test_effort_from_thinking_budget_tiers(budget_tokens, expected_effort): """The budget -> effort mapping pins each tier boundary so a shifted threshold is caught.""" - assert AmazonAnthropicClaudeMessagesConfig._effort_from_thinking_budget(budget_tokens) == expected_effort + assert ( + AmazonAnthropicClaudeMessagesConfig._effort_from_thinking_budget(budget_tokens) + == expected_effort + ) def test_inject_adaptive_thinking_preserves_existing_effort(): @@ -2597,7 +2657,9 @@ def test_inject_adaptive_thinking_preserves_existing_effort(): cfg = AmazonAnthropicClaudeMessagesConfig() request = {"output_config": {"effort": "max", "other": "keep"}} - cfg._inject_adaptive_thinking_for_clear_thinking(request, budget_tokens=24000, model="us.anthropic.claude-fable-5") + cfg._inject_adaptive_thinking_for_clear_thinking( + request, budget_tokens=24000, model="us.anthropic.claude-fable-5" + ) assert request["thinking"] == {"type": "adaptive"} assert request["output_config"] == {"effort": "max", "other": "keep"} @@ -2610,7 +2672,9 @@ def test_bedrock_clear_thinking_noops_when_thinking_already_adaptive(): request = { "max_tokens": 32000, "thinking": {"type": "adaptive"}, - "context_management": {"edits": [{"type": "clear_thinking_20251015", "keep": "all"}]}, + "context_management": { + "edits": [{"type": "clear_thinking_20251015", "keep": "all"}] + }, } changed = cfg._ensure_thinking_for_clear_thinking_context_management( @@ -2630,7 +2694,9 @@ def test_bedrock_clear_thinking_replaces_disabled_thinking_on_adaptive_model(): request = { "max_tokens": 32000, "thinking": {"type": "disabled"}, - "context_management": {"edits": [{"type": "clear_thinking_20251015", "keep": "all"}]}, + "context_management": { + "edits": [{"type": "clear_thinking_20251015", "keep": "all"}] + }, } changed = cfg._ensure_thinking_for_clear_thinking_context_management( @@ -2650,7 +2716,9 @@ def test_bedrock_clear_thinking_leaves_enabled_thinking_on_non_adaptive_model(): request = { "max_tokens": 32000, "thinking": {"type": "enabled", "budget_tokens": 8000}, - "context_management": {"edits": [{"type": "clear_thinking_20251015", "keep": "all"}]}, + "context_management": { + "edits": [{"type": "clear_thinking_20251015", "keep": "all"}] + }, } changed = cfg._ensure_thinking_for_clear_thinking_context_management( @@ -2685,7 +2753,9 @@ def test_bedrock_messages_preserves_clear_tool_uses_context_management_and_adds_ messages = [{"role": "user", "content": [{"type": "text", "text": "Hi"}]}] optional_params = { "max_tokens": 4096, - "context_management": {"edits": [{"type": "clear_tool_uses_20250919"}]}, + "context_management": { + "edits": [{"type": "clear_tool_uses_20250919"}] + }, } result = cfg.transform_anthropic_messages_request( @@ -2696,11 +2766,12 @@ def test_bedrock_messages_preserves_clear_tool_uses_context_management_and_adds_ headers={}, ) - assert result.get("context_management") == {"edits": [{"type": "clear_tool_uses_20250919"}]}, ( - "clear_tool_uses_20250919 edit must reach Bedrock InvokeModel body" - ) + assert result.get("context_management") == { + "edits": [{"type": "clear_tool_uses_20250919"}] + }, "clear_tool_uses_20250919 edit must reach Bedrock InvokeModel body" assert "context-management-2025-06-27" in result.get("anthropic_beta", []), ( - "context-management-2025-06-27 beta must reach the InvokeModel body so the tool-call-clearing edit is accepted" + "context-management-2025-06-27 beta must reach the InvokeModel body so " + "the tool-call-clearing edit is accepted" ) @@ -2777,9 +2848,9 @@ def test_bedrock_messages_filters_clear_thinking_keeps_clear_tool_uses( cm = result.get("context_management") assert cm is not None - assert [e.get("type") for e in cm["edits"]] == ["clear_tool_uses_20250919"], ( - "clear_thinking_20251015 must still be stripped (LiteLLM-internal)" - ) + assert [e.get("type") for e in cm["edits"]] == [ + "clear_tool_uses_20250919" + ], "clear_thinking_20251015 must still be stripped (LiteLLM-internal)" betas = result.get("anthropic_beta", []) assert "context-management-2025-06-27" in betas @@ -2920,7 +2991,9 @@ def test_bedrock_messages_tool_search_follows_claude_tool_search_rule(local_mode assert cfg._supports_tool_search_on_bedrock(model) is expected -def test_bedrock_messages_thinking_shape_follows_exact_bedrock_entry_flag(local_model_cost_map, monkeypatch): +def test_bedrock_messages_thinking_shape_follows_exact_bedrock_entry_flag( + local_model_cost_map, monkeypatch +): """The outbound thinking payload must follow the exact Bedrock cost-map entry. Before threading the caller's provider through the capability probes, the probe was pinned to ``"anthropic"``: the exact ``global.anthropic.claude-opus-4-8`` @@ -3330,14 +3403,22 @@ def test_bedrock_invoke_eager_input_streaming_beta_not_duplicated_with_client_he def _bedrock_event_frame(payload: dict) -> bytes: def _header(name: str, value: str) -> bytes: - return bytes([len(name)]) + name.encode() + bytes([7]) + struct.pack(">H", len(value)) + value.encode() + return ( + bytes([len(name)]) + + name.encode() + + bytes([7]) + + struct.pack(">H", len(value)) + + value.encode() + ) headers: Final = ( _header(":message-type", "event") + _header(":event-type", "chunk") + _header(":content-type", "application/json") ) - body: Final = json.dumps({"bytes": base64.b64encode(json.dumps(payload).encode()).decode()}).encode() + body: Final = json.dumps( + {"bytes": base64.b64encode(json.dumps(payload).encode()).decode()} + ).encode() prelude: Final = struct.pack(">II", 12 + len(headers) + len(body) + 4, len(headers)) prelude_crc: Final = struct.pack(">I", zlib.crc32(prelude)) message_crc: Final = struct.pack(">I", zlib.crc32(prelude + prelude_crc + headers + body)) From acadf8f6fd33e1fa9d7f9d933d3a7bf3bde2e658 Mon Sep 17 00:00:00 2001 From: mateo Date: Tue, 22 Sep 2026 23:08:40 +0000 Subject: [PATCH 4/7] test(bedrock): collect streamed chunks into a tuple in passthrough regression test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_anthropic_claude3_transformation.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index c1db6e25af75..71e2f59457b5 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -7,7 +7,8 @@ import zlib from datetime import datetime from types import SimpleNamespace -from typing import AsyncIterator, Final +from collections.abc import AsyncIterator +from typing import Final from unittest.mock import Mock import httpx @@ -3490,8 +3491,6 @@ async def test_get_async_streaming_response_iterator_yields_small_frame_before_u assert first.startswith(b"event: message_start\n"), first gate.set() - remaining: list[bytes] = [] - async for chunk in iterator: - remaining.append(chunk) + remaining: Final = tuple([chunk async for chunk in iterator]) assert any(chunk.startswith(b"event: message_stop\n") for chunk in remaining), remaining await iterator.aclose() From 61f29258d67a09b559d289bdaa3dbcea44da7d9e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 22 Sep 2026 18:40:09 -0700 Subject: [PATCH 5/7] test(bedrock): give the passthrough regression test a 10s first-chunk budget --- .../test_anthropic_claude3_transformation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index 71e2f59457b5..2bc398405426 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -3487,7 +3487,7 @@ async def test_get_async_streaming_response_iterator_yields_small_frame_before_u ), ) - first: Final = await asyncio.wait_for(anext(iterator), timeout=1) + first: Final = await asyncio.wait_for(anext(iterator), timeout=10) assert first.startswith(b"event: message_start\n"), first gate.set() From b503762289db16644e592114ba14dd6ca1dd3ac0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 22 Sep 2026 19:03:04 -0700 Subject: [PATCH 6/7] test(bedrock): type the eventstream frame helper's payload as Mapping[str, object] --- .../test_anthropic_claude3_transformation.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index 2bc398405426..208ffaa00734 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -7,7 +7,7 @@ import zlib from datetime import datetime from types import SimpleNamespace -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Mapping from typing import Final from unittest.mock import Mock @@ -3402,7 +3402,7 @@ def test_bedrock_invoke_eager_input_streaming_beta_not_duplicated_with_client_he assert result["anthropic_beta"] == [FINE_GRAINED_TOOL_STREAMING_BETA] -def _bedrock_event_frame(payload: dict) -> bytes: +def _bedrock_event_frame(payload: Mapping[str, object]) -> bytes: def _header(name: str, value: str) -> bytes: return ( bytes([len(name)]) From 67c156f2134578de5bf447b453cee83fce403414 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 22 Sep 2026 19:16:17 -0700 Subject: [PATCH 7/7] test(bedrock): take the gated byte stream's chunks as an immutable Sequence --- .../test_anthropic_claude3_transformation.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index 208ffaa00734..f4d51d975bb5 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -7,7 +7,7 @@ import zlib from datetime import datetime from types import SimpleNamespace -from collections.abc import AsyncIterator, Mapping +from collections.abc import AsyncIterator, Mapping, Sequence from typing import Final from unittest.mock import Mock @@ -3427,7 +3427,7 @@ def _header(name: str, value: str) -> bytes: class _GatedAsyncByteStream(httpx.AsyncByteStream): - def __init__(self, chunks: list[bytes], gate: asyncio.Event) -> None: + def __init__(self, chunks: Sequence[bytes], gate: asyncio.Event) -> None: self._chunks = chunks self._gate = gate @@ -3447,7 +3447,7 @@ async def test_get_async_streaming_response_iterator_yields_small_frame_before_u response: Final = httpx.Response( 200, stream=_GatedAsyncByteStream( - chunks=[ + chunks=( _bedrock_event_frame( { "type": "message_start", @@ -3467,7 +3467,7 @@ async def test_get_async_streaming_response_iterator_yields_small_frame_before_u "usage": {"input_tokens": 3, "output_tokens": 9}, } ), - ], + ), gate=gate, ), )