fix(responses): preserve identity across replacement deltas

litellm/responses/litellm_completion_transformation/streaming_iterator.py

@@ -119,8 +119,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
         self._tool_call_id_by_index: dict[int, str] = {}
         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._ambiguous_tool_call_indexes: set[int] = set()
+        self._streamed_call_id_by_provider_id: dict[str, str] = {}  # mutable-ok: provider identity correlation
         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
@@ -159,8 +158,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
             return None
 
     def _streamed_tool_call_id_at_position(self, position: int) -> str | None:
-        if position in self._ambiguous_tool_call_indexes:
-            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
@@ -177,8 +174,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
         """Match a terminal aggregate tool call to the identity emitted while streaming."""
         tool_call_index: Final = self._normalize_tool_call_index(tool_call)
         if tool_call_index is not None:
-            if tool_call_index in self._ambiguous_tool_call_indexes:
-                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
@@ -186,6 +181,23 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
                 return None
         return self._streamed_tool_call_id_at_position(position)
 
+    def _resolve_streamed_tool_call_id(self, tool_call_index: int | None, call_id_raw: object) -> str | None:
+        if tool_call_index is None:
+            return str(call_id_raw) if call_id_raw else None
+
+        indexed_call_id: Final = self._tool_call_id_by_index.get(tool_call_index)
+        if indexed_call_id is not None:
+            if call_id_raw:
+                self._streamed_call_id_by_provider_id[str(call_id_raw)] = indexed_call_id
+            return indexed_call_id
+
+        if not call_id_raw:
+            return None
+
+        call_id: Final = str(call_id_raw)
+        self._tool_call_id_by_index[tool_call_index] = call_id
+        return call_id
+
     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:
@@ -255,25 +267,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
         for tc in tool_calls:
             tc_index = self._normalize_tool_call_index(tc)
             call_id_raw = tc.get("id") if isinstance(tc, dict) else getattr(tc, "id", None)
-            call_id = ""
-
-            if call_id_raw:
-                call_id = str(call_id_raw)
-                if tc_index is not None:
-                    existing_call_id = self._tool_call_id_by_index.get(tc_index)
-                    if existing_call_id is not None and existing_call_id != call_id:
-                        # Reusing the same index for multiple call_ids is ambiguous for id-less deltas.
-                        # Guard against silent misrouting by disabling index fallback for this index.
-                        self._ambiguous_tool_call_indexes.add(tc_index)
-                    self._tool_call_id_by_index[tc_index] = call_id
-            elif tc_index is not None:
-                if tc_index in self._ambiguous_tool_call_indexes:
-                    continue
-                mapped_call_id = self._tool_call_id_by_index.get(tc_index)
-                if mapped_call_id:
-                    call_id = mapped_call_id
-
-            if not call_id:
+            call_id = self._resolve_streamed_tool_call_id(tc_index, call_id_raw)
+            if call_id is None:
                 continue
 
             fn = tc.get("function") if isinstance(tc, dict) else getattr(tc, "function", None)
@@ -338,7 +333,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
             terminal_call_id = str(call_id_raw)
             call_id = self._streamed_tool_call_id_for_terminal_call(tc, position) or terminal_call_id
             self._resolved_tool_call_id_by_position[position] = call_id
-            self._streamed_call_id_by_terminal_id[terminal_call_id] = call_id
+            self._streamed_call_id_by_provider_id[terminal_call_id] = call_id
             output_index = self._get_or_assign_tool_output_index(call_id)
 
             fn = tc.get("function") if isinstance(tc, dict) else getattr(tc, "function", None)
@@ -1216,7 +1211,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
     def _output_item_with_streamed_tool_identity(self, item: object, tool_position: int) -> object:
         terminal_call_id: Final = getattr(item, "call_id", None)
         resolved_by_call_id: Final = (
-            self._streamed_call_id_by_terminal_id.get(terminal_call_id) if isinstance(terminal_call_id, str) else None
+            self._streamed_call_id_by_provider_id.get(terminal_call_id) if isinstance(terminal_call_id, str) else None
         )
         resolved_by_position: Final = self._resolved_tool_call_id_by_position.get(tool_position)
         streamed_call_id: Final = (

tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py

@@ -804,7 +804,7 @@ def test_completed_snapshot_correlates_function_after_server_tool_replacement():
     assert function_call.arguments == '{"city":"Paris"}'
 
 
-def test_reused_index_with_new_call_id_marks_fallback_ambiguous():
+def test_reused_index_with_new_call_id_preserves_first_streamed_identity():
     iterator = LiteLLMCompletionStreamingIterator(
         model="test-model",
         litellm_custom_stream_wrapper=AsyncMock(),
@@ -828,39 +828,77 @@ def test_reused_index_with_new_call_id_marks_fallback_ambiguous():
                 "index": 0,
                 "id": "call_b",
                 "type": "function",
-                "function": {"name": "tool_b", "arguments": '{"b":'},
+                "function": {"name": "tool_a", "arguments": "1"},
             }
         ]
     )
-    # Ambiguous chunk: index reused and id missing. We should skip fallback rather than misroute.
     iterator._queue_tool_call_delta_events(
         [
             {
                 "index": 0,
                 "type": "function",
-                "function": {"arguments": "1}"},
+                "function": {"arguments": "}"},
             }
         ]
     )
+    streamed_argument_events = [
+        event
+        for event in iterator._pending_tool_events
+        if event.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA
+    ]
+
+    terminal_response = ModelResponse(
+        id="chatcmpl-terminal",
+        created=123,
+        model="test-model",
+        object="chat.completion",
+        choices=[
+            {
+                "index": 0,
+                "finish_reason": "tool_calls",
+                "message": {
+                    "role": "assistant",
+                    "content": None,
+                    "tool_calls": [
+                        {
+                            "index": 0,
+                            "id": "call_b",
+                            "type": "function",
+                            "function": {"name": "tool_a", "arguments": '{"a":1}'},
+                        }
+                    ],
+                },
+            }
+        ],
+    )
+    iterator._queue_final_tool_call_done_events(terminal_response)
+    completed = iterator._emit_response_completed_event(terminal_response)
 
     all_events = []
     while iterator._pending_tool_events:
         all_events.append(iterator._pending_tool_events.pop(0))
 
-    delta_events = [
-        evt
-        for evt in all_events
-        if evt.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA
+    added_items = [event.item for event in all_events if event.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED]
+    argument_events = [
+        event
+        for event in all_events
+        if event.type
+        in {
+            ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA,
+            ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DONE,
+        }
     ]
-    arguments_by_call_id = {}
-    for evt in delta_events:
-        arguments_by_call_id.setdefault(evt.item_id, "")
-        arguments_by_call_id[evt.item_id] += evt.delta
+    done_item = next(event.item for event in all_events if event.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE)
 
-    assert arguments_by_call_id["fc_call_a"] == '{"a":'
-    assert arguments_by_call_id["fc_call_b"] == '{"b":'
-    assert arguments_by_call_id["fc_call_a"] != '{"a":1}'
-    assert arguments_by_call_id["fc_call_b"] != '{"b":1}'
+    assert completed is not None
+    assert [(item.id, item.call_id) for item in added_items] == [("fc_call_a", "call_a")]
+    assert {event.item_id for event in streamed_argument_events} == {"fc_call_a"}
+    assert "".join(event.delta for event in streamed_argument_events) == '{"a":1}'
+    assert {event.item_id for event in argument_events} == {"fc_call_a"}
+    assert argument_events[-1].arguments == '{"a":1}'
+    assert (done_item.id, done_item.call_id) == ("fc_call_a", "call_a")
+    completed_call = next(item for item in completed.response.output if item.type == "function_call")
+    assert (completed_call.id, completed_call.call_id) == ("fc_call_a", "call_a")
 
 
 @pytest.mark.asyncio