fix(gemini): assign unique stream indices to parallel tool calls

The streaming translator in agent/gemini_cloudcode_adapter.py keyed OpenAI
tool-call indices by function name, so when the model emitted multiple
parallel functionCall parts with the same name in a single turn (e.g.
three read_file calls in one response), they all collapsed onto index 0.
Downstream aggregators that key chunks by index would overwrite or drop
all but the first call.

Replace the name-keyed dict with a per-stream counter that persists across
SSE events. Each functionCall part now gets a fresh, unique index,
matching the non-streaming path which already uses enumerate(parts).

Add TestTranslateStreamEvent covering parallel-same-name calls, index
persistence across events, and finish-reason promotion to tool_calls.
This commit is contained in:
haileymarshall 2026-04-18 01:15:50 +01:00 committed by Teknium
parent d990fa52ed
commit 49282b6e04
2 changed files with 77 additions and 6 deletions

View file

@ -505,9 +505,16 @@ def _iter_sse_events(response: httpx.Response) -> Iterator[Dict[str, Any]]:
def _translate_stream_event(
event: Dict[str, Any],
model: str,
tool_call_indices: Dict[str, int],
tool_call_counter: List[int],
) -> List[_GeminiStreamChunk]:
"""Unwrap Code Assist envelope and emit OpenAI-shaped chunk(s)."""
"""Unwrap Code Assist envelope and emit OpenAI-shaped chunk(s).
``tool_call_counter`` is a single-element list used as a mutable counter
across events in the same stream. Each ``functionCall`` part gets a
fresh, unique OpenAI ``index`` keying by function name would collide
whenever the model issues parallel calls to the same tool (e.g. reading
three files in one turn).
"""
inner = event.get("response") if isinstance(event.get("response"), dict) else event
candidates = inner.get("candidates") or []
if not candidates:
@ -533,7 +540,8 @@ def _translate_stream_event(
fc = part.get("functionCall")
if isinstance(fc, dict) and fc.get("name"):
name = str(fc["name"])
idx = tool_call_indices.setdefault(name, len(tool_call_indices))
idx = tool_call_counter[0]
tool_call_counter[0] += 1
try:
args_str = json.dumps(fc.get("args") or {}, ensure_ascii=False)
except (TypeError, ValueError):
@ -550,7 +558,7 @@ def _translate_stream_event(
finish_reason_raw = str(cand.get("finishReason") or "")
if finish_reason_raw:
mapped = _map_gemini_finish_reason(finish_reason_raw)
if tool_call_indices:
if tool_call_counter[0] > 0:
mapped = "tool_calls"
chunks.append(_make_stream_chunk(model=model, finish_reason=mapped))
return chunks
@ -734,9 +742,9 @@ class GeminiCloudCodeClient:
# Materialize error body for better diagnostics
response.read()
raise _gemini_http_error(response)
tool_call_indices: Dict[str, int] = {}
tool_call_counter: List[int] = [0]
for event in _iter_sse_events(response):
for chunk in _translate_stream_event(event, model, tool_call_indices):
for chunk in _translate_stream_event(event, model, tool_call_counter):
yield chunk
except httpx.HTTPError as exc:
raise CodeAssistError(