Skip to content

Commit 6143c8b

Browse files
committed
fix(streaming): handle tool use metadata in contentBlockDelta for non-standard models
Some Bedrock models (e.g., Kimi K2.5) send toolUseId and name in contentBlockDelta rather than in contentBlockStart. The parser only extracted input from the delta, leaving current_tool_use without those fields and crashing in handle_content_block_stop. - Capture toolUseId and name from delta when not already set by start - Use .get("input", "") to avoid KeyError when input key is absent - Add warning + skip for incomplete tool use blocks in stop handler Fixes #1646
1 parent 94fc8dd commit 6143c8b

2 files changed

Lines changed: 127 additions & 3 deletions

File tree

src/strands/event_loop/streaming.py

Lines changed: 20 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -210,10 +210,19 @@ def handle_content_block_delta(
210210
typed_event: ModelStreamEvent = ModelStreamEvent({})
211211

212212
if "toolUse" in delta_content:
213+
tool_use_delta = delta_content["toolUse"]
213214
if "input" not in state["current_tool_use"]:
214215
state["current_tool_use"]["input"] = ""
215216

216-
state["current_tool_use"]["input"] += delta_content["toolUse"]["input"]
217+
state["current_tool_use"]["input"] += tool_use_delta.get("input", "")
218+
219+
# Some models provide toolUseId and name in contentBlockDelta instead of contentBlockStart.
220+
# Capture them here if not already set from a prior contentBlockStart event.
221+
if "toolUseId" not in state["current_tool_use"] and "toolUseId" in tool_use_delta:
222+
state["current_tool_use"]["toolUseId"] = tool_use_delta["toolUseId"]
223+
if "name" not in state["current_tool_use"] and "name" in tool_use_delta:
224+
state["current_tool_use"]["name"] = tool_use_delta["name"]
225+
217226
typed_event = ToolUseStreamEvent(delta_content, state["current_tool_use"])
218227

219228
elif "text" in delta_content:
@@ -281,8 +290,16 @@ def handle_content_block_stop(state: dict[str, Any]) -> dict[str, Any]:
281290
except ValueError:
282291
current_tool_use["input"] = {}
283292

284-
tool_use_id = current_tool_use["toolUseId"]
285-
tool_use_name = current_tool_use["name"]
293+
tool_use_id = current_tool_use.get("toolUseId", "")
294+
tool_use_name = current_tool_use.get("name", "")
295+
296+
if not tool_use_id or not tool_use_name:
297+
logger.warning(
298+
"Incomplete tool use block (missing toolUseId or name); skipping content block. "
299+
"The model may be using a non-standard streaming format."
300+
)
301+
state["current_tool_use"] = {}
302+
return state
286303

287304
tool_use = ToolUse(
288305
toolUseId=tool_use_id,

tests/strands/event_loop/test_streaming.py

Lines changed: 107 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1422,3 +1422,110 @@ async def test_process_stream_keeps_tool_use_stop_reason_unchanged(agenerator, a
14221422
last_event = cast(ModelStopReason, (await alist(stream))[-1])
14231423

14241424
assert last_event["stop"][0] == "tool_use"
1425+
1426+
1427+
def test_handle_content_block_delta_captures_tool_use_id_and_name_from_delta():
1428+
"""Delta events that include toolUseId and name should populate current_tool_use."""
1429+
event = {"delta": {"toolUse": {"input": '{"x": 1}', "toolUseId": "abc123", "name": "output_slide"}}}
1430+
state = {"current_tool_use": {}}
1431+
1432+
updated_state, _ = strands.event_loop.streaming.handle_content_block_delta(event, state)
1433+
1434+
assert updated_state["current_tool_use"]["toolUseId"] == "abc123"
1435+
assert updated_state["current_tool_use"]["name"] == "output_slide"
1436+
assert updated_state["current_tool_use"]["input"] == '{"x": 1}'
1437+
1438+
1439+
def test_handle_content_block_delta_does_not_override_existing_tool_use_id_and_name():
1440+
"""toolUseId and name from contentBlockStart should not be overridden by a later delta."""
1441+
event = {"delta": {"toolUse": {"input": '{"x": 1}', "toolUseId": "from_delta", "name": "from_delta"}}}
1442+
state = {"current_tool_use": {"toolUseId": "from_start", "name": "from_start", "input": ""}}
1443+
1444+
updated_state, _ = strands.event_loop.streaming.handle_content_block_delta(event, state)
1445+
1446+
assert updated_state["current_tool_use"]["toolUseId"] == "from_start"
1447+
assert updated_state["current_tool_use"]["name"] == "from_start"
1448+
1449+
1450+
def test_handle_content_block_delta_tool_use_without_input_key():
1451+
"""A toolUse delta missing the input key should not raise KeyError."""
1452+
event = {"delta": {"toolUse": {}}}
1453+
state = {"current_tool_use": {"toolUseId": "t1", "name": "tool"}}
1454+
1455+
updated_state, _ = strands.event_loop.streaming.handle_content_block_delta(event, state)
1456+
1457+
assert updated_state["current_tool_use"]["input"] == ""
1458+
1459+
1460+
def test_handle_content_block_stop_skips_incomplete_tool_use_missing_id(caplog):
1461+
"""A tool use block missing toolUseId is skipped with a warning."""
1462+
import logging
1463+
1464+
state = {
1465+
"content": [],
1466+
"current_tool_use": {"name": "output_slide", "input": '{"x": 1}'},
1467+
"text": "",
1468+
"reasoningText": "",
1469+
"citationsContent": [],
1470+
}
1471+
1472+
with caplog.at_level(logging.WARNING, logger="strands.event_loop.streaming"):
1473+
updated_state = strands.event_loop.streaming.handle_content_block_stop(state)
1474+
1475+
assert updated_state["content"] == []
1476+
assert updated_state["current_tool_use"] == {}
1477+
assert "Incomplete tool use block" in caplog.text
1478+
1479+
1480+
def test_handle_content_block_stop_skips_incomplete_tool_use_missing_name(caplog):
1481+
"""A tool use block missing name is skipped with a warning."""
1482+
import logging
1483+
1484+
state = {
1485+
"content": [],
1486+
"current_tool_use": {"toolUseId": "abc123", "input": '{"x": 1}'},
1487+
"text": "",
1488+
"reasoningText": "",
1489+
"citationsContent": [],
1490+
}
1491+
1492+
with caplog.at_level(logging.WARNING, logger="strands.event_loop.streaming"):
1493+
updated_state = strands.event_loop.streaming.handle_content_block_stop(state)
1494+
1495+
assert updated_state["content"] == []
1496+
assert updated_state["current_tool_use"] == {}
1497+
assert "Incomplete tool use block" in caplog.text
1498+
1499+
1500+
@pytest.mark.asyncio
1501+
async def test_process_stream_tool_use_info_in_delta(agenerator, alist):
1502+
"""Models that provide toolUseId and name in contentBlockDelta (not contentBlockStart) work correctly."""
1503+
response = [
1504+
{"messageStart": {"role": "assistant"}},
1505+
{"contentBlockStart": {"start": {}}},
1506+
{
1507+
"contentBlockDelta": {
1508+
"delta": {"toolUse": {"input": '{"title": "Test"}', "toolUseId": "xyz789", "name": "output_slide"}}
1509+
}
1510+
},
1511+
{"contentBlockStop": {}},
1512+
{"messageStop": {"stopReason": "tool_use"}},
1513+
{
1514+
"metadata": {
1515+
"usage": {"inputTokens": 5, "outputTokens": 10, "totalTokens": 15},
1516+
"metrics": {"latencyMs": 50},
1517+
}
1518+
},
1519+
]
1520+
1521+
stream = strands.event_loop.streaming.process_stream(agenerator(response))
1522+
events = await alist(stream)
1523+
last_event = cast(ModelStopReason, events[-1])
1524+
1525+
stop_reason, message, _, _ = last_event["stop"]
1526+
assert stop_reason == "tool_use"
1527+
assert len(message["content"]) == 1
1528+
tool_use = message["content"][0]["toolUse"]
1529+
assert tool_use["toolUseId"] == "xyz789"
1530+
assert tool_use["name"] == "output_slide"
1531+
assert tool_use["input"] == {"title": "Test"}

0 commit comments

Comments
 (0)