From c6bfb2225c98ef75b2d54fdff9e82a9c6423657a Mon Sep 17 00:00:00 2001 From: Mateusz Cikowski Date: Fri, 30 Oct 2026 13:03:45 -0400 Subject: [PATCH] fix(agents): send Gemini thought signatures back through the chat-completions bridge Co-Authored-By: Claude Opus 5.5 --- .../agents/agents/gemini/gemini_bridge.py | 17 +++++---- tst/unit/agents/gemini/test_litellm_proxy.py | 35 +++++++++++++++++++ 2 files changed, 46 insertions(+), 6 deletions(-) diff --git a/src/agent_env/agents/agents/gemini/gemini_bridge.py b/src/agent_env/agents/agents/gemini/gemini_bridge.py index a86927102..b98ff2059 100644 --- a/src/agent_env/agents/agents/gemini/gemini_bridge.py +++ b/src/agent_env/agents/agents/gemini/gemini_bridge.py @@ -751,13 +751,17 @@ def gemini_request_to_openai(greq: dict, model: str) -> dict: cid = f"call_{counter}" counter += 1 pending_ids.setdefault(fc.get("name"), []).append(cid) - tool_calls.append({ + tool_call: dict = { "id": cid, "type": "function", "function": { "name": fc.get("name"), "arguments": json.dumps(fc.get("args") or fc.get("arguments") or {}), }, - }) + } + sig = p.get("thoughtSignature") or p.get("thought_signature") + if sig: + tool_call["provider_specific_fields"] = {"thought_signature": sig} + tool_calls.append(tool_call) msg: dict = {"role": "assistant", "content": "\n".join(text_bits) if text_bits else None} if tool_calls: @@ -826,8 +830,8 @@ def gemini_request_to_openai(greq: dict, model: str) -> dict: } -def _thought_signature(msg: dict, tc: dict) -> str | None: - for holder in (tc, msg): +def _thought_signature(*holders: dict) -> str | None: + for holder in holders: psf = holder.get("provider_specific_fields") or {} sig = psf.get("thought_signature") or psf.get("thoughtSignature") if isinstance(sig, str) and sig: @@ -846,14 +850,15 @@ def openai_response_to_gemini(oai: dict) -> dict: parts: list[dict] = [] if msg.get("content"): parts.append({"text": msg["content"]}) - for tc in msg.get("tool_calls") or []: + for i, tc in enumerate(msg.get("tool_calls") or []): fn = tc.get("function") or {} try: args = json.loads(fn.get("arguments") or "{}") except (json.JSONDecodeError, TypeError): args = {} fcpart: dict = {"functionCall": {"name": fn.get("name"), "args": args}} - sig = _thought_signature(msg, tc) + # Gemini signs only the first call of a parallel batch, and bills each replayed copy. + sig = _thought_signature(tc, msg) if i == 0 else _thought_signature(tc) if sig: fcpart["thoughtSignature"] = sig parts.append(fcpart) diff --git a/tst/unit/agents/gemini/test_litellm_proxy.py b/tst/unit/agents/gemini/test_litellm_proxy.py index 352c7030c..1835c0679 100644 --- a/tst/unit/agents/gemini/test_litellm_proxy.py +++ b/tst/unit/agents/gemini/test_litellm_proxy.py @@ -246,3 +246,38 @@ def test_bridge_upstream_errors(vertex, client): assert (limited.status_code, limited.content) == (429, b'{"error":"slow down"}') assert (garbled.status_code, garbled.json()) == (502, {"detail": "vertex bridge: non-JSON upstream response"}) assert (unreachable.status_code, unreachable.json()) == (502, {"detail": "vertex bridge upstream failed: refused"}) + + + +def _tool_turn(*calls: tuple[str, str | None]) -> httpx.Response: + tool_calls = [{"id": f"call_{i}", "type": "function", "function": {"name": name, "arguments": "{}"}, + **({"provider_specific_fields": {"thought_signature": sig}} if sig else {})} + for i, (name, sig) in enumerate(calls)] + return httpx.Response(200, json={"choices": [{"finish_reason": "tool_calls", + "message": {"content": None, "tool_calls": tool_calls}}]}) + + +def test_each_tool_calls_thought_signature_survives_the_round_trip(vertex, client): + contents = list(_GREQ["contents"]) + for reply in (_tool_turn(("get_time", "sig-1")), _tool_turn(("get_time", "sig-2"), ("get_weather", None))): + vertex.response = reply + model_turn = client.post(GENERATE, json={"contents": contents}).json()["candidates"][0]["content"] + results = [{"functionResponse": {"name": p["functionCall"]["name"], "response": {}}} for p in model_turn["parts"]] + contents += [model_turn, {"role": "user", "parts": results}] + + client.post(GENERATE, json={"contents": contents}) + + assistant = [m for m in json.loads(vertex.requests[-1].content)["messages"] if m["role"] == "assistant"] + sent = [tc.get("provider_specific_fields") for m in assistant for tc in m["tool_calls"]] + assert sent == [{"thought_signature": "sig-1"}, {"thought_signature": "sig-2"}, None] + + +def test_a_message_level_signature_goes_only_to_the_first_parallel_call(vertex, client): + vertex.response = httpx.Response(200, json={"choices": [{"finish_reason": "tool_calls", "message": { + "content": None, "provider_specific_fields": {"thought_signatures": ["sig-a"]}, + "tool_calls": [{"id": f"call_{n}", "type": "function", "function": {"name": n, "arguments": "{}"}} for n in "ab"], + }}]}) + + parts = client.post(GENERATE, json=_GREQ).json()["candidates"][0]["content"]["parts"] + + assert [p.get("thoughtSignature") for p in parts] == ["sig-a", None]