@@ -772,6 +772,12 @@ class _StreamState:
772772 """
773773
774774 parts : list [types .Part ] = dataclasses .field (default_factory = list )
775+ # Maps a function-call step's ``index`` to the part started at that step, so
776+ # interleaved calls route their argument deltas and stops to the matching
777+ # step instead of always landing on the most recently started call.
778+ fc_parts_by_index : dict [int , types .Part ] = dataclasses .field (
779+ default_factory = dict
780+ )
775781 web_search_queries : list [str ] = dataclasses .field (default_factory = list )
776782 grounding_chunks : list [types .GroundingChunk ] = dataclasses .field (
777783 default_factory = list
@@ -836,23 +842,43 @@ def _handle_media(
836842 return _partial_part_response (part , interaction_id )
837843
838844
845+ def _resolve_streaming_function_call_part (
846+ index : int | None , state : _StreamState
847+ ) -> types .Part | None :
848+ """Resolve the function-call part a streaming event applies to.
849+
850+ Streaming events carry the ``index`` of the step they belong to. When that
851+ index maps to a known function-call part we use it directly, so argument
852+ deltas and step stops for interleaved calls route to the correct step instead
853+ of always landing on the most recently started call. Events without an index
854+ (or from builds that don't track one) fall back to the last started function
855+ call to preserve the previous behavior.
856+ """
857+ if index is not None and index in state .fc_parts_by_index :
858+ return state .fc_parts_by_index [index ]
859+ if state .parts and state .parts [- 1 ].function_call :
860+ return state .parts [- 1 ]
861+ return None
862+
863+
839864def _handle_arguments_delta (
840- delta : StepDeltaData , state : _StreamState , interaction_id : str | None
865+ delta : StepDeltaData ,
866+ state : _StreamState ,
867+ interaction_id : str | None ,
868+ index : int | None = None ,
841869) -> LlmResponse | None :
842- if not state .parts :
843- return None
844- last_part = state .parts [- 1 ]
845- if not last_part .function_call :
870+ target_part = _resolve_streaming_function_call_part (index , state )
871+ if target_part is None or not target_part .function_call :
846872 return None
847873 delta_args = delta .arguments
848- if delta_args is None or last_part .function_call .partial_args is None :
874+ if delta_args is None or target_part .function_call .partial_args is None :
849875 return None
850- last_part .function_call .partial_args .append (
876+ target_part .function_call .partial_args .append (
851877 types .PartialArg (string_value = delta_args )
852878 )
853879 chunk_part = types .Part (
854880 function_call = types .FunctionCall (
855- name = last_part .function_call .name ,
881+ name = target_part .function_call .name ,
856882 partial_args = [types .PartialArg (string_value = delta_args )],
857883 )
858884 )
@@ -1066,6 +1092,8 @@ def convert_interaction_event_to_llm_response(
10661092 )
10671093 part = types .Part (function_call = fc )
10681094 state .parts .append (part )
1095+ if event .index is not None :
1096+ state .fc_parts_by_index [event .index ] = part
10691097
10701098 return LlmResponse (
10711099 content = types .Content (role = 'model' , parts = [part ]),
@@ -1087,7 +1115,7 @@ def convert_interaction_event_to_llm_response(
10871115 elif delta_type in ('image' , 'audio' , 'video' , 'document' ):
10881116 return _handle_media (delta , state , interaction_id )
10891117 elif delta_type == 'arguments_delta' :
1090- return _handle_arguments_delta (delta , state , interaction_id )
1118+ return _handle_arguments_delta (delta , state , interaction_id , event . index )
10911119 elif delta_type == 'code_execution_call' :
10921120 return _handle_code_execution_call (delta , state , interaction_id )
10931121 elif delta_type == 'code_execution_result' :
@@ -1104,8 +1132,9 @@ def convert_interaction_event_to_llm_response(
11041132 return _handle_unknown_delta (delta , state , interaction_id )
11051133
11061134 elif isinstance (event , StepStop ):
1107- if state .parts and state .parts [- 1 ].function_call :
1108- fc = state .parts [- 1 ].function_call
1135+ target_part = _resolve_streaming_function_call_part (event .index , state )
1136+ if target_part is not None and target_part .function_call :
1137+ fc = target_part .function_call
11091138 if fc .partial_args is not None :
11101139 arg_str = '' .join (pa .string_value or '' for pa in fc .partial_args )
11111140
0 commit comments