From 53c6939ad406b501d75d2f3cb13f0e1fc2ca4c5d Mon Sep 17 00:00:00 2001 From: zoveki <349028329+zoveki@users.noreply.github.com> Date: Wed, 28 Oct 2026 04:24:50 -0700 Subject: [PATCH] test(anthropic, bedrock): audit #47489 (#48782) * test(anthropic, bedrock): audit #47489 thinking replay across Claude routes and Bedrock Converse * test(bedrock): compare the whole replayed Converse turn in the merged-message and cache point cells --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- ...bedrock_converse_redacted_thinking_wire.py | 91 +++++++- ...est_claude_thinking_content_replay_wire.py | 195 ++++++++++++++++++ ...k_converse_redacted_thinking_router_sdk.py | 39 ++++ 3 files changed, 324 insertions(+), 1 deletion(-) create mode 100644 tests/integration/providers/test_claude_thinking_content_replay_wire.py diff --git a/tests/integration/providers/test_bedrock_converse_redacted_thinking_wire.py b/tests/integration/providers/test_bedrock_converse_redacted_thinking_wire.py index 79ec0545c44..efb0609d560 100644 --- a/tests/integration/providers/test_bedrock_converse_redacted_thinking_wire.py +++ b/tests/integration/providers/test_bedrock_converse_redacted_thinking_wire.py @@ -6,7 +6,7 @@ from typing import Final from urllib.parse import unquote import pytest -from integration._support.client import Gateway +from integration._support.client import Gateway, Scenario from integration._support.upstream import _aws_event_frame from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import BaseModel, ConfigDict, JsonValue @@ -414,3 +414,92 @@ def test_converse_redacted_thinking_replays_to_the_anthropic_api_as_claude_nativ {"type": "redacted_thinking", "data": NATIVE_REDACTED}, {"type": "text", "text": FIRST_ANSWER}, ) + + +_SIGNED_THINKING: Final[dict[str, JsonValue]] = { + "type": "thinking", + "thinking": "Checking the forecast.", + "signature": "scripted-signature-" + "s" * 32, +} +_CONVERSE_REDACTED_PART: Final[dict[str, JsonValue]] = {"reasoningContent": {"redactedContent": CONVERSE_REDACTED}} +_CONVERSE_SIGNED_PART: Final[dict[str, JsonValue]] = { + "reasoningContent": { + "reasoningText": {"text": _SIGNED_THINKING["thinking"], "signature": _SIGNED_THINKING["signature"]} + } +} + + +def _replayed_to_converse( + gateway: Gateway, wire: Wire, model: str, assistants: list[JsonValue], **fields: JsonValue +) -> tuple[dict[str, JsonValue], ...]: + reply: Final = _chat( + gateway, + { + "model": model, + "thinking": THINKING, + "max_tokens": 2048, + "messages": [{"role": "user", "content": QUESTION}, *assistants, {"role": "user", "content": FOLLOW_UP}], + **fields, + }, + ) + assert reply.content == FINAL_ANSWER + return _replayed_assistant_turn(wire).content + + +def _claude_on_converse(scenario: Scenario, wire: Wire) -> str: + return scenario.model( + model=f"bedrock/converse/{CLAUDE_MODEL_ID}", + api_key=TOKEN, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint=wire.url, + ) + + +@pytest.mark.parametrize( + "thinking_blocks", [None, [_redacted(NATIVE_REDACTED)]], ids=["content", "content-and-thinking-blocks"] +) +def test_claude_redacted_thinking_replayed_in_content_reaches_converse_once( + gateway: Gateway, thinking_blocks: list[JsonValue] | None +) -> None: + assistant: Final[dict[str, JsonValue]] = { + "role": "assistant", + "content": [_redacted(NATIVE_REDACTED), {"type": "text", "text": FIRST_ANSWER}], + **({} if thinking_blocks is None else {"thinking_blocks": thinking_blocks}), + } + with wire_server(_bedrock_peer()) as wire, gateway.scenario() as scenario: + sent: Final = _replayed_to_converse(gateway, wire, _claude_on_converse(scenario, wire), [assistant]) + assert sent == (_CONVERSE_REDACTED_PART, {"text": FIRST_ANSWER}) + + +@pytest.mark.parametrize( + "block, converse_part", + [(_redacted(NATIVE_REDACTED), _CONVERSE_REDACTED_PART), (_SIGNED_THINKING, _CONVERSE_SIGNED_PART)], + ids=["redacted", "signed"], +) +def test_block_split_across_merged_assistant_messages_reaches_converse_once( + gateway: Gateway, block: dict[str, JsonValue], converse_part: dict[str, JsonValue] +) -> None: + assistants: Final[list[JsonValue]] = [ + {"role": "assistant", "content": [block, {"type": "text", "text": FIRST_ANSWER}]}, + {"role": "assistant", "content": "Checking again.", "thinking_blocks": [block]}, + ] + with wire_server(_bedrock_peer()) as wire, gateway.scenario() as scenario: + sent: Final = _replayed_to_converse(gateway, wire, _claude_on_converse(scenario, wire), assistants) + assert sent == (converse_part, {"text": FIRST_ANSWER}, {"text": "Checking again."}) + + +@pytest.mark.parametrize("drop", [True, False], ids=["drop-params", "no-drop-params"]) +def test_converse_cache_point_after_replayed_thinking_follows_drop_params(gateway: Gateway, drop: bool) -> None: + assistant: Final[dict[str, JsonValue]] = { + "role": "assistant", + "content": [ + {**_redacted(NATIVE_REDACTED), "cache_control": {"type": "ephemeral"}}, + {"type": "text", "text": FIRST_ANSWER}, + ], + } + with wire_server(_bedrock_peer()) as wire, gateway.scenario() as scenario: + sent: Final = _replayed_to_converse( + gateway, wire, _claude_on_converse(scenario, wire), [assistant], drop_params=drop + ) + replayed: Final = (_CONVERSE_REDACTED_PART, {"text": FIRST_ANSWER}) + assert sent == (replayed if drop else (*replayed, {"cachePoint": {"type": "default"}})) diff --git a/tests/integration/providers/test_claude_thinking_content_replay_wire.py b/tests/integration/providers/test_claude_thinking_content_replay_wire.py new file mode 100644 index 00000000000..4f2f3cc66c4 --- /dev/null +++ b/tests/integration/providers/test_claude_thinking_content_replay_wire.py @@ -0,0 +1,195 @@ +import uuid +from collections.abc import Mapping +from typing import Final, Literal + +import pytest +from integration._support.anthropic_thinking import ( + BEDROCK_MODEL, + MODEL, + NO_CACHE, + SIGNATURE, + THINKING, + answer, + marker_of, + prompt, + standard_peer, + thinking_block, +) +from integration._support.client import Gateway, Scenario +from integration._support.vertex import service_account_json +from integration._support.wire import Wire, wire_server +from pydantic import BaseModel, ConfigDict, JsonValue + +Route = Literal["anthropic", "bedrock_invoke", "claude_platform", "vertex_ai", "azure_ai"] +Endpoint = Literal["chat", "messages"] + +_VERTEX_PROJECT: Final = "scripted-project" +_VERTEX_LOCATION: Final = "us-east5" +_ENABLED: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": 1024} +_EPHEMERAL: Final[dict[str, JsonValue]] = {"type": "ephemeral"} +_SIGNED: Final = thinking_block(THINKING, SIGNATURE) +_REDACTED: Final[dict[str, JsonValue]] = {"type": "redacted_thinking", "data": "scripted-redacted-" + "r" * 32} +_BLANK_REDACTED: Final[dict[str, JsonValue]] = {"type": "redacted_thinking", "data": " "} +_ROUTES: Final[tuple[Route, ...]] = ("anthropic", "bedrock_invoke", "claude_platform", "vertex_ai", "azure_ai") + + +class _Frozen(BaseModel): + model_config = ConfigDict(frozen=True) + + +class _SentMessage(_Frozen): + role: str + content: str | tuple[dict[str, JsonValue], ...] + + +class _SentRequest(_Frozen): + messages: tuple[_SentMessage, ...] + + +def _deployment(scenario: Scenario, route: Route, wire_url: str, upstream_url: str) -> str: + match route: + case "anthropic": + return scenario.model(model=f"anthropic/{MODEL}", api_base=wire_url, api_key="scripted-key") + case "bedrock_invoke": + return scenario.model( + model=f"bedrock/invoke/{BEDROCK_MODEL}", + api_base=wire_url, + aws_access_key_id="AKIASCRIPTEDPROVIDER", + aws_secret_access_key="scripted-secret", + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint=wire_url, + ) + case "claude_platform": + return scenario.model( + model=f"bedrock/claude_platform/{MODEL}", + api_base=wire_url, + api_key="scripted-platform-key", + aws_region_name="us-east-1", + workspace_id="scripted-workspace", + ) + case "vertex_ai": + return scenario.model( + model=f"vertex_ai/{MODEL}", + api_base=( + f"{wire_url}/v1/projects/{_VERTEX_PROJECT}/locations/{_VERTEX_LOCATION}" + f"/publishers/anthropic/models/{MODEL}" + ), + api_key=None, + vertex_project=_VERTEX_PROJECT, + vertex_location=_VERTEX_LOCATION, + vertex_credentials=service_account_json(_VERTEX_PROJECT, upstream_url.rstrip("/")), + ) + case "azure_ai": + return scenario.model(model=f"azure_ai/{MODEL}", api_base=wire_url, api_key="scripted-azure-key") + + +def _text(text: str) -> dict[str, JsonValue]: + return {"type": "text", "text": text} + + +def _replayed_turn( + gateway: Gateway, + wire: Wire, + *, + endpoint: Endpoint = "chat", + model: str, + marker: str, + assistant: Mapping[str, JsonValue], + **fields: JsonValue, +) -> tuple[dict[str, JsonValue], ...]: + follow_up: Final = uuid.uuid4().hex + body: Final[dict[str, JsonValue]] = { + "model": model, + "max_tokens": 2048, + "thinking": _ENABLED, + "messages": [ + {"role": "user", "content": prompt(marker)}, + {"role": "assistant", **assistant}, + {"role": "user", "content": prompt(follow_up)}, + ], + **NO_CACHE, + **fields, + } + path: Final = "/v1/chat/completions" if endpoint == "chat" else "/v1/messages" + response: Final = gateway.request("POST", path, body) + assert response.status_code == 200, response.text + assert answer(follow_up) in response.text, response.text + received: Final = wire.drain() + assert len(received) == 1, [request.body for request in received] + assert marker_of(received[0]) == follow_up, received[0].body + sent: Final = _SentRequest.model_validate_json(received[0].body).messages[1] + assert sent.role == "assistant", received[0].body + assert not isinstance(sent.content, str), received[0].body + return sent.content + + +@pytest.mark.parametrize("route", _ROUTES) +def test_redacted_thinking_replayed_in_assistant_content_reaches_claude_once(gateway: Gateway, route: Route) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(standard_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, route, wire.url, gateway.upstream_url) + sent: Final = _replayed_turn( + gateway, wire, model=model, marker=marker, assistant={"content": [_REDACTED, _text(answer(marker))]} + ) + assert sent == (_REDACTED, _text(answer(marker))) + + +def test_block_in_content_and_thinking_blocks_reaches_claude_once_in_order(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(standard_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", wire.url, gateway.upstream_url) + sent: Final = _replayed_turn( + gateway, + wire, + model=model, + marker=marker, + assistant={"content": [_REDACTED, _text(answer(marker))], "thinking_blocks": [_SIGNED, _REDACTED]}, + ) + assert sent == (_SIGNED, _REDACTED, _text(answer(marker))) + + +def test_streamed_thinking_replayed_in_content_reaches_claude_without_its_index(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(standard_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", wire.url, gateway.upstream_url) + sent: Final = _replayed_turn( + gateway, + wire, + model=model, + marker=marker, + assistant={"content": [{**_SIGNED, "index": 0}, _text(answer(marker))]}, + ) + assert sent == (_SIGNED, _text(answer(marker))) + + +@pytest.mark.parametrize("drop", [True, False], ids=["drop-params", "no-drop-params"]) +def test_replayed_thinking_cache_control_and_blank_data_follow_drop_params(gateway: Gateway, drop: bool) -> None: + marker: Final = uuid.uuid4().hex + content: Final = ({**_SIGNED, "cache_control": _EPHEMERAL}, _BLANK_REDACTED, _REDACTED, _text(answer(marker))) + with wire_server(standard_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", wire.url, gateway.upstream_url) + sent: Final = _replayed_turn( + gateway, wire, model=model, marker=marker, assistant={"content": list(content)}, drop_params=drop + ) + assert sent == ((_SIGNED, _REDACTED, _text(answer(marker))) if drop else content) + + +@pytest.mark.parametrize("endpoint", ["chat", "messages"]) +def test_assistant_cache_injection_point_never_marks_a_thinking_part(gateway: Gateway, endpoint: Endpoint) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(standard_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"anthropic/{MODEL}", + api_base=wire.url, + api_key="scripted-key", + cache_control_injection_points=[{"location": "message", "role": "assistant"}], + ) + sent: Final = _replayed_turn( + gateway, + wire, + endpoint=endpoint, + model=model, + marker=marker, + assistant={"content": [_text(answer(marker)), _SIGNED]}, + ) + assert sent == ({**_text(answer(marker)), "cache_control": _EPHEMERAL}, _SIGNED) diff --git a/tests/integration/sdk/test_bedrock_converse_redacted_thinking_router_sdk.py b/tests/integration/sdk/test_bedrock_converse_redacted_thinking_router_sdk.py index 6b2f4cb8eea..2988f9b5f53 100644 --- a/tests/integration/sdk/test_bedrock_converse_redacted_thinking_router_sdk.py +++ b/tests/integration/sdk/test_bedrock_converse_redacted_thinking_router_sdk.py @@ -88,3 +88,42 @@ def test_sync_router_reads_base_model_to_replay_redacted_thinking_to_an_applicat assistant: Final = next(message for message in received[-1].messages if message.role == "assistant") redacted: Final = ({"reasoningContent": {"redactedContent": REDACTED}},) if replays_redacted else () assert assistant.content == (*redacted, {"text": FIRST_ANSWER}) + + +def test_sync_router_replays_a_block_in_content_and_thinking_blocks_once() -> None: + redacted: Final[dict[str, JsonValue]] = {"type": "redacted_thinking", "data": REDACTED} + with wire_server(_respond) as wire: + router: Final = Router( + model_list=[ + { + "model_name": "claude-profile", + "litellm_params": { + "model": f"bedrock/converse/{APPLICATION_PROFILE_ARN}", + "api_base": wire.url, + "api_key": "synthetic-bedrock-bearer", + "aws_region_name": "us-east-1", + }, + "model_info": {"base_model": "bedrock/us.anthropic.claude-sonnet-4-6"}, + } + ], + num_retries=0, + ) + + response: Final = router.completion( + model="claude-profile", + max_tokens=256, + messages=[ + {"role": "user", "content": "What's the weather in Paris?"}, + { + "role": "assistant", + "content": [redacted, {"type": "text", "text": FIRST_ANSWER}], + "thinking_blocks": [redacted], + }, + {"role": "user", "content": "Reply with the word ok."}, + ], + ) + + assert response.choices[0].message.content == FINAL_ANSWER + received: Final = tuple(_ConverseRequest.model_validate_json(request.body) for request in wire.drain()) + assistant: Final = next(message for message in received[-1].messages if message.role == "assistant") + assert assistant.content == ({"reasoningContent": {"redactedContent": REDACTED}}, {"text": FIRST_ANSWER})