fix(bedrock): degrade to empty args on malformed streamed tool JSON#2230
Open
anxkhn wants to merge 1 commit into
Open
fix(bedrock): degrade to empty args on malformed streamed tool JSON#2230anxkhn wants to merge 1 commit into
anxkhn wants to merge 1 commit into
Conversation
In the Bedrock streaming path, tool-call arguments are accumulated across
contentBlockDelta chunks and parsed with json.loads, but without the
JSONDecodeError guard that every sibling provider uses (OpenAI streaming and
non-streaming, SAP AI Core streaming). A truncated or malformed streamed
tool-argument JSON therefore raises inside the generator and aborts the whole
turn as an API_ERROR, instead of degrading to empty args like the other
providers and the Bedrock non-streaming path.
Wrap the json.loads call in try/except json.JSONDecodeError and fall back to
{}, matching the existing behavior across the models package. Add a regression
test that streams a truncated toolUse input and asserts the turn completes with
empty args rather than an error.
Signed-off-by: Anas Khan <83116240+anxkhn@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Pull request overview
This PR aligns the Bedrock streaming tool-call parsing behavior with the other model providers by preventing malformed/truncated streamed tool-argument JSON from aborting an entire turn. Instead, malformed tool args degrade to {} so the turn can complete successfully.
Changes:
- Guard Bedrock streaming
json.loads(...)withtry/except json.JSONDecodeErrorand fall back to{}. - Add a regression test that simulates a truncated streamed tool
inputand asserts the turn completes with empty args.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| python/packages/kagent-adk/src/kagent/adk/models/_bedrock.py | Wraps streamed tool-argument JSON parsing to degrade to {} on JSONDecodeError instead of failing the turn. |
| python/packages/kagent-adk/tests/unittests/models/test_bedrock.py | Adds a streaming regression test for malformed tool input; includes a small cleanup opportunity (nit) around unused to_thread patching. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+324
to
+336
| async def fake_to_thread(fn, **kwargs): | ||
| return fn(**kwargs) | ||
|
|
||
| request = mock.MagicMock() | ||
| request.model = "us.anthropic.claude-sonnet-4-20250514-v1:0" | ||
| request.contents = [] | ||
| request.config = None | ||
|
|
||
| with ( | ||
| mock.patch("kagent.adk.models._bedrock._get_bedrock_client", return_value=mock_client), | ||
| mock.patch("kagent.adk.models._bedrock.asyncio.to_thread", side_effect=fake_to_thread), | ||
| ): | ||
| responses = [r async for r in llm.generate_content_async(request, stream=True)] |
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
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.
What
In the Bedrock streaming path, tool-call arguments are accumulated across
contentBlockDeltachunks and then parsed withjson.loads, but without thejson.JSONDecodeErrorguard that every other provider uses. When a streamedtool-argument JSON is truncated or malformed, the
json.loadsraises inside theresponse generator, gets caught by the broad
except Exceptionthat wraps thewhole method, and the entire turn is returned as an
API_ERROR.This PR wraps that single
json.loadsintry/except json.JSONDecodeErrorandfalls back to
{}, so a recoverable-but-malformed streamed tool call degrades toempty args instead of aborting the turn.
Why
Every sibling provider already handles this exact situation gracefully, so
Bedrock streaming is the odd one out:
_openai.py(try/except json.JSONDecodeError: args = {})_openai.py(same guard on the accumulated tool call)_sap_ai_core.py(same guard)_bedrock.pyreads an already-parsed dict(
tool.get("input", {})), so it never raises here.Only the Bedrock streaming branch parsed the accumulated
input_jsonwith a barejson.loads. Streamed tool arguments are a realistic place for truncation (earlystream termination, or a model emitting slightly malformed JSON), so the
inconsistency is user-visible: the same partial tool call that recovers on
OpenAI/SAP fails the whole turn on Bedrock streaming. This change restores the
behavior the rest of the models package already standardized on.
How
python/packages/kagent-adk/src/kagent/adk/models/_bedrock.py:The
try/exceptis copied verbatim from the sibling providers so the behavior isidentical across the package.
Testing
Added a colocated regression test,
test_streaming_malformed_tool_input_degrades_to_empty_args, inpython/packages/kagent-adk/tests/unittests/models/test_bedrock.py. It drives thereal
generate_content_async(..., stream=True)path with a mockedconverse_streamthat emits atoolUsewhose accumulated input is truncated(
{"city": "Paris"), and asserts the turn completes cleanly:error_code is None, the function-callname/idare preserved, andargs == {}. Bedrock (AWS) is fully mocked, so no live credentials are needed.API_ERROR("Expecting ',' delimiter").
No CRD or Go API change, so no
make generateand no E2E changes are required.