From 06ff5e211aebeda51d2f5398ad2fba5c1bf2cfd5 Mon Sep 17 00:00:00 2001 From: Jason Date: Thu, 16 Apr 2026 22:02:10 +0800 Subject: [PATCH] fix(proxy/gemini): reconcile synthesized tool-call ids with later real ids + preserve thoughtSignature MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three related findings on `streaming_gemini.rs` for Gemini's cumulative `streamGenerateContent` stream, all centered on `merge_tool_call_snapshots`: 1. (P1) Match upgraded tool-call IDs by position. When Gemini delivers a `functionCall` without an id on chunk 1 (cc-switch synthesizes `gemini_synth_*`) and then upgrades it to a real id on chunk 2, the `Some(incoming_id)` branch only matched by id and missed the existing synthesized snapshot. A second entry would be pushed, yielding duplicate `tool_use` content blocks at stream end — one with the synthesized id, one with the real id — which could trigger duplicate tool execution and break tool_result correlation. Add a positional fallback: when no id match exists but the same-position slot holds a synthesized id, merge into it. `or(preserved_id)` already lets the real id win the merge. 2. (P2) Preserve prior thoughtSignature when merging snapshots. `tool_call_snapshots[index] = tool_call` overwrote the slot entirely, dropping any `thoughtSignature` captured on an earlier chunk if the current cumulative snapshot omitted it. Since `build_shadow_assistant_parts` writes `thoughtSignature` into the shadow turn from `tool_call.thought_signature`, a dropped signature would cause later replay requests to Gemini to be rejected with invalid-signature errors. Preserve the existing signature when the incoming chunk does not carry one. 3. (P2) Document the part-order streaming trade-off. All `tool_use` content blocks are emitted after the final text `content_block_stop`, so interleaved [text, functionCall, text, functionCall] parts arrive at the Anthropic client as [text(concat), tool_use, tool_use] — different from the non-streaming transformer, which preserves part order. This is intentional given the cumulative snapshot model and the consumers we target (claude-code-like clients don't depend on strict interleaving for tool execution correctness). Add a block comment at the flush site describing the trade-off and what a strict-order fix would entail, so this isn't rediscovered as a bug later. Regression tests: - upgraded_real_id_merges_into_existing_synthesized_snapshot - thought_signature_preserved_when_later_chunk_omits_it Test count: 868 -> 870. clippy 1.95 clean. fmt clean. --- .../src/proxy/providers/streaming_gemini.rs | 139 +++++++++++++++++- 1 file changed, 138 insertions(+), 1 deletion(-) diff --git a/src-tauri/src/proxy/providers/streaming_gemini.rs b/src-tauri/src/proxy/providers/streaming_gemini.rs index 0a81d0c4c..815c19ed1 100644 --- a/src-tauri/src/proxy/providers/streaming_gemini.rs +++ b/src-tauri/src/proxy/providers/streaming_gemini.rs @@ -122,7 +122,28 @@ fn merge_tool_call_snapshots( let existing_index = match tool_call.id.as_deref() { Some(incoming_id) => tool_call_snapshots .iter() - .position(|existing| existing.id.as_deref() == Some(incoming_id)), + .position(|existing| existing.id.as_deref() == Some(incoming_id)) + .or_else(|| { + // Fallback for the "synth -> real id upgrade" case: + // Gemini's cumulative stream may deliver the first chunk + // of a tool call without an id (we synthesize one) and + // then upgrade it to a genuine id on a later chunk. A + // pure id-match would miss the existing synthesized + // snapshot and push a second entry, yielding duplicate + // `tool_use` content blocks at stream end. If the + // same-position slot currently holds a synthesized id, + // merge into it — `or(preserved_id)` below will keep + // the real id, dropping the synthesized one. + tool_call_snapshots + .get(position) + .filter(|existing| { + matches!( + existing.id.as_deref(), + Some(id) if is_synthesized_tool_call_id(id) + ) + }) + .map(|_| position) + }), None => tool_call_snapshots .get(position) .filter(|existing| match existing.id.as_deref() { @@ -140,8 +161,24 @@ fn merge_tool_call_snapshots( if let Some(index) = existing_index { // Preserve any synthesized id assigned on a previous chunk so the // Anthropic-visible id stays stable across the whole stream. + // When incoming carries a real Gemini id and the slot holds a + // synthesized one, `Some(real).or(Some(synth)) == Some(real)` + // so the upgrade wins naturally. let preserved_id = tool_call_snapshots[index].id.clone(); tool_call.id = tool_call.id.or(preserved_id); + + // Preserve `thought_signature` across chunks. Gemini's cumulative + // stream may include `thoughtSignature` on one chunk and omit it + // on a subsequent cumulative snapshot of the same part, even + // though the signature still belongs to the call. A blind + // overwrite would drop it, so the shadow turn we record (and + // later replay) would be missing `thoughtSignature` and the + // upstream would reject the follow-up for invalid signature. + if tool_call.thought_signature.is_none() { + tool_call + .thought_signature + .clone_from(&tool_call_snapshots[index].thought_signature); + } } if tool_call.id.is_none() { tool_call.id = Some(synthesize_tool_call_id()); @@ -455,6 +492,33 @@ pub fn create_anthropic_sse_stream_from_gemini a `gemini_synth_*` id is assigned. + "data: {\"responseId\":\"rupg\",\"modelVersion\":\"gemini-2.5-pro\",\"candidates\":[{\"content\":{\"parts\":[{\"functionCall\":{\"name\":\"get_weather\",\"args\":{\"city\":\"Tokyo\"}}}]}}],\"usageMetadata\":{\"promptTokenCount\":4,\"totalTokenCount\":6}}\n\n", + // Chunk 2: cumulative snapshot upgrades the same call to a + // real Gemini id. Must merge into the existing slot, not + // spawn a second snapshot. + "data: {\"responseId\":\"rupg\",\"modelVersion\":\"gemini-2.5-pro\",\"candidates\":[{\"finishReason\":\"STOP\",\"content\":{\"parts\":[{\"functionCall\":{\"id\":\"real_id_abc\",\"name\":\"get_weather\",\"args\":{\"city\":\"Tokyo\",\"units\":\"c\"}}}]}}],\"usageMetadata\":{\"promptTokenCount\":4,\"totalTokenCount\":9}}\n\n", + ]); + + // Exactly one tool_use block (not two). + assert_eq!( + output.matches("\"type\":\"tool_use\"").count(), + 1, + "id upgrade must merge into the synthesized snapshot, not duplicate it: {output}" + ); + // The emitted tool_use id is the real Gemini id, not the synthesized one. + assert!( + output.contains("\"id\":\"real_id_abc\""), + "expected real id to win after upgrade: {output}" + ); + assert!( + !output.contains("\"id\":\"gemini_synth_"), + "synthesized id must be dropped when a real id arrives: {output}" + ); + // Args from the final cumulative snapshot are emitted. + assert!(output.contains("units")); + } + + /// Regression for Codex P2: Gemini's cumulative stream may include + /// `thoughtSignature` on one chunk and omit it on a later cumulative + /// snapshot of the same call. A blind `tool_call_snapshots[index] = + /// tool_call` overwrite would drop the signature, so the shadow turn + /// recorded (and later replayed to Gemini) would miss + /// `thoughtSignature` and the upstream would reject the follow-up. + /// `merge_tool_call_snapshots` must retain the prior signature when + /// the incoming chunk does not carry one. + #[test] + fn thought_signature_preserved_when_later_chunk_omits_it() { + let store = Arc::new(GeminiShadowStore::with_limits(8, 4)); + collect_stream_output_with_shadow( + vec![ + // Chunk 1: carries thoughtSignature "sig-keep". + "data: {\"responseId\":\"rsig\",\"modelVersion\":\"gemini-2.5-pro\",\"candidates\":[{\"content\":{\"parts\":[{\"functionCall\":{\"id\":\"call_1\",\"name\":\"get_weather\",\"args\":{\"city\":\"Tokyo\"}},\"thoughtSignature\":\"sig-keep\"}]}}],\"usageMetadata\":{\"promptTokenCount\":4,\"totalTokenCount\":6}}\n\n", + // Chunk 2: cumulative update for the same call, but + // thoughtSignature is omitted — common for Gemini's + // one-shot signature fields. + "data: {\"responseId\":\"rsig\",\"modelVersion\":\"gemini-2.5-pro\",\"candidates\":[{\"finishReason\":\"STOP\",\"content\":{\"parts\":[{\"functionCall\":{\"id\":\"call_1\",\"name\":\"get_weather\",\"args\":{\"city\":\"Tokyo\",\"units\":\"c\"}}}]}}],\"usageMetadata\":{\"promptTokenCount\":4,\"totalTokenCount\":9}}\n\n", + ], + store.clone(), + "provider-sig", + "session-sig", + ); + + let shadow = store + .latest_assistant_content("provider-sig", "session-sig") + .expect("shadow turn must be recorded"); + assert_eq!(shadow["parts"][0]["functionCall"]["id"], "call_1"); + assert_eq!( + shadow["parts"][0]["thoughtSignature"], "sig-keep", + "prior thoughtSignature must survive a later chunk that omits it: {shadow}" + ); + } }