fix(bedrock): parse streamed Converse tool input like the non-streamed output

py/src/braintrust/integrations/bedrock_runtime/test_bedrock_runtime.py

@@ -289,6 +289,14 @@ def _converse_tool_kwargs(model_id=CONVERSE_MODEL):
     }
 
 
+def test_parse_streamed_tool_input_keeps_input_the_decoder_rejects():
+    from braintrust.integrations.bedrock_runtime.tracing import _parse_streamed_tool_input
+
+    assert _parse_streamed_tool_input('{"city": "Par') == '{"city": "Par'
+    deeply_nested = "[" * 100_000 + "]" * 100_000
+    assert _parse_streamed_tool_input(deeply_nested) == deeply_nested
+
+
 @pytest.mark.vcr
 def test_wrap_bedrock_converse_normalizes_tool_config_and_output(memory_logger):
     assert not memory_logger.pop()
@@ -331,6 +339,28 @@ def test_wrap_bedrock_converse_normalizes_tool_config_and_output(memory_logger):
     assert isinstance(tool_use.get("input"), dict)
 
 
+@pytest.mark.vcr
+def test_wrap_bedrock_converse_stream_parses_tool_input(memory_logger):
+    assert not memory_logger.pop()
+    client = wrap_bedrock(_bedrock_client())
+
+    response = client.converse_stream(**_converse_tool_kwargs(CONVERSE_STREAM_MODEL))
+    tool_input = ""
+    for event in response["stream"]:
+        delta = event.get("contentBlockDelta", {}).get("delta", {})
+        if "toolUse" in delta:
+            tool_input += delta["toolUse"]["input"]
+
+    spans = memory_logger.pop()
+    assert len(spans) == 1
+    content = spans[0]["output"][0]["content"]
+    tool_uses = [block for block in content if block.get("type") == "tool_use"]
+    assert tool_uses, f"expected a tool_use block in {content!r}"
+    # Same shape as the non-streamed converse output: the streamed JSON string is parsed.
+    assert tool_uses[0]["input"] == json.loads(tool_input)
+    assert isinstance(tool_uses[0]["input"], dict)
+
+
 @pytest.mark.vcr
 def test_wrap_bedrock_converse_with_image_input_materializes_attachment(memory_logger):
     assert not memory_logger.pop()

py/src/braintrust/integrations/bedrock_runtime/tracing.py

@@ -560,6 +560,18 @@ def _anthropic_invoke_usage_metrics(model_id: Any, output: Any) -> dict[str, Any
     return metrics
 
 
+def _parse_streamed_tool_input(text: str) -> Any:
+    """Parse streamed tool input into the JSON value the non-streamed converse output has."""
+    if not text:
+        return {}
+    try:
+        return json.loads(text)
+    except (ValueError, RecursionError):
+        # e.g. cut off by maxTokens, or nested deeper than the decoder allows:
+        # keep what was streamed
+        return text
+
+
 class _TracedConverseStream:
     """Wrap a botocore EventStream and log the final span when it drains."""
 
@@ -656,7 +668,7 @@ def _finish(self, error: BaseException | None) -> None:
             block = self._content_blocks.setdefault(idx, {})
             text = "".join(pieces)
             if block.get("type") == "tool_use":
-                block["input"] = text
+                block["input"] = _parse_streamed_tool_input(text)
             else:
                 block["text"] = text
         content = [self._content_blocks[idx] for idx in sorted(self._content_blocks)]