fix(bedrock/realtime): content-block session state lifecycle

Nova Sonic streams discrete content blocks. Reset part-local state
(delta chunks/type/item) on every contentEnd, mint a new item per
assistant contentStart, and only emit response.created for a new
response so text/tool/audio blocks cannot leak into each other
This commit is contained in:
mubashir1osmani 2026-08-07 16:20:55 -07:00
parent 7d60bf6126
commit a098330b51
2 changed files with 194 additions and 38 deletions

View file

@ -754,6 +754,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
current_response_id: str | None,
current_output_item_id: str | None,
current_conversation_id: str | None,
current_delta_type: ALL_DELTA_TYPES | None = None,
) -> tuple[
list[OpenAIRealtimeEvents],
str | None,
@ -764,14 +765,9 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
"""
Transform Bedrock contentStart event to OpenAI response events.
Args:
event: Bedrock contentStart event
current_response_id: Current response ID
current_output_item_id: Current output item ID
current_conversation_id: Current conversation ID
Returns:
Tuple of (events, response_id, output_item_id, conversation_id, delta_type)
Bedrock streams one content block at a time (TEXT, AUDIO, TOOL, …). Only
ASSISTANT blocks open an OpenAI response/item lifecycle. Non-assistant
blocks must not clobber in-flight assistant part state.
"""
content_start: Final = event["contentStart"]
role: Final = content_start.get("role")
@ -782,40 +778,37 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
current_response_id,
current_output_item_id,
current_conversation_id,
None,
current_delta_type,
)
verbose_logger.debug("Handling ASSISTANT contentStart")
# Initialize IDs if needed
is_new_response: Final = not current_response_id
if not current_response_id:
current_response_id = f"resp_{uuid.uuid4()}"
if not current_output_item_id:
current_output_item_id = f"item_{uuid.uuid4()}"
current_output_item_id = f"item_{uuid.uuid4()}"
if not current_conversation_id:
current_conversation_id = f"conv_{uuid.uuid4()}"
# Determine content type
content_type: Final = content_start.get("type", "TEXT")
current_delta_type: Final[ALL_DELTA_TYPES] = "text" if content_type == "TEXT" else "audio"
next_delta_type: Final[ALL_DELTA_TYPES] = "text" if content_type == "TEXT" else "audio"
returned_messages: Final[list[OpenAIRealtimeEvents]] = []
# Send response.created
response_created: Final = OpenAIRealtimeStreamResponseBaseObject(
type="response.created",
event_id=f"event_{uuid.uuid4()}",
response={
"object": "realtime.response",
"id": current_response_id,
"status": "in_progress",
"output": [],
"conversation_id": current_conversation_id,
},
)
returned_messages.append(response_created)
if is_new_response:
response_created: Final = OpenAIRealtimeStreamResponseBaseObject(
type="response.created",
event_id=f"event_{uuid.uuid4()}",
response={
"object": "realtime.response",
"id": current_response_id,
"status": "in_progress",
"output": [],
"conversation_id": current_conversation_id,
},
)
returned_messages.append(response_created)
# Send response.output_item.added
output_item_added: Final = OpenAIRealtimeStreamResponseOutputItemAdded(
type="response.output_item.added",
response_id=current_response_id,
@ -831,7 +824,6 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
)
returned_messages.append(output_item_added)
# Send response.content_part.added
content_part_added: Final = OpenAIRealtimeResponseContentPartAdded(
type="response.content_part.added",
content_index=0,
@ -839,7 +831,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
event_id=f"event_{uuid.uuid4()}",
item_id=current_output_item_id,
part=(
{"type": "text", "text": ""} if current_delta_type == "text" else {"type": "audio", "transcript": ""}
{"type": "text", "text": ""} if next_delta_type == "text" else {"type": "audio", "transcript": ""}
),
response_id=current_response_id,
)
@ -850,7 +842,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
current_response_id,
current_output_item_id,
current_conversation_id,
current_delta_type,
next_delta_type,
)
def transform_text_output_event(
@ -955,10 +947,10 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
verbose_logger.debug("Handling contentEnd: %s", content_end)
if not current_output_item_id or not current_response_id:
return [], current_delta_chunks
return [], None
if content_end.get("type") == "TOOL" or current_delta_type not in ("text", "audio"):
return [], current_delta_chunks
return [], None
returned_messages: Final[list[OpenAIRealtimeEvents]] = []
@ -1273,8 +1265,11 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
current_response_id,
current_output_item_id,
current_conversation_id,
current_delta_type,
)
returned_messages.extend(events)
if events:
current_delta_chunks = None
elif "textOutput" in event:
events, current_delta_chunks = self.transform_text_output_event(
@ -1299,10 +1294,11 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
current_delta_chunks,
)
returned_messages.extend(events)
current_delta_chunks = None
current_delta_type = None
current_output_item_id = None
if content_end.get("type") == "TOOL":
current_output_item_id = None # rebind-ok: tool block ends so next ASSISTANT mints fresh ids
current_response_id = None # rebind-ok: tool block ends so next ASSISTANT mints fresh ids
current_delta_type = None # rebind-ok: tool block ends so next ASSISTANT mints fresh ids
current_response_id = None
if BedrockContentEnd.model_validate(content_end).stopReason == "END_TURN":
(
done_events,
@ -1311,6 +1307,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
current_delta_type,
) = self._response_done_events(current_response_id, current_conversation_id)
returned_messages.extend(done_events)
current_delta_chunks = None
elif "toolUse" in event:
(
@ -1321,8 +1318,10 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
tool_response_id,
) = self.transform_tool_use_event(event, current_output_item_id, current_response_id)
returned_messages.extend(events)
current_output_item_id = tool_output_item_id # rebind-ok: persist minted tool item id into session
current_response_id = tool_response_id # rebind-ok: persist minted tool response id into session
current_output_item_id = tool_output_item_id
current_response_id = tool_response_id
current_delta_chunks = None
current_delta_type = None
verbose_logger.debug("Tool use event: %s (ID: %s)", tool_name, tool_call_id)
elif "promptEnd" in event or "completionEnd" in event:
@ -1333,6 +1332,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
current_delta_type,
) = self.transform_prompt_end_event(event, current_response_id, current_conversation_id)
returned_messages.extend(events)
current_delta_chunks = None
return {
"response": returned_messages,

View file

@ -1090,6 +1090,162 @@ class TestBedrockRealtimeSessionEvents:
assert event["session"]["modalities"] == ["text", "audio"]
class TestBedrockRealtimeContentBlockLifecycle:
"""
Bedrock streams discrete content blocks. Session state must follow block
boundaries so text/audio/tool blocks cannot leak into each other.
"""
def _state(self, **overrides):
base = {
"session_configuration_request": json.dumps({"configured": True}),
"current_output_item_id": None,
"current_response_id": None,
"current_conversation_id": "conv_1",
"current_delta_chunks": None,
"current_item_chunks": [],
"current_delta_type": None,
}
base.update(overrides)
return base
def _apply(self, config, logging_obj, state, message):
result = config.transform_realtime_response(
json.dumps(message),
"amazon.nova-2-sonic-v1:0",
logging_obj,
realtime_response_transform_input=state,
)
state.update(
{
"current_output_item_id": result["current_output_item_id"],
"current_response_id": result["current_response_id"],
"current_conversation_id": result["current_conversation_id"],
"current_delta_chunks": result["current_delta_chunks"],
"current_item_chunks": result["current_item_chunks"],
"current_delta_type": result["current_delta_type"],
}
)
return result
def test_tool_block_does_not_leak_prior_text_into_next_assistant_turn(self):
config = BedrockRealtimeConfig()
logging_obj = MagicMock()
logging_obj.litellm_trace_id = "trace_123"
state = self._state()
self._apply(config, logging_obj, state, {"event": {"contentStart": {"role": "ASSISTANT", "type": "TEXT"}}})
first_response_id = state["current_response_id"]
self._apply(
config,
logging_obj,
state,
{"event": {"textOutput": {"content": "I will check the weather."}}},
)
assert state["current_delta_chunks"] is not None
assert len(state["current_delta_chunks"]) == 1
self._apply(
config,
logging_obj,
state,
{"event": {"contentEnd": {"stopReason": "PARTIAL_TURN", "type": "TEXT"}}},
)
assert state["current_delta_chunks"] is None
assert state["current_delta_type"] is None
assert state["current_output_item_id"] is None
assert state["current_response_id"] == first_response_id
self._apply(config, logging_obj, state, {"event": {"contentStart": {"role": "TOOL", "type": "TOOL"}}})
assert state["current_delta_chunks"] is None
assert state["current_response_id"] == first_response_id
tool_result = self._apply(
config,
logging_obj,
state,
{
"event": {
"toolUse": {
"toolUseId": "tool_1",
"toolName": "get_weather",
"content": json.dumps({"location": "Seattle"}),
}
}
},
)
assert tool_result["response"][0]["type"] == "response.function_call_arguments.done"
assert state["current_delta_chunks"] is None
assert state["current_delta_type"] is None
self._apply(
config,
logging_obj,
state,
{"event": {"contentEnd": {"stopReason": "TOOL_USE", "type": "TOOL"}}},
)
assert state["current_response_id"] is None
assert state["current_output_item_id"] is None
assert state["current_delta_chunks"] is None
assert state["current_delta_type"] is None
post_tool = self._apply(
config,
logging_obj,
state,
{"event": {"contentStart": {"role": "ASSISTANT", "type": "TEXT"}}},
)
assert state["current_response_id"] != first_response_id
assert state["current_delta_chunks"] is None
assert [msg["type"] for msg in post_tool["response"]].count("response.created") == 1
self._apply(
config,
logging_obj,
state,
{"event": {"textOutput": {"content": "It is sunny in Seattle."}}},
)
done = self._apply(
config,
logging_obj,
state,
{"event": {"contentEnd": {"stopReason": "END_TURN", "type": "TEXT"}}},
)
text_done = [msg for msg in done["response"] if msg["type"] == "response.text.done"][0]
assert text_done["text"] == "It is sunny in Seattle."
assert "I will check the weather." not in text_done["text"]
assert any(msg["type"] == "response.done" for msg in done["response"])
assert state["current_response_id"] is None
assert state["current_delta_chunks"] is None
def test_second_assistant_content_block_reuses_response_not_item(self):
config = BedrockRealtimeConfig()
logging_obj = MagicMock()
logging_obj.litellm_trace_id = "trace_123"
state = self._state()
first = self._apply(
config, logging_obj, state, {"event": {"contentStart": {"role": "ASSISTANT", "type": "TEXT"}}}
)
response_id = state["current_response_id"]
first_item = state["current_output_item_id"]
assert sum(1 for msg in first["response"] if msg["type"] == "response.created") == 1
self._apply(
config,
logging_obj,
state,
{"event": {"contentEnd": {"stopReason": "PARTIAL_TURN", "type": "TEXT"}}},
)
second = self._apply(
config, logging_obj, state, {"event": {"contentStart": {"role": "ASSISTANT", "type": "AUDIO"}}}
)
assert state["current_response_id"] == response_id
assert state["current_output_item_id"] != first_item
assert sum(1 for msg in second["response"] if msg["type"] == "response.created") == 0
assert sum(1 for msg in second["response"] if msg["type"] == "response.output_item.added") == 1
class TestBedrockRealtimeUsageAccounting:
def _usage_event(
self,