scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

fix(responses): reject tool metadata mutation

Browse source
This commit is contained in:
Dan Loftus 2026-09-01 23:34:56 -04:00
parent cddb23dc5d
commit bd795bf5d9
3 changed files with 129 additions and 78 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

118
litellm/responses/litellm_completion_transformation/streaming_iterator.py
View file

@ -57,7 +57,6 @@ class _StreamedToolCallMetadata:
call_type: str | None
tool_name: str | None
tool_namespace: str | None
ambiguous: bool = False
def matches(self, incoming: "_StreamedToolCallMetadata") -> bool:
return all(
@ -73,15 +72,6 @@ class _StreamedToolCallMetadata:
call_type=self.call_type or incoming.call_type,
tool_name=self.tool_name or incoming.tool_name,
tool_namespace=self.tool_namespace or incoming.tool_namespace,
ambiguous=self.ambiguous,
)
def marked_ambiguous(self) -> "_StreamedToolCallMetadata":
return _StreamedToolCallMetadata(
call_type=self.call_type,
tool_name=self.tool_name,
tool_namespace=self.tool_namespace,
ambiguous=True,
)
@ -157,6 +147,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
self._streamed_tool_call_ids_in_order: list[str] = [] # mutable-ok: accumulates call ids across stream chunks
self._resolved_tool_call_id_by_position: dict[int, str] = {} # mutable-ok: terminal correlation state
self._streamed_call_id_by_terminal_id: dict[str, str] = {} # mutable-ok: terminal identity correlation
self._pending_tool_call_error: litellm.InternalServerError | None = None
self._next_tool_output_index: int = 1 # output_index=0 reserved for the message item
self._final_tool_events_queued: bool = False
self._sequence_number: int = 0
@ -195,9 +186,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
return None
def _streamed_tool_call_id_at_position(self, position: int) -> str | None:
indexed_metadata: Final = self._tool_call_metadata_by_index.get(position)
if indexed_metadata is not None and indexed_metadata.ambiguous:
return None
indexed_call_id: Final = self._tool_call_id_by_index.get(position)
if indexed_call_id is not None:
return indexed_call_id
@ -212,17 +200,23 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
def _streamed_tool_call_id_for_terminal_call(self, tool_call: object, position: int) -> str | None:
"""Match a terminal aggregate tool call to the identity emitted while streaming."""
tool_call_index: Final = self._normalize_tool_call_index(tool_call)
metadata_index: Final = tool_call_index if tool_call_index is not None else position
streamed_metadata: Final = self._tool_call_metadata_by_index.get(metadata_index)
terminal_metadata: Final = self._tool_call_metadata(tool_call)
if streamed_metadata is not None and not streamed_metadata.matches(terminal_metadata):
streamed_call_id: Final = self._tool_call_id_by_index.get(metadata_index)
if streamed_call_id is not None:
self._queue_tool_call_metadata_error(streamed_call_id, streamed_metadata, metadata_index)
return None
call_id_raw: Final = tool_call.get("id") if isinstance(tool_call, dict) else getattr(tool_call, "id", None)
if call_id_raw:
call_id_match: Final = self._streamed_call_id_by_terminal_id.get(str(call_id_raw))
if call_id_match is not None:
return call_id_match
tool_call_index: Final = self._normalize_tool_call_index(tool_call)
if tool_call_index is not None:
indexed_metadata: Final = self._tool_call_metadata_by_index.get(tool_call_index)
if indexed_metadata is not None and indexed_metadata.ambiguous:
return None
indexed_call_id: Final = self._tool_call_id_by_index.get(tool_call_index)
if indexed_call_id is not None:
return indexed_call_id
@ -250,25 +244,23 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
indexed_call_id: Final = self._tool_call_id_by_index.get(tool_call_index)
if indexed_call_id is not None:
if known_call_id is not None and known_call_id != indexed_call_id:
return known_call_id
indexed_metadata: Final = self._tool_call_metadata_by_index.get(
tool_call_index,
_StreamedToolCallMetadata(None, None, None),
)
if known_call_id is not None and known_call_id != indexed_call_id:
self._queue_tool_call_metadata_error(indexed_call_id, indexed_metadata, tool_call_index)
return None
if incoming_call_id is None:
return None if indexed_metadata.ambiguous else indexed_call_id
return indexed_call_id
if indexed_metadata.matches(metadata):
self._tool_call_metadata_by_index[tool_call_index] = indexed_metadata.merged_with(metadata)
self._streamed_call_id_by_terminal_id[incoming_call_id] = indexed_call_id
return indexed_call_id
self._tool_call_metadata_by_index[tool_call_index] = indexed_metadata.marked_ambiguous()
if incoming_call_id == indexed_call_id:
return None
self._streamed_call_id_by_terminal_id[incoming_call_id] = incoming_call_id
return incoming_call_id
self._queue_tool_call_metadata_error(indexed_call_id, indexed_metadata, tool_call_index)
return None
if incoming_call_id is None:
return None
@ -280,6 +272,64 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
self._streamed_call_id_by_terminal_id[incoming_call_id] = incoming_call_id
return incoming_call_id
def _tool_call_metadata(self, tool_call: object) -> _StreamedToolCallMetadata:
function: Final = (
tool_call.get("function") if isinstance(tool_call, dict) else getattr(tool_call, "function", None)
)
function_name_raw: Final = (
function.get("name") if isinstance(function, dict) else getattr(function, "name", None)
)
tool_name, tool_namespace = self._responses_namespace_tool_call_fields(str(function_name_raw or ""))
call_type_raw: Final = (
tool_call.get("type") if isinstance(tool_call, dict) else getattr(tool_call, "type", None)
)
return _StreamedToolCallMetadata(
call_type=str(call_type_raw) if call_type_raw else None,
tool_name=tool_name or None,
tool_namespace=tool_namespace,
)
def _queue_tool_call_metadata_error(
self,
call_id: str,
metadata: _StreamedToolCallMetadata,
tool_call_index: int,
) -> None:
if self._pending_tool_call_error is not None:
return
output_index: Final = self._get_or_assign_tool_output_index(call_id)
arguments: Final = self._tool_args_by_call_id.get(call_id, "")
item_kwargs: Final = build_tool_call_item_kwargs(
call_id,
metadata.tool_name or "",
arguments,
"incomplete",
self._custom_tool_names,
)
item_kwargs["id"] = self._tool_item_id_by_call_id.get(call_id, item_kwargs["id"])
if metadata.tool_namespace:
item_kwargs["namespace"] = metadata.tool_namespace
self._sequence_number += 1
self._pending_tool_events.append(
OutputItemDoneEvent(
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
output_index=output_index,
sequence_number=self._sequence_number,
item=BaseLiteLLMOpenAIResponseObject(**item_kwargs),
)
)
self._pending_tool_call_error = litellm.InternalServerError(
message=f"Provider changed tool metadata at tool call index {tool_call_index}",
llm_provider=self.custom_llm_provider or "",
model=self.model,
)
def _raise_pending_tool_call_error(self) -> None:
if self._pending_tool_call_error is not None:
raise self._pending_tool_call_error
def _responses_namespace_tool_call_fields(self, fn_name: str) -> tuple[str, str | None]:
mapped: Final = self._namespace_tool_names.get(fn_name)
if mapped:
@ -343,7 +393,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
Note: Some providers (like Bedrock) send tool call arguments in one large chunk.
We split these into smaller deltas to match OpenAI's token-by-token streaming behavior.
"""
if not isinstance(tool_calls, list):
if self._pending_tool_call_error is not None or not isinstance(tool_calls, list):
return
for tc in tool_calls:
@ -359,14 +409,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
fn_name = str(getattr(fn, "name", "") or "")
fn_args_delta = serialize_tool_call_arguments(getattr(fn, "arguments", ""))
tool_name, tool_namespace = self._responses_namespace_tool_call_fields(fn_name)
call_type_raw = tc.get("type") if isinstance(tc, dict) else getattr(tc, "type", None)
metadata = _StreamedToolCallMetadata(
call_type=str(call_type_raw) if call_type_raw else None,
tool_name=tool_name or None,
tool_namespace=tool_namespace,
)
metadata = self._tool_call_metadata(tc)
call_id = self._resolve_streamed_tool_call_id(tc_index, call_id_raw, metadata)
if call_id is None:
if self._pending_tool_call_error is not None:
return
continue
output_index = self._get_or_assign_tool_output_index(call_id)
@ -400,7 +447,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
"""
Ensure tool calls that were not streamed as deltas still get emitted before response.completed.
"""
if self._final_tool_events_queued:
if self._final_tool_events_queued or self._pending_tool_call_error is not None:
return
self._final_tool_events_queued = True
@ -419,6 +466,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
continue
terminal_call_id = str(call_id_raw)
call_id = self._streamed_tool_call_id_for_terminal_call(tc, position) or terminal_call_id
if self._pending_tool_call_error is not None:
return
self._resolved_tool_call_id_by_position[position] = call_id
self._streamed_call_id_by_terminal_id[terminal_call_id] = call_id
output_index = self._get_or_assign_tool_output_index(call_id)
@ -937,6 +986,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
self._queue_final_tool_call_done_events(self.litellm_model_response)
if self._pending_tool_events:
return self._pending_tool_events.pop(0)
self._raise_pending_tool_call_error()
done_event: Final = self.return_default_done_events(self.litellm_model_response)
if done_event:
@ -1044,6 +1094,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
# Emit any pending tool events before reading a new chunk
if self._pending_tool_events:
return self._pending_tool_events.pop(0)
self._raise_pending_tool_call_error()
try:
chunk = self._take_buffered_chunk()
@ -1151,6 +1202,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
# Emit any pending tool events before reading a new chunk
if self._pending_tool_events:
return self._pending_tool_events.pop(0)
self._raise_pending_tool_call_error()
try:
buffered_chunk = self._take_buffered_chunk()
if buffered_chunk is not None:

1
tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py
View file

@ -3414,6 +3414,7 @@ class TestEnsureOutputItemContentPartAdded:
iterator._streamed_tool_call_ids_in_order = []
iterator._resolved_tool_call_id_by_position = {}
iterator._streamed_call_id_by_terminal_id = {}
iterator._pending_tool_call_error = None
iterator._ambiguous_tool_call_indexes = set()
iterator._next_tool_output_index = 1
iterator._final_tool_events_queued = False

88
tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py
View file

@ -15,6 +15,7 @@ from unittest.mock import AsyncMock, MagicMock
import pytest
import litellm
from litellm.responses.litellm_completion_transformation.streaming_iterator import (
LiteLLMCompletionStreamingIterator,
)
@ -915,16 +916,11 @@ def test_reused_index_with_new_call_id_preserves_first_streamed_identity(
("replacement_type", "replacement_name"),
(("function", "tool_b"), ("custom", "tool_a")),
)
def test_reused_index_with_changed_tool_metadata_starts_separate_identity(
def test_reused_index_with_changed_tool_metadata_fails_closed(
replacement_type: str,
replacement_name: str,
):
iterator = LiteLLMCompletionStreamingIterator(
model="test-model",
litellm_custom_stream_wrapper=AsyncMock(),
request_input="Test input",
responses_api_request={},
)
iterator = _build_iterator([])
iterator._queue_tool_call_delta_events(
[
@ -956,11 +952,34 @@ def test_reused_index_with_changed_tool_metadata_starts_separate_identity(
}
]
)
events = []
with pytest.raises(litellm.InternalServerError, match="changed tool metadata at tool call index 0"):
while True:
events.append(next(iterator))
added_items = [event.item for event in events if event.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED]
done_items = [event.item for event in events if event.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE]
delta_events = [event for event in events if event.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA]
assert [(item.id, item.call_id, item.name) for item in added_items] == [("fc_call_a", "call_a", "tool_a")]
assert [(item.id, item.call_id, item.name, item.status) for item in done_items] == [
("fc_call_a", "call_a", "tool_a", "incomplete")
]
assert "".join(event.delta for event in delta_events) == '{"safe":'
assert all(event.type != ResponsesAPIStreamEvents.RESPONSE_COMPLETED for event in events)
def test_terminal_tool_metadata_drift_fails_closed():
iterator = _build_iterator([])
iterator._queue_tool_call_delta_events(
[{"index": 0, "id": "call_b", "function": {"arguments": "true}"}}]
)
iterator._queue_tool_call_delta_events(
[{"index": 0, "function": {"arguments": "ignored"}}]
[
{
"index": 0,
"id": "call_stream",
"type": "function",
"function": {"name": "tool_a", "arguments": '{"safe":true}'},
}
]
)
terminal_response = ModelResponse(
id="chatcmpl-terminal",
@ -977,12 +996,9 @@ def test_reused_index_with_changed_tool_metadata_starts_separate_identity(
"tool_calls": [
{
"index": 0,
"id": "call_b",
"id": "call_terminal",
"type": "function",
"function": {
"name": replacement_name,
"arguments": '{"privileged":true}',
},
"function": {"name": "tool_b", "arguments": '{"privileged":true}'},
}
],
},
@ -990,38 +1006,20 @@ def test_reused_index_with_changed_tool_metadata_starts_separate_identity(
],
)
iterator._queue_final_tool_call_done_events(terminal_response)
completed = iterator._emit_response_completed_event(terminal_response)
added_items = [
event.item
for event in iterator._pending_tool_events
if event.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED
]
delta_events = [
event
for event in iterator._pending_tool_events
if event.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA
]
done_item = next(
event.item
for event in iterator._pending_tool_events
if event.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE
)
events = []
with pytest.raises(litellm.InternalServerError, match="changed tool metadata at tool call index 0"):
while True:
events.append(next(iterator))
assert completed is not None
assert [(item.id, item.call_id) for item in added_items] == [
("fc_call_a", "call_a"),
("fc_call_b", "call_b"),
added_items = [event.item for event in events if event.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED]
done_items = [event.item for event in events if event.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE]
assert [(item.id, item.call_id, item.name) for item in added_items] == [("fc_call_stream", "call_stream", "tool_a")]
assert [(item.id, item.call_id, item.name, item.status) for item in done_items] == [
("fc_call_stream", "call_stream", "tool_a", "incomplete")
]
assert "".join(event.delta for event in delta_events if event.item_id == "fc_call_a") == '{"safe":'
assert "".join(event.delta for event in delta_events if event.item_id == "fc_call_b") == '{"privileged":true}'
assert (done_item.id, done_item.call_id, done_item.name) == ("fc_call_b", "call_b", replacement_name)
completed_call = next(item for item in completed.response.output if item.type == "function_call")
assert (completed_call.id, completed_call.call_id, completed_call.name) == (
"fc_call_b",
"call_b",
replacement_name,
)
assert all(event.type != ResponsesAPIStreamEvents.RESPONSE_COMPLETED for event in events)
@pytest.mark.asyncio