fix(agents): send Gemini thought signatures back through the chat-completions bridge

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)

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]