From 0cf2c7c32676d4fda6889e98bb3b2b900af936bb Mon Sep 17 00:00:00 2001 From: Jiri Sacha Date: Sun, 30 Aug 2026 19:16:28 +0200 Subject: [PATCH] fix(agents): preserve Gemini thought_signature across tool-calling turns Gemini 3 models reason before they act, and when that reasoning produces a function call the response carries an opaque signature at tool_calls[].extra_content.google.thought_signature which must be sent back with the same tool call on the next turn. The reasoning is never returned as text, so the signature is the only thing carrying it forward, and Gemini rejects a request that omits it: 400 INVALID_ARGUMENT - "Function call is missing a thought_signature in functionCall parts. This is required for tools to work correctly." The field sits outside the OpenAI schema, so the OpenAI .NET SDK and Microsoft.Extensions.AI drop it while remapping the response, and the follow-up request is rebuilt without it. The first call of a conversation succeeds and the second fails, which in practice means an agent dies on its first tool result -- catalog and document generation never get past their opening step. Handle it at the HTTP layer, where the raw JSON is still intact, following the pattern FinishReasonNormalizingHandler already uses for a different Gemini/OpenAI mismatch on the same layer. ThoughtSignatureHandler records each signature as the response streams past, keyed by tool call id, and restores it on any later request naming the same call. Responses are only read, never modified; only requests bound for Gemini's host are touched; parsing failures are swallowed so a malformed chunk can never break a response; and the cache is bounded, since a signature runs to several kilobytes. Verified against the live API: without the handler the second turn of a tool conversation returns 400 with the signature error, with it the same conversation returns 200 and the model's answer. --- src/OpenDeepWiki/Agents/AgentFactory.cs | 18 +- .../Agents/ThoughtSignatureHandler.cs | 395 ++++++++++++++++++ .../Agents/ThoughtSignatureHandlerTests.cs | 224 ++++++++++ 3 files changed, 631 insertions(+), 6 deletions(-) create mode 100644 src/OpenDeepWiki/Agents/ThoughtSignatureHandler.cs create mode 100644 tests/OpenDeepWiki.Tests/Agents/ThoughtSignatureHandlerTests.cs diff --git a/src/OpenDeepWiki/Agents/AgentFactory.cs b/src/OpenDeepWiki/Agents/AgentFactory.cs index 101cef8e..7bd405fb 100644 --- a/src/OpenDeepWiki/Agents/AgentFactory.cs +++ b/src/OpenDeepWiki/Agents/AgentFactory.cs @@ -64,19 +64,25 @@ public class AgentFactory(IOptions options) /// /// Creates an HttpClient with the full handler chain: - /// FinishReasonNormalizingHandler -> LoggingHttpHandler -> HttpClientHandler + /// ThoughtSignatureHandler -> FinishReasonNormalizingHandler -> LoggingHttpHandler + /// -> HttpClientHandler /// - /// FinishReasonNormalizingHandler is outermost so it transforms the SSE response - /// AFTER LoggingHttpHandler's retry logic has delivered the final response. - /// This ensures Gemini's non-OpenAI finish_reason values (STOP, MAX_TOKENS, - /// SAFETY, etc.) are mapped to the OpenAI SDK's expected set before deserialization. + /// FinishReasonNormalizingHandler sits above LoggingHttpHandler so it transforms the + /// SSE response AFTER the retry logic has delivered the final response. This ensures + /// Gemini's non-OpenAI finish_reason values (STOP, MAX_TOKENS, SAFETY, etc.) are + /// mapped to the OpenAI SDK's expected set before deserialization. + /// + /// ThoughtSignatureHandler is outermost so it sees the request before it is sent and + /// the response after everything below has settled. It carries Gemini 3's + /// thought_signature from one turn to the next, which the OpenAI SDK would otherwise + /// drop, making the second call of any tool conversation fail with 400. /// private static HttpClient CreateHttpClient() { var handler = new FinishReasonNormalizingHandler( new LoggingHttpHandler( new HttpClientHandler())); - return new HttpClient(handler) + return new HttpClient(new ThoughtSignatureHandler(handler)) { Timeout = TimeSpan.FromSeconds(300) }; diff --git a/src/OpenDeepWiki/Agents/ThoughtSignatureHandler.cs b/src/OpenDeepWiki/Agents/ThoughtSignatureHandler.cs new file mode 100644 index 00000000..bb43d338 --- /dev/null +++ b/src/OpenDeepWiki/Agents/ThoughtSignatureHandler.cs @@ -0,0 +1,395 @@ +using System.Net; +using System.Text; +using System.Text.Json; +using System.Text.Json.Nodes; + +namespace OpenDeepWiki.Agents; + +/// +/// A that round-trips Gemini's thought_signature +/// across the turns of a tool-calling conversation. +/// +/// Gemini 3 models reason before they act. When that reasoning produces a function call, +/// the response carries an opaque signature at +/// choices[].message.tool_calls[].extra_content.google.thought_signature, and the API +/// requires it to be sent back with that same tool call on the following turn. The reasoning +/// itself is never returned as text, so the signature is the only thing carrying it forward; +/// without it the request is rejected outright: +/// +/// 400 INVALID_ARGUMENT - "Function call is missing a thought_signature in functionCall +/// parts. This is required for tools to work correctly." +/// +/// The field lives outside the OpenAI schema, so the OpenAI .NET SDK and +/// Microsoft.Extensions.AI drop it while remapping the response onto their own model, and +/// the follow-up request is rebuilt without it. The first call of a conversation therefore +/// succeeds and the second fails - which shows up as an agent dying on its first tool result. +/// +/// This handler closes the gap at the HTTP layer, where the raw JSON is still intact: it +/// records each signature as the response streams past, keyed by tool call id, and restores +/// it on any later request that references the same call. Responses are only read, never +/// modified, and only requests bound for Gemini's OpenAI-compatible endpoint are touched. +/// +/// This mirrors , which already corrects a +/// different Gemini/OpenAI mismatch on the same layer. +/// +public sealed class ThoughtSignatureHandler : DelegatingHandler +{ + private const string GeminiHost = "generativelanguage.googleapis.com"; + + /// + /// Upper bound on remembered signatures. A signature runs to several kilobytes and + /// AgentFactory builds a fresh per agent, so this only has + /// to span a single conversation; the cap is a safety net against a very long one. + /// + private const int MaxCachedSignatures = 512; + + private static readonly Serilog.ILogger Logger = Serilog.Log.ForContext(); + + private readonly Dictionary _signatures = new(StringComparer.Ordinal); + private readonly Queue _insertionOrder = new(); + private readonly object _gate = new(); + + public ThoughtSignatureHandler(HttpMessageHandler innerHandler) + : base(innerHandler) + { + } + + protected override async Task SendAsync( + HttpRequestMessage request, + CancellationToken cancellationToken) + { + var isGemini = request.RequestUri?.Host.EndsWith(GeminiHost, StringComparison.OrdinalIgnoreCase) == true; + + if (isGemini) + { + await TryRestoreSignaturesAsync(request, cancellationToken); + } + + var response = await base.SendAsync(request, cancellationToken); + + if (!isGemini || !response.IsSuccessStatusCode) + { + return response; + } + + try + { + var mediaType = response.Content.Headers.ContentType?.MediaType; + + if (string.Equals(mediaType, "text/event-stream", StringComparison.OrdinalIgnoreCase)) + { + response.Content = WrapSseContent(response.Content); + } + else if (string.Equals(mediaType, "application/json", StringComparison.OrdinalIgnoreCase)) + { + // Non-streaming completion: buffer it, harvest, hand back an equivalent body. + var json = await response.Content.ReadAsStringAsync(cancellationToken); + RecordSignatures(json); + + var replacement = new StringContent(json, Encoding.UTF8); + foreach (var header in response.Content.Headers) + { + replacement.Headers.TryAddWithoutValidation(header.Key, header.Value); + } + + response.Content = replacement; + } + } + catch (Exception ex) + { + // Never break the call - fall through with the original response. + Logger.Warning(ex, "ThoughtSignatureHandler: failed to record signatures; response left untouched."); + } + + return response; + } + + /// + /// Rewrites the outgoing body so every tool call we hold a signature for carries it again. + /// The body is replaced with a fresh , which stays re-readable + /// and therefore survives retries performed further down the handler chain. + /// + private async Task TryRestoreSignaturesAsync(HttpRequestMessage request, CancellationToken cancellationToken) + { + try + { + if (request.Content is null || CachedCount == 0) + { + return; + } + + var body = await request.Content.ReadAsStringAsync(cancellationToken); + if (body.Length == 0 || !body.Contains("\"tool_calls\"", StringComparison.Ordinal)) + { + return; + } + + if (JsonNode.Parse(body) is not JsonObject root) + { + return; + } + + var restored = RestoreSignatures(root, TryGetSignature); + if (restored == 0) + { + return; + } + + var rewritten = new StringContent(root.ToJsonString(), Encoding.UTF8, "application/json"); + foreach (var header in request.Content.Headers) + { + if (!string.Equals(header.Key, "Content-Type", StringComparison.OrdinalIgnoreCase) && + !string.Equals(header.Key, "Content-Length", StringComparison.OrdinalIgnoreCase)) + { + rewritten.Headers.TryAddWithoutValidation(header.Key, header.Value); + } + } + + request.Content = rewritten; + + Logger.Debug("ThoughtSignatureHandler: restored {Count} thought_signature value(s) on the request.", restored); + } + catch (Exception ex) + { + Logger.Warning(ex, "ThoughtSignatureHandler: failed to restore signatures; request left untouched."); + } + } + + private HttpContent WrapSseContent(HttpContent original) + { + var observing = new ObservingStreamContent(original, this); + + foreach (var header in original.Headers) + { + observing.Headers.TryAddWithoutValidation(header.Key, header.Value); + } + + return observing; + } + + /// + /// Attaches a cached signature to every tool call in that + /// names one and does not already carry extra_content. Returns how many were added. + /// + /// Exposed as internal static so unit tests can call it without an HTTP stack. + /// + internal static int RestoreSignatures(JsonObject requestRoot, Func lookup) + { + if (requestRoot["messages"] is not JsonArray messages) + { + return 0; + } + + var restored = 0; + + foreach (var message in messages.OfType()) + { + if (message["tool_calls"] is not JsonArray toolCalls) + { + continue; + } + + foreach (var toolCall in toolCalls.OfType()) + { + // Leave anything the caller already supplied alone. + if (toolCall["extra_content"] is not null) + { + continue; + } + + var id = toolCall["id"]?.GetValue(); + if (id is null) + { + continue; + } + + var signature = lookup(id); + if (signature is null) + { + continue; + } + + toolCall["extra_content"] = new JsonObject + { + ["google"] = new JsonObject + { + ["thought_signature"] = signature + } + }; + + restored++; + } + } + + return restored; + } + + /// + /// Pulls every id/signature pair out of one SSE data line or one complete JSON completion + /// and remembers it. Streaming splits a tool call across chunks - the id arrives in + /// the first chunk for a given index and later chunks carry only that index - so + /// unresolved indices are matched through , which the caller + /// keeps for the lifetime of one response stream. + /// + /// Exposed as internal so unit tests can call it without an HTTP stack. + /// + internal void RecordSignatures(string payload, Dictionary? indexToId = null) + { + if (!payload.Contains("thought_signature", StringComparison.Ordinal) && + !payload.Contains("\"id\"", StringComparison.Ordinal)) + { + return; + } + + JsonObject? root; + try + { + root = JsonNode.Parse(payload) as JsonObject; + } + catch (JsonException) + { + // A chunk we cannot parse is never worth failing the response over. + return; + } + + if (root?["choices"] is not JsonArray choices) + { + return; + } + + foreach (var choice in choices.OfType()) + { + // "delta" for streaming chunks, "message" for a complete completion. + var container = choice["delta"] as JsonObject ?? choice["message"] as JsonObject; + + if (container?["tool_calls"] is not JsonArray toolCalls) + { + continue; + } + + foreach (var toolCall in toolCalls.OfType()) + { + var id = toolCall["id"]?.GetValue(); + var index = toolCall["index"] is JsonValue indexValue && indexValue.TryGetValue(out var i) + ? i + : (int?)null; + + if (id is not null && index is not null && indexToId is not null) + { + indexToId[index.Value] = id; + } + + var signature = toolCall["extra_content"]?["google"]?["thought_signature"]?.GetValue(); + if (signature is null) + { + continue; + } + + var resolvedId = id; + if (resolvedId is null && index is not null && indexToId is not null) + { + indexToId.TryGetValue(index.Value, out resolvedId); + } + + if (resolvedId is not null) + { + Remember(resolvedId, signature); + } + } + } + } + + internal int CachedCount + { + get + { + lock (_gate) + { + return _signatures.Count; + } + } + } + + private void Remember(string toolCallId, string signature) + { + lock (_gate) + { + if (!_signatures.TryAdd(toolCallId, signature)) + { + return; + } + + _insertionOrder.Enqueue(toolCallId); + + while (_insertionOrder.Count > MaxCachedSignatures) + { + _signatures.Remove(_insertionOrder.Dequeue()); + } + } + } + + internal string? TryGetSignature(string toolCallId) + { + lock (_gate) + { + return _signatures.TryGetValue(toolCallId, out var signature) ? signature : null; + } + } + + /// + /// Streams the original SSE body through untouched while reading each data line for + /// signatures. Nothing is buffered beyond the line in hand. + /// + private sealed class ObservingStreamContent(HttpContent inner, ThoughtSignatureHandler owner) : HttpContent + { + protected override async Task SerializeToStreamAsync(Stream stream, TransportContext? context) + { + var innerStream = await inner.ReadAsStreamAsync(); + await ObserveSseStreamAsync(innerStream, stream, CancellationToken.None); + } + + protected override async Task SerializeToStreamAsync( + Stream stream, + TransportContext? context, + CancellationToken cancellationToken) + { + var innerStream = await inner.ReadAsStreamAsync(cancellationToken); + await ObserveSseStreamAsync(innerStream, stream, cancellationToken); + } + + protected override bool TryComputeLength(out long length) + { + length = -1; + return false; + } + + private async Task ObserveSseStreamAsync(Stream source, Stream destination, CancellationToken cancellationToken) + { + var reader = new StreamReader(source, Encoding.UTF8, detectEncodingFromByteOrderMarks: false, + bufferSize: 4096, leaveOpen: true); + var writer = new StreamWriter(destination, new UTF8Encoding(encoderShouldEmitUTF8Identifier: false), + bufferSize: 4096, leaveOpen: true) { AutoFlush = true }; + + // Maps a streamed tool call's index to the id announced in its first chunk. + var indexToId = new Dictionary(); + + string? line; + while ((line = await reader.ReadLineAsync(cancellationToken)) is not null) + { + if (line.StartsWith("data:", StringComparison.Ordinal)) + { + var payload = line["data:".Length..].Trim(); + if (payload.Length > 0 && payload != "[DONE]") + { + owner.RecordSignatures(payload, indexToId); + } + } + + // Pass the line on exactly as it arrived. + await writer.WriteLineAsync(line.AsMemory(), cancellationToken); + } + + await writer.FlushAsync(cancellationToken); + } + } +} diff --git a/tests/OpenDeepWiki.Tests/Agents/ThoughtSignatureHandlerTests.cs b/tests/OpenDeepWiki.Tests/Agents/ThoughtSignatureHandlerTests.cs new file mode 100644 index 00000000..bc73916a --- /dev/null +++ b/tests/OpenDeepWiki.Tests/Agents/ThoughtSignatureHandlerTests.cs @@ -0,0 +1,224 @@ +using System.Text.Json.Nodes; +using OpenDeepWiki.Agents; +using Xunit; + +namespace OpenDeepWiki.Tests.Agents; + +/// +/// Unit tests for . +/// Exercises the internal helpers directly so no HTTP stack is required. +/// +public class ThoughtSignatureHandlerTests +{ + private static ThoughtSignatureHandler NewHandler() => new(new HttpClientHandler()); + + /// One SSE data payload carrying a tool call and its signature. + private static string StreamingChunk(string id, string signature, int index = 0) => + "{\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"tool_calls\":[" + + "{\"index\":" + index + ",\"id\":\"" + id + "\",\"type\":\"function\"," + + "\"function\":{\"name\":\"ReadFile\",\"arguments\":\"{}\"}," + + "\"extra_content\":{\"google\":{\"thought_signature\":\"" + signature + "\"}}}]}}]}"; + + private static JsonObject RequestWithToolCall(string id, bool withExtraContent = false) + { + var toolCall = new JsonObject + { + ["id"] = id, + ["type"] = "function", + ["function"] = new JsonObject { ["name"] = "ReadFile", ["arguments"] = "{}" } + }; + + if (withExtraContent) + { + toolCall["extra_content"] = new JsonObject + { + ["google"] = new JsonObject { ["thought_signature"] = "caller-supplied" } + }; + } + + return new JsonObject + { + ["model"] = "gemini-3-pro-preview", + ["messages"] = new JsonArray( + new JsonObject { ["role"] = "user", ["content"] = "read it" }, + new JsonObject + { + ["role"] = "assistant", + ["content"] = null, + ["tool_calls"] = new JsonArray(toolCall) + }, + new JsonObject { ["role"] = "tool", ["tool_call_id"] = id, ["content"] = "file body" }) + }; + } + + private static string? SignatureOf(JsonObject request, int messageIndex = 1, int toolCallIndex = 0) => + request["messages"]?[messageIndex]?["tool_calls"]?[toolCallIndex]? + ["extra_content"]?["google"]?["thought_signature"]?.GetValue(); + + // ------------------------------------------------------------------ + // RecordSignatures - reading them off the response + // ------------------------------------------------------------------ + + [Fact] + public void RecordSignatures_Reads_Streaming_Delta() + { + var handler = NewHandler(); + + handler.RecordSignatures(StreamingChunk("call_1", "sig-abc"), new Dictionary()); + + Assert.Equal("sig-abc", handler.TryGetSignature("call_1")); + } + + [Fact] + public void RecordSignatures_Reads_NonStreaming_Message() + { + var handler = NewHandler(); + const string payload = """ + {"choices":[{"message":{"role":"assistant","tool_calls":[ + {"id":"call_2","type":"function","function":{"name":"ReadFile","arguments":"{}"}, + "extra_content":{"google":{"thought_signature":"sig-xyz"}}}]}}]} + """; + + handler.RecordSignatures(payload); + + Assert.Equal("sig-xyz", handler.TryGetSignature("call_2")); + } + + [Fact] + public void RecordSignatures_Pairs_By_Index_When_Chunk_Omits_Id() + { + var handler = NewHandler(); + var indexToId = new Dictionary(); + + // First chunk announces the id but carries no signature yet. + handler.RecordSignatures( + """{"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_3","type":"function","function":{"name":"ReadFile"}}]}}]}""", + indexToId); + + // A later chunk carries the signature under the same index, with no id of its own. + handler.RecordSignatures( + """{"choices":[{"delta":{"tool_calls":[{"index":0,"extra_content":{"google":{"thought_signature":"sig-late"}}}]}}]}""", + indexToId); + + Assert.Equal("sig-late", handler.TryGetSignature("call_3")); + } + + [Fact] + public void RecordSignatures_Ignores_Payload_Without_Signature() + { + var handler = NewHandler(); + + handler.RecordSignatures("""{"choices":[{"delta":{"content":"hello"}}]}"""); + + Assert.Equal(0, handler.CachedCount); + } + + [Fact] + public void RecordSignatures_Does_Not_Throw_On_Malformed_Json() + { + var handler = NewHandler(); + + // A truncated chunk must never take the response down with it. + var ex = Record.Exception(() => + handler.RecordSignatures("""{"choices":[{"delta":{"tool_calls":[{"id":"call_4","extra_content" """)); + + Assert.Null(ex); + Assert.Equal(0, handler.CachedCount); + } + + // ------------------------------------------------------------------ + // RestoreSignatures - putting them back on the request + // ------------------------------------------------------------------ + + [Fact] + public void RestoreSignatures_Attaches_Known_Signature() + { + var request = RequestWithToolCall("call_5"); + + var restored = ThoughtSignatureHandler.RestoreSignatures( + request, id => id == "call_5" ? "sig-restored" : null); + + Assert.Equal(1, restored); + Assert.Equal("sig-restored", SignatureOf(request)); + } + + [Fact] + public void RestoreSignatures_Leaves_Unknown_Tool_Call_Untouched() + { + var request = RequestWithToolCall("call_6"); + + var restored = ThoughtSignatureHandler.RestoreSignatures(request, _ => null); + + Assert.Equal(0, restored); + Assert.Null(SignatureOf(request)); + } + + [Fact] + public void RestoreSignatures_Does_Not_Overwrite_Existing_ExtraContent() + { + var request = RequestWithToolCall("call_7", withExtraContent: true); + + var restored = ThoughtSignatureHandler.RestoreSignatures(request, _ => "sig-ours"); + + Assert.Equal(0, restored); + Assert.Equal("caller-supplied", SignatureOf(request)); + } + + [Fact] + public void RestoreSignatures_Handles_Request_Without_Messages() + { + var request = new JsonObject { ["model"] = "gemini-3-pro-preview" }; + + var restored = ThoughtSignatureHandler.RestoreSignatures(request, _ => "sig"); + + Assert.Equal(0, restored); + } + + // ------------------------------------------------------------------ + // Round trip and cache behaviour + // ------------------------------------------------------------------ + + [Fact] + public void Signature_Survives_A_Full_Round_Trip() + { + var handler = NewHandler(); + + // Response arrives carrying the signature... + handler.RecordSignatures(StreamingChunk("call_8", "sig-round-trip"), new Dictionary()); + + // ...and the next request, rebuilt by the OpenAI SDK without it, gets it back. + var request = RequestWithToolCall("call_8"); + var restored = ThoughtSignatureHandler.RestoreSignatures(request, handler.TryGetSignature); + + Assert.Equal(1, restored); + Assert.Equal("sig-round-trip", SignatureOf(request)); + } + + [Fact] + public void Cache_Is_Bounded_And_Evicts_Oldest_First() + { + var handler = NewHandler(); + const int cap = 512; + + for (var i = 0; i < cap + 10; i++) + { + handler.RecordSignatures(StreamingChunk($"call_{i}", $"sig-{i}"), new Dictionary()); + } + + Assert.Equal(cap, handler.CachedCount); + Assert.Null(handler.TryGetSignature("call_0")); // evicted + Assert.Equal($"sig-{cap + 9}", handler.TryGetSignature($"call_{cap + 9}")); // newest kept + } + + [Fact] + public void First_Signature_Wins_For_A_Repeated_Tool_Call_Id() + { + var handler = NewHandler(); + + handler.RecordSignatures(StreamingChunk("call_9", "sig-first"), new Dictionary()); + handler.RecordSignatures(StreamingChunk("call_9", "sig-second"), new Dictionary()); + + Assert.Equal("sig-first", handler.TryGetSignature("call_9")); + Assert.Equal(1, handler.CachedCount); + } +}