#11538

google-generative-ai: streamed parts with separate thought signatures are merged or dropped

Openyanqiangluopened 11d ago

What happened?

The Google stream parser coalesces all consecutive text parts (and all thinking parts) into one block and keeps only the last thoughtSignature via retainThoughtSignature. Parts without a text field are skipped entirely. On replay, convertMessages emits one part per block.

So the next request does not carry the signatures the way the model returned them:

  • A signature that arrived on its own trailing empty-text part is moved onto the preceding text part.
  • When two parts each carry a signature, only the last one survives.
  • A part with only thoughtSignature (no text key) is dropped, together with its signature.

Google's thought-signature docs say to send each signature back inside its original Part, not to merge a signed part with an unsigned one, and not to combine two signed parts. They also note that when streaming, the signature may arrive on a part with empty text. Text-part signatures are not validated, so this does not cause a 400. Per the docs, the cost is reasoning quality on later turns.

I have seen the first shape (unsigned text, then { text: "", thoughtSignature }) from a current Gemini model in normal multi-turn tool use. The other two I have only reproduced offline.

Steps to reproduce

Offline, against @earendil-works/pi-ai@1.4.0: a local server streams the chunks below to api/google-generative-ai stream(), then the resulting message goes through api/google-shared convertMessages().

== trailing empty-text signature
received : [{"text":"Hello "},{"text":"world"},{"text":"","thoughtSignature":"QUFBQQ=="}]
replayed : [{"text":"Hello world","thoughtSignature":"QUFBQQ=="}]

== two signed text parts
received : [{"text":"First.","thoughtSignature":"QUFBQQ=="},{"text":" Second.","thoughtSignature":"QkJCQg=="}]
replayed : [{"text":"First. Second.","thoughtSignature":"QkJCQg=="}]

== signature-only part
received : [{"text":"Answer"},{"thoughtSignature":"Q0NDQw=="}]
replayed : [{"text":"Answer"}]
Repro script (no network; npm i @earendil-works/pi-ai@1.4.0, then node repro.mjs)
import http from "node:http";
import { stream } from "@earendil-works/pi-ai/api/google-generative-ai";
import { convertMessages } from "@earendil-works/pi-ai/api/google-shared";

const SIG_A = "QUFBQQ==", SIG_B = "QkJCQg==", SIG_C = "Q0NDQw==";
const scenarios = {
  // Real-world shape: text streamed unsigned, signature on a trailing empty-text part.
  trailingEmptyTextSignature: [
    [{ text: "Hello " }], [{ text: "world" }], [{ text: "", thoughtSignature: SIG_A }],
  ],
  // Two separately signed text parts in one response.
  twoSignedTextParts: [
    [{ text: "First.", thoughtSignature: SIG_A }], [{ text: " Second.", thoughtSignature: SIG_B }],
  ],
  // Signature-only part (no text field).
  signatureOnlyPart: [
    [{ text: "Answer" }], [{ thoughtSignature: SIG_C }],
  ],
};
let current;
const server = http.createServer((req, res) => {
  let body = ""; req.on("data", d => body += d); req.on("end", () => {
    res.writeHead(200, { "content-type": "text/event-stream" });
    const chunks = scenarios[current];
    chunks.forEach((parts, i) => {
      const c = { candidates: [{ content: { role: "model", parts }, ...(i === chunks.length - 1 && { finishReason: "STOP" }) }] };
      res.write(`data: ${JSON.stringify(c)}\r\n\r\n`);
    });
    res.end();
  });
});
await new Promise(r => server.listen(0, r));
const model = { id: "gemini-3-pro-preview", name: "g", api: "google-generative-ai", provider: "google",
  baseUrl: `http://127.0.0.1:${server.address().port}/v1beta`, reasoning: true, input: ["text"],
  cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, contextWindow: 1e6, maxTokens: 8192 };
for (const name of Object.keys(scenarios)) {
  current = name;
  const s = stream(model, { messages: [{ role: "user", content: "hi", timestamp: Date.now() }] }, { apiKey: "x" });
  const msg = await s.result();
  const replay = convertMessages(model, { messages: [{ role: "user", content: "hi", timestamp: 0 }, msg] });
  console.log(`\n== ${name}`);
  console.log("received :", JSON.stringify(scenarios[name].flat()));
  console.log("stored   :", JSON.stringify(msg.content), msg.stopReason, msg.errorMessage ?? "");
  console.log("replayed :", JSON.stringify(replay.at(-1).parts));
}
server.close();

Expected behavior

The replayed parts keep the signature boundaries as received, for the same provider and model:

[{"text":"Hello world"},{"text":"","thoughtSignature":"QUFBQQ=="}]
[{"text":"First.","thoughtSignature":"QUFBQQ=="},{"text":" Second.","thoughtSignature":"QkJCQg=="}]
[{"text":"Answer"},{"thoughtSignature":"Q0NDQw=="}]

Unsigned deltas can still be coalesced. Only the parts that carry a signature need to stay separate. Displayed text does not need to change.

Related: #7356 / #7362 fixed the replay side for signed empty blocks. This is the stream-parser side.

Version

1.4.0 (same code on main at f715a15c)

Activity

0

No activity yet.