Python: fix(bedrock): keep reasoning content so extended thinking works with tool calls

python/packages/bedrock/agent_framework_bedrock/_chat_client.py

@@ -4,12 +4,14 @@
 from __future__ import annotations
 
 import asyncio
+import base64
 import copy
 import json
 import logging
 import sys
 from collections import deque
 from collections.abc import AsyncIterable, Awaitable, Mapping, MutableMapping, Sequence
+from dataclasses import dataclass
 from typing import Any, ClassVar, Generic, Literal, TypedDict
 from uuid import uuid4
 
@@ -58,6 +60,18 @@
 
 logger = logging.getLogger("agent_framework.bedrock")
 
+# Marks text_reasoning content whose protected_data is base64 of a Bedrock ``redactedContent`` blob,
+# as opposed to the signature of a ``reasoningText`` block.
+_REDACTED_REASONING_KEY = "bedrock.redacted_reasoning"
+
+
+@dataclass
+class _StreamedReasoningBlock:
+    """State for one reasoning content block while its ConverseStream deltas arrive."""
+
+    id: str
+    redacted: bytes = b""
+
 
 __all__ = [
     "BedrockChatClient",
@@ -382,9 +396,12 @@ async def _stream() -> AsyncIterable[ChatResponseUpdate]:
                 events = response["stream"]
                 event_iterator = iter(events)
                 tool_call_ids: dict[int, str] = {}
+                reasoning_blocks: dict[int, _StreamedReasoningBlock] = {}
                 try:
                     while (event := await asyncio.to_thread(next, event_iterator, None)) is not None:
-                        if update := self._process_converse_stream_event(event, request["modelId"], tool_call_ids):
+                        if update := self._process_converse_stream_event(
+                            event, request["modelId"], tool_call_ids, reasoning_blocks
+                        ):
                             yield update
                 finally:
                     events.close()
@@ -570,7 +587,11 @@ def _convert_message_to_content_blocks(self, message: Message) -> list[dict[str,
         blocks: list[dict[str, Any]] = []
         for content in message.contents:
             block = self._convert_content_to_bedrock_block(content)
-            if block is None or ("image" in block and message.role != "user"):
+            if (
+                block is None
+                or ("image" in block and message.role != "user")
+                or ("reasoningContent" in block and message.role != "assistant")
+            ):
                 logger.debug("Skipping unsupported content type for Bedrock: %s", type(content))
                 continue
             blocks.append(block)
@@ -587,6 +608,21 @@ def _convert_content_to_bedrock_block(self, content: Content) -> dict[str, Any]
                     logger.warning("Skipping %s image: Bedrock accepts gif, jpeg, png or webp.", content.media_type)
                     return None
                 return {"image": {"format": image_format, "source": {"bytes": _get_data_bytes(content)}}}
+            case "text_reasoning":
+                # Extended thinking with tool use requires reasoning blocks to be sent back unchanged.
+                if content.additional_properties.get(_REDACTED_REASONING_KEY):
+                    try:
+                        redacted = base64.b64decode(content.protected_data or "", validate=True)
+                    except ValueError:
+                        logger.warning("Skipping redacted reasoning for Bedrock: protected_data is not valid base64")
+                        return None
+                    return {"reasoningContent": {"redactedContent": redacted}} if redacted else None
+                if content.text is None:
+                    return None
+                reasoning_text: dict[str, Any] = {"text": content.text}
+                if content.protected_data:
+                    reasoning_text["signature"] = content.protected_data
+                return {"reasoningContent": {"reasoningText": reasoning_text}}
             case "function_call":
                 arguments = content.parse_arguments() or {}
                 return {
@@ -685,6 +721,19 @@ def _prepare_tools(self, tools: list[FunctionTool | MutableMapping[str, Any]] |
     def _generate_tool_call_id() -> str:
         return f"tool-call-{uuid4().hex}"
 
+    @staticmethod
+    def _generate_reasoning_id() -> str:
+        return f"reasoning-{uuid4().hex}"
+
+    @staticmethod
+    def _redacted_reasoning_content(redacted: bytes, *, id: str, raw_representation: Any) -> Content:
+        return Content.from_text_reasoning(
+            id=id,
+            protected_data=base64.b64encode(redacted).decode("utf-8"),
+            additional_properties={_REDACTED_REASONING_KEY: True},
+            raw_representation=raw_representation,
+        )
+
     def _process_converse_response(
         self, response: dict[str, Any], options: Mapping[str, Any] | None = None
     ) -> ChatResponse:
@@ -710,7 +759,11 @@ def _process_converse_response(
         )
 
     def _process_converse_stream_event(
-        self, event: Mapping[str, Any], model: str, tool_call_ids: dict[int, str]
+        self,
+        event: Mapping[str, Any],
+        model: str,
+        tool_call_ids: dict[int, str],
+        reasoning_blocks: dict[int, _StreamedReasoningBlock],
     ) -> ChatResponseUpdate | None:
         """Convert a single Bedrock ConverseStream event to a ChatResponseUpdate."""
         contents: list[Content] = []
@@ -733,6 +786,29 @@ def _process_converse_stream_event(
                         arguments=tool_use_delta.get("input", ""),
                     )
                 )
+            elif reasoning_delta := delta.get("reasoningContent"):
+                # The deltas of one reasoning block share an id, so they merge into a single text_reasoning
+                # content in the final response while consecutive blocks keep their own text and signature.
+                block = reasoning_blocks.setdefault(
+                    block_delta.get("contentBlockIndex", 0),
+                    _StreamedReasoningBlock(id=self._generate_reasoning_id()),
+                )
+                if isinstance(redacted := reasoning_delta.get("redactedContent"), (bytes, bytearray)):
+                    # Merging keeps only the latest protected_data, so send everything received so far.
+                    block.redacted += bytes(redacted)
+                    contents.append(
+                        self._redacted_reasoning_content(block.redacted, id=block.id, raw_representation=delta)
+                    )
+                elif "text" in reasoning_delta or "signature" in reasoning_delta:
+                    contents.append(
+                        Content.from_text_reasoning(
+                            id=block.id,
+                            # "" rather than None, so a signed block without text is still replayed.
+                            text=reasoning_delta.get("text") or "",
+                            protected_data=reasoning_delta.get("signature"),
+                            raw_representation=delta,
+                        )
+                    )
         elif message_stop := event.get("messageStop"):
             finish_reason = self._map_finish_reason(message_stop.get("stopReason"))
         elif (metadata := event.get("metadata")) and (usage_details := self._parse_usage(metadata.get("usage"))):
@@ -771,6 +847,26 @@ def _parse_message_contents(self, content_blocks: Sequence[dict[str, Any]]) -> l
                     Content.from_text(text=json.dumps(json_value, ensure_ascii=False), raw_representation=block)
                 )
                 continue
+            if isinstance(reasoning := block.get("reasoningContent"), Mapping):
+                # Each block gets its own id so consecutive reasoning blocks are never merged into one.
+                if isinstance(reasoning_text := reasoning.get("reasoningText"), Mapping):
+                    contents.append(
+                        Content.from_text_reasoning(
+                            id=self._generate_reasoning_id(),
+                            text=reasoning_text.get("text") or "",
+                            protected_data=reasoning_text.get("signature"),
+                            raw_representation=block,
+                        )
+                    )
+                elif isinstance(redacted := reasoning.get("redactedContent"), (bytes, bytearray)):
+                    contents.append(
+                        self._redacted_reasoning_content(
+                            bytes(redacted), id=self._generate_reasoning_id(), raw_representation=block
+                        )
+                    )
+                else:
+                    logger.debug("Ignoring unsupported Bedrock reasoning block: %s", block)
+                continue
             tool_use_value = block.get("toolUse")
             tool_use = (
                 tool_use_value

python/packages/bedrock/tests/test_bedrock_client.py

@@ -152,6 +152,214 @@ async def test_stream_yields_updates_as_converse_stream_events_arrive() -> None:
     assert stub.stream.closed
 
 
+def test_process_converse_response_parses_reasoning_content() -> None:
+    client = _make_client()
+
+    response = client._process_converse_response({
+        "output": {
+            "message": {
+                "role": "assistant",
+                "content": [
+                    {"reasoningContent": {"reasoningText": {"text": "Need the weather tool.", "signature": "sig-1"}}},
+                    {"toolUse": {"toolUseId": "call-1", "name": "get_weather", "input": {"city": "Paris"}}},
+                ],
+            }
+        },
+        "stopReason": "tool_use",
+    })
+
+    reasoning, function_call = response.messages[0].contents
+    assert reasoning.type == "text_reasoning"
+    assert reasoning.text == "Need the weather tool."
+    assert reasoning.protected_data == "sig-1"
+    assert function_call.type == "function_call"
+
+
+async def test_stream_merges_reasoning_text_and_signature() -> None:
+    stub = _StubBedrockStreamRuntime([
+        {"contentBlockDelta": {"delta": {"reasoningContent": {"text": "Let me "}}, "contentBlockIndex": 0}},
+        {"contentBlockDelta": {"delta": {"reasoningContent": {"text": "think."}}, "contentBlockIndex": 0}},
+        {"contentBlockDelta": {"delta": {"reasoningContent": {"signature": "sig-1"}}, "contentBlockIndex": 0}},
+        {"contentBlockStop": {"contentBlockIndex": 0}},
+        {"contentBlockDelta": {"delta": {"text": "Done."}, "contentBlockIndex": 1}},
+        {"messageStop": {"stopReason": "end_turn"}},
+    ])
+    client = BedrockChatClient(
+        model="anthropic.claude-sonnet-4",
+        region="us-east-1",
+        client=stub,  # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] # pyright: ignore[reportArgumentType]
+    )
+
+    stream = client._inner_get_response(
+        messages=[Message(role="user", contents=[Content.from_text(text="hi")])], options={}, stream=True
+    )
+    assert isinstance(stream, ResponseStream)
+    async for _ in stream:
+        pass
+    response = await stream.get_final_response()
+
+    reasoning = [content for content in response.messages[0].contents if content.type == "text_reasoning"]
+    assert len(reasoning) == 1
+    assert reasoning[0].text == "Let me think."
+    assert reasoning[0].protected_data == "sig-1"
+    assert response.text == "Done."
+
+
+def test_prepare_bedrock_messages_sends_reasoning_back_with_tool_use() -> None:
+    """Extended thinking with tool use fails unless the signed reasoning is replayed before the toolUse block."""
+    client = _make_client()
+    messages = [
+        Message(role="user", contents=[Content.from_text(text="Weather in Paris?")]),
+        Message(
+            role="assistant",
+            contents=[
+                Content.from_text_reasoning(text="Need the weather tool.", protected_data="sig-1"),
+                Content.from_function_call(call_id="call-1", name="get_weather", arguments={"city": "Paris"}),
+            ],
+        ),
+        Message(role="tool", contents=[Content.from_function_result(call_id="call-1", result="88F")]),
+    ]
+
+    _, conversation = client._prepare_bedrock_messages(messages)
+
+    assert conversation[1]["content"][0] == {
+        "reasoningContent": {"reasoningText": {"text": "Need the weather tool.", "signature": "sig-1"}}
+    }
+    assert "toolUse" in conversation[1]["content"][1]
+
+
+def test_prepare_bedrock_messages_skips_reasoning_outside_assistant_messages() -> None:
+    client = _make_client()
+    messages = [
+        Message(
+            role="user",
+            contents=[Content.from_text(text="hi"), Content.from_text_reasoning(text="not from the model")],
+        )
+    ]
+
+    _, conversation = client._prepare_bedrock_messages(messages)
+
+    assert conversation[0]["content"] == [{"text": "hi"}]
+
+
+async def test_stream_keeps_consecutive_reasoning_blocks_separate() -> None:
+    """Each streamed reasoning block must keep its own text and signature when replayed."""
+    stub = _StubBedrockStreamRuntime([
+        {"contentBlockDelta": {"delta": {"reasoningContent": {"text": "First "}}, "contentBlockIndex": 0}},
+        {"contentBlockDelta": {"delta": {"reasoningContent": {"text": "thought."}}, "contentBlockIndex": 0}},
+        {"contentBlockDelta": {"delta": {"reasoningContent": {"signature": "sig-1"}}, "contentBlockIndex": 0}},
+        {"contentBlockStop": {"contentBlockIndex": 0}},
+        {"contentBlockDelta": {"delta": {"reasoningContent": {"text": "Second thought."}}, "contentBlockIndex": 1}},
+        {"contentBlockDelta": {"delta": {"reasoningContent": {"signature": "sig-2"}}, "contentBlockIndex": 1}},
+        {"contentBlockStop": {"contentBlockIndex": 1}},
+        {
+            "contentBlockStart": {
+                "start": {"toolUse": {"toolUseId": "call-1", "name": "get_weather"}},
+                "contentBlockIndex": 2,
+            }
+        },
+        {"contentBlockDelta": {"delta": {"toolUse": {"input": '{"city": "Paris"}'}}, "contentBlockIndex": 2}},
+        {"messageStop": {"stopReason": "tool_use"}},
+    ])
+    client = BedrockChatClient(
+        model="anthropic.claude-sonnet-4",
+        region="us-east-1",
+        client=stub,  # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] # pyright: ignore[reportArgumentType]
+    )
+
+    stream = client._inner_get_response(
+        messages=[Message(role="user", contents=[Content.from_text(text="Weather in Paris?")])],
+        options={},
+        stream=True,
+    )
+    assert isinstance(stream, ResponseStream)
+    async for _ in stream:
+        pass
+    response = await stream.get_final_response()
+
+    reasoning = [content for content in response.messages[0].contents if content.type == "text_reasoning"]
+    assert [(content.text, content.protected_data) for content in reasoning] == [
+        ("First thought.", "sig-1"),
+        ("Second thought.", "sig-2"),
+    ]
+
+    _, conversation = client._prepare_bedrock_messages(response.messages)
+    assert conversation[0]["content"][:2] == [
+        {"reasoningContent": {"reasoningText": {"text": "First thought.", "signature": "sig-1"}}},
+        {"reasoningContent": {"reasoningText": {"text": "Second thought.", "signature": "sig-2"}}},
+    ]
+    assert "toolUse" in conversation[0]["content"][2]
+
+
+def test_redacted_reasoning_content_is_replayed_unchanged() -> None:
+    client = _make_client()
+    redacted = b"\x00\x01encrypted-reasoning\xff"
+
+    response = client._process_converse_response({
+        "output": {
+            "message": {
+                "role": "assistant",
+                "content": [
+                    {"reasoningContent": {"redactedContent": redacted}},
+                    {"toolUse": {"toolUseId": "call-1", "name": "get_weather", "input": {"city": "Paris"}}},
+                ],
+            }
+        },
+        "stopReason": "tool_use",
+    })
+
+    reasoning = response.messages[0].contents[0]
+    assert reasoning.type == "text_reasoning"
+    assert reasoning.text is None
+
+    _, conversation = client._prepare_bedrock_messages(response.messages)
+    assert conversation[0]["content"][0] == {"reasoningContent": {"redactedContent": redacted}}
+    assert "toolUse" in conversation[0]["content"][1]
+
+
+async def test_stream_redacted_reasoning_content_is_replayed_unchanged() -> None:
+    stub = _StubBedrockStreamRuntime([
+        {
+            "contentBlockDelta": {
+                "delta": {"reasoningContent": {"redactedContent": b"\x00\x01part"}},
+                "contentBlockIndex": 0,
+            }
+        },
+        {
+            "contentBlockDelta": {
+                "delta": {"reasoningContent": {"redactedContent": b"-two\xff"}},
+                "contentBlockIndex": 0,
+            }
+        },
+        {"contentBlockStop": {"contentBlockIndex": 0}},
+        {"contentBlockDelta": {"delta": {"reasoningContent": {"text": "Visible."}}, "contentBlockIndex": 1}},
+        {"contentBlockDelta": {"delta": {"reasoningContent": {"signature": "sig-1"}}, "contentBlockIndex": 1}},
+        {"contentBlockStop": {"contentBlockIndex": 1}},
+        {"contentBlockDelta": {"delta": {"text": "Done."}, "contentBlockIndex": 2}},
+        {"messageStop": {"stopReason": "end_turn"}},
+    ])
+    client = BedrockChatClient(
+        model="anthropic.claude-3-7-sonnet",
+        region="us-east-1",
+        client=stub,  # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] # pyright: ignore[reportArgumentType]
+    )
+
+    stream = client._inner_get_response(
+        messages=[Message(role="user", contents=[Content.from_text(text="hi")])], options={}, stream=True
+    )
+    assert isinstance(stream, ResponseStream)
+    async for _ in stream:
+        pass
+    response = await stream.get_final_response()
+
+    _, conversation = client._prepare_bedrock_messages(response.messages)
+    assert conversation[0]["content"] == [
+        {"reasoningContent": {"redactedContent": b"\x00\x01part-two\xff"}},
+        {"reasoningContent": {"reasoningText": {"text": "Visible.", "signature": "sig-1"}}},
+        {"text": "Done."},
+    ]
+
+
 def test_build_request_requires_non_system_messages() -> None:
     client = BedrockChatClient(
         model="amazon.titan-text",