From ece9145d07c113c70db191a3756a4be3c8f5a8c4 Mon Sep 17 00:00:00 2001 From: NoWaY233851 <212446833+NoWaY233851@users.noreply.github.com> Date: Thu, 3 Sep 2026 23:22:38 -0400 Subject: [PATCH 1/2] Carry Gemini thought signatures on the transcript MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Gemini attaches a `thought_signature` to each `functionCall` part and rejects a request whose function calls have lost it: 400 INVALID_ARGUMENT "Function call is missing a thought_signature in functionCall parts." `GeminiPart` decoded the signature away, and `Transcript.ToolCall` had nowhere to keep it, so every request rebuilt from the transcript dropped it — the follow-up inside the tool loop, and the replay of that tool call on later turns. Give `Transcript.ToolCall` a `providerMetadata` dictionary for opaque state that a provider requires back verbatim, and have the Gemini adapter put the signature there. The field is optional, so transcripts encoded before this change still decode, and it is marked as an AnyLanguageModel extension rather than something carried over from Foundation Models. Other providers need the same shape: Anthropic's `encrypted_content` for server tool results, and encrypted reasoning items on the OpenAI Responses API. --- .../Models/GeminiLanguageModel.swift | 34 +++++- Sources/AnyLanguageModel/Transcript.swift | 18 ++- .../GeminiThoughtSignatureTests.swift | 106 ++++++++++++++++++ .../Shared/StubURLProtocol.swift | 92 +++++++++++++++ .../TranscriptTests.swift | 32 ++++++ 5 files changed, 277 insertions(+), 5 deletions(-) create mode 100644 Tests/AnyLanguageModelTests/GeminiThoughtSignatureTests.swift create mode 100644 Tests/AnyLanguageModelTests/Shared/StubURLProtocol.swift diff --git a/Sources/AnyLanguageModel/Models/GeminiLanguageModel.swift b/Sources/AnyLanguageModel/Models/GeminiLanguageModel.swift index b074ddca..14ee82a5 100644 --- a/Sources/AnyLanguageModel/Models/GeminiLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/GeminiLanguageModel.swift @@ -613,7 +613,8 @@ private func resolveFunctionCalls( Transcript.ToolCall( id: callID, toolName: call.name, - arguments: args + arguments: args, + providerMetadata: call.thoughtSignature.map { [thoughtSignatureMetadataKey: $0] } ) ) } @@ -786,7 +787,13 @@ extension Transcript { // Add model's response with function calls let functionCallParts: [GeminiPart] = toolCalls.map { call in let args = try? fromGeneratedContent(call.arguments) - return .functionCall(GeminiFunctionCall(name: call.toolName, args: args)) + return .functionCall( + GeminiFunctionCall( + name: call.toolName, + args: args, + thoughtSignature: call.providerMetadata?[thoughtSignatureMetadataKey] + ) + ) } messages.append( .init( @@ -895,8 +902,11 @@ private enum GeminiPart: Codable, Sendable { let text = try container.decode(String.self, forKey: .text) self = .text(GeminiTextPart(text: text)) } else if container.contains(.functionCall) { - // Note: thoughtSignature may be present but is ignored - self = .functionCall(try container.decode(GeminiFunctionCall.self, forKey: .functionCall)) + // `thoughtSignature` is a sibling of `functionCall` within the part, not a member of it. + // Thinking models require it to be echoed back verbatim, so it travels with the call. + var call = try container.decode(GeminiFunctionCall.self, forKey: .functionCall) + call.thoughtSignature = try container.decodeIfPresent(String.self, forKey: .thoughtSignature) + self = .functionCall(call) } else if container.contains(.functionResponse) { self = .functionResponse(try container.decode(GeminiFunctionResponse.self, forKey: .functionResponse)) } else if container.contains(.inlineData) { @@ -920,6 +930,7 @@ private enum GeminiPart: Codable, Sendable { try container.encode(part.text, forKey: .text) case .functionCall(let call): try container.encode(call, forKey: .functionCall) + try container.encodeIfPresent(call.thoughtSignature, forKey: .thoughtSignature) case .functionResponse(let response): try container.encode(response, forKey: .functionResponse) case .inlineData(let data): @@ -973,10 +984,25 @@ private func convertSegmentsToGeminiParts(_ segments: [Transcript.Segment]) -> [ return parts } +/// Key under which a thought signature is kept in ``Transcript/ToolCall/providerMetadata``. +private let thoughtSignatureMetadataKey = "thoughtSignature" + private struct GeminiFunctionCall: Codable, Sendable { let name: String let args: [String: JSONValue]? + /// The opaque thought signature Gemini attached to this call, if any. + /// + /// Deliberately absent from ``CodingKeys``: on the wire the signature sits on the enclosing + /// part, so ``GeminiPart`` is what reads and writes it. + var thoughtSignature: String? + + init(name: String, args: [String: JSONValue]?, thoughtSignature: String? = nil) { + self.name = name + self.args = args + self.thoughtSignature = thoughtSignature + } + enum CodingKeys: String, CodingKey { case name case args diff --git a/Sources/AnyLanguageModel/Transcript.swift b/Sources/AnyLanguageModel/Transcript.swift index dcd1fadc..54d32da5 100644 --- a/Sources/AnyLanguageModel/Transcript.swift +++ b/Sources/AnyLanguageModel/Transcript.swift @@ -332,10 +332,26 @@ public struct Transcript: Sendable, Equatable, Codable { /// Arguments to pass to the invoked tool. public var arguments: GeneratedContent - public init(id: String, toolName: String, arguments: GeneratedContent) { + /// Opaque, provider-specific state attached to this call. + /// + /// Some providers hand back values alongside a tool call that must be echoed + /// verbatim on every later request — Gemini's thought signatures, for instance. + /// The values are meaningful only to the model that produced them; nothing else + /// should interpret or modify them. + /// + /// - Note: This property is exclusive to AnyLanguageModel + public var providerMetadata: [String: String]? + + public init( + id: String, + toolName: String, + arguments: GeneratedContent, + providerMetadata: [String: String]? = nil + ) { self.id = id self.toolName = toolName self.arguments = arguments + self.providerMetadata = providerMetadata } } diff --git a/Tests/AnyLanguageModelTests/GeminiThoughtSignatureTests.swift b/Tests/AnyLanguageModelTests/GeminiThoughtSignatureTests.swift new file mode 100644 index 00000000..ba3e8f7b --- /dev/null +++ b/Tests/AnyLanguageModelTests/GeminiThoughtSignatureTests.swift @@ -0,0 +1,106 @@ +import Foundation +import Testing + +@testable import AnyLanguageModel + +#if canImport(Darwin) && !canImport(AsyncHTTPClient) + + @Suite("GeminiLanguageModel thought signatures", .serialized) + struct GeminiThoughtSignatureTests { + private static let signature = "CvsBAdHtim8n5xQK1pVX2H0lPQeXAMPLEsignature==" + + private func makeModel() -> GeminiLanguageModel { + GeminiLanguageModel( + apiKey: "test-key", + model: "gemini-3.6-flash", + session: StubURLProtocol.makeSession() + ) + } + + private func functionCallResponse(signature: String?) -> String { + let signatureField = signature.map { ", \"thoughtSignature\": \"\($0)\"" } ?? "" + return """ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "functionCall": { "name": "getWeather", "args": { "city": "Paris" } }\(signatureField) + } + ] + }, + "finishReason": "STOP" + } + ] + } + """ + } + + private func textResponse(_ text: String) -> String { + """ + { + "candidates": [ + { + "content": { "role": "model", "parts": [{ "text": "\(text)" }] }, + "finishReason": "STOP" + } + ] + } + """ + } + + /// Signatures of every `functionCall` part in a request body, in order. + /// A part without a signature contributes `nil`. + private func functionCallSignatures(in body: Data) throws -> [String?] { + let json = try JSONSerialization.jsonObject(with: body) as? [String: Any] + let contents = json?["contents"] as? [[String: Any]] ?? [] + return contents.flatMap { content -> [String?] in + let parts = content["parts"] as? [[String: Any]] ?? [] + return parts.compactMap { part -> String?? in + guard part["functionCall"] != nil else { return nil } + return .some(part["thoughtSignature"] as? String) + } + } + } + + @Test("echoes the thought signature back with the function results") + func echoesSignatureOnFollowUpRequest() async throws { + StubURLProtocol.reset() + StubURLProtocol.enqueue(json: functionCallResponse(signature: Self.signature)) + StubURLProtocol.enqueue(json: textResponse("It is sunny in Paris.")) + + let session = LanguageModelSession(model: makeModel(), tools: [WeatherTool()]) + let response = try await session.respond(to: "What is the weather in Paris?") + + #expect(response.content == "It is sunny in Paris.") + + let bodies = StubURLProtocol.recordedBodies + try #require(bodies.count == 2) + #expect(try functionCallSignatures(in: bodies[1]) == [Self.signature]) + } + + @Test("keeps the signature on the tool call when the conversation continues") + func keepsSignatureOnLaterTurn() async throws { + StubURLProtocol.reset() + StubURLProtocol.enqueue(json: functionCallResponse(signature: Self.signature)) + StubURLProtocol.enqueue(json: textResponse("It is sunny in Paris.")) + StubURLProtocol.enqueue(json: textResponse("You asked about the weather in Paris.")) + + let session = LanguageModelSession(model: makeModel(), tools: [WeatherTool()]) + _ = try await session.respond(to: "What is the weather in Paris?") + _ = try await session.respond(to: "What did I just ask about?") + + let bodies = StubURLProtocol.recordedBodies + try #require(bodies.count == 3) + + // The tool call is replayed as history on the next turn, and Gemini rejects a request + // whose functionCall parts have lost their signatures. + let signatures = try functionCallSignatures(in: bodies[2]) + #expect(!signatures.isEmpty) + #expect(signatures.allSatisfy { $0 == Self.signature }) + } + } + +#endif diff --git a/Tests/AnyLanguageModelTests/Shared/StubURLProtocol.swift b/Tests/AnyLanguageModelTests/Shared/StubURLProtocol.swift new file mode 100644 index 00000000..49da63fc --- /dev/null +++ b/Tests/AnyLanguageModelTests/Shared/StubURLProtocol.swift @@ -0,0 +1,92 @@ +import Foundation + +@testable import AnyLanguageModel + +#if canImport(Darwin) && !canImport(AsyncHTTPClient) + + /// A `URLProtocol` that answers requests from a queue of canned responses and records + /// every request body it sees, so request/response round trips can be asserted offline. + final class StubURLProtocol: URLProtocol { + struct Exchange: Sendable { + var statusCode: Int = 200 + var body: Data + } + + private struct State: Sendable { + var pending: [Exchange] = [] + var recordedBodies: [Data] = [] + } + + private static let state = Locked(State()) + + /// Discards queued responses and recorded bodies. + static func reset() { + state.withLock { $0 = State() } + } + + /// Queues one JSON response, returned to the next request that arrives. + static func enqueue(json: String, statusCode: Int = 200) { + state.withLock { $0.pending.append(Exchange(statusCode: statusCode, body: Data(json.utf8))) } + } + + /// The bodies of the requests seen so far, in order. + static var recordedBodies: [Data] { + state.withLock { $0.recordedBodies } + } + + /// A session that routes every request to this protocol. + static func makeSession() -> URLSession { + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [StubURLProtocol.self] + return URLSession(configuration: configuration) + } + + override class func canInit(with request: URLRequest) -> Bool { true } + + override class func canonicalRequest(for request: URLRequest) -> URLRequest { request } + + override func startLoading() { + // URLSession moves `httpBody` to `httpBodyStream` before the protocol sees the request. + let body = request.httpBody ?? request.httpBodyStream.map(Self.readAll) ?? Data() + + let exchange = Self.state.withLock { state -> Exchange? in + state.recordedBodies.append(body) + return state.pending.isEmpty ? nil : state.pending.removeFirst() + } + + guard let exchange, let url = request.url else { + client?.urlProtocol(self, didFailWithError: URLError(.resourceUnavailable)) + return + } + + let response = HTTPURLResponse( + url: url, + statusCode: exchange.statusCode, + httpVersion: "HTTP/1.1", + headerFields: ["Content-Type": "application/json"] + )! + + client?.urlProtocol(self, didReceive: response, cacheStoragePolicy: .notAllowed) + client?.urlProtocol(self, didLoad: exchange.body) + client?.urlProtocolDidFinishLoading(self) + } + + override func stopLoading() {} + + private static func readAll(_ stream: InputStream) -> Data { + stream.open() + defer { stream.close() } + + var data = Data() + let bufferSize = 4096 + var buffer = [UInt8](repeating: 0, count: bufferSize) + while true { + let read = stream.read(&buffer, maxLength: bufferSize) + if read <= 0 { break } + data.append(buffer, count: read) + } + return data + } + } + +#endif diff --git a/Tests/AnyLanguageModelTests/TranscriptTests.swift b/Tests/AnyLanguageModelTests/TranscriptTests.swift index 037c70a0..4f3eb65e 100644 --- a/Tests/AnyLanguageModelTests/TranscriptTests.swift +++ b/Tests/AnyLanguageModelTests/TranscriptTests.swift @@ -137,4 +137,36 @@ struct TranscriptTests { ) #expect(firstToolDefinition == secondToolDefinition) } + @Test func toolCallOmitsProviderMetadataWhenAbsent() throws { + // Guards decoding of transcripts encoded before `providerMetadata` existed. Making the + // property non-optional breaks this: synthesized `Codable` demands a key for every + // non-optional property and never consults its default value. + let arguments = try GeneratedContent(json: #"{"city":"Cupertino"}"#) + let call = Transcript.ToolCall(id: "call-id", toolName: "getWeather", arguments: arguments) + + let data = try JSONEncoder().encode(call) + var object = try #require(try JSONSerialization.jsonObject(with: data) as? [String: Any]) + #expect(object["providerMetadata"] == nil) + + object.removeValue(forKey: "providerMetadata") + let encodedBeforeTheFieldExisted = try JSONSerialization.data(withJSONObject: object) + let decoded = try JSONDecoder().decode(Transcript.ToolCall.self, from: encodedBeforeTheFieldExisted) + #expect(decoded.providerMetadata == nil) + #expect(decoded.toolName == "getWeather") + } + + @Test func toolCallRoundTripsProviderMetadata() throws { + let arguments = try GeneratedContent(json: #"{"city":"Cupertino"}"#) + let call = Transcript.ToolCall( + id: "call-id", + toolName: "getWeather", + arguments: arguments, + providerMetadata: ["thoughtSignature": "opaque-signature"] + ) + + let data = try JSONEncoder().encode(call) + let decoded = try JSONDecoder().decode(Transcript.ToolCall.self, from: data) + + #expect(decoded.providerMetadata == ["thoughtSignature": "opaque-signature"]) + } } From a7ca2162fc227bfb64923c5fe1ec6eb0807c6488 Mon Sep 17 00:00:00 2001 From: NoWaY233851 <212446833+NoWaY233851@users.noreply.github.com> Date: Fri, 4 Sep 2026 08:34:02 -0400 Subject: [PATCH 2/2] Report only the entries this response added MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `GeminiLanguageModel.respond` returned `transcriptEntries: ArraySlice(transcript)` — the whole conversation rather than the entries it appended — where the other adapters accumulate a local array. `LanguageModelSession.respond` appends that back into the session transcript, so history doubled on every turn: request `contents` went 1 -> 3 -> 6 -> 14 over three turns. Until thought signatures were echoed, that was a token cost. It now also replays a signed `functionCall` part at a position its signature was never issued for, so the signature fix does not hold past the second turn without this. Collect the tool call and tool output entries in a local `entries` array, as `AnthropicLanguageModel` and the other adapters do, and report that. `transcript` still carries the whole conversation, because each iteration of the tool loop rebuilds the request from it. --- .../Models/GeminiLanguageModel.swift | 22 ++++++++++---- .../GeminiThoughtSignatureTests.swift | 30 +++++++++++++++++++ 2 files changed, 46 insertions(+), 6 deletions(-) diff --git a/Sources/AnyLanguageModel/Models/GeminiLanguageModel.swift b/Sources/AnyLanguageModel/Models/GeminiLanguageModel.swift index 14ee82a5..9f7d25ff 100644 --- a/Sources/AnyLanguageModel/Models/GeminiLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/GeminiLanguageModel.swift @@ -282,6 +282,10 @@ public struct GeminiLanguageModel: LanguageModel { var transcript = session.transcript + // The entries this call adds, which is what the response reports. `transcript` keeps the + // full conversation because each iteration rebuilds the request from it. + var entries: [Transcript.Entry] = [] + // Multi-turn conversation loop for tool calling while true { let params = try createGenerateContentParams( @@ -318,20 +322,26 @@ public struct GeminiLanguageModel: LanguageModel { switch resolution { case .stop(let calls): if !calls.isEmpty { - transcript.append(.toolCalls(Transcript.ToolCalls(calls))) + entries.append(.toolCalls(Transcript.ToolCalls(calls))) } let empty = try emptyResponseContent(for: type) return LanguageModelSession.Response( content: empty.content, rawContent: empty.rawContent, - transcriptEntries: ArraySlice(transcript) + transcriptEntries: ArraySlice(entries) ) case .invocations(let invocations): if !invocations.isEmpty { - transcript.append(.toolCalls(Transcript.ToolCalls(invocations.map(\.call)))) + let calls = Transcript.Entry.toolCalls( + Transcript.ToolCalls(invocations.map(\.call)) + ) + transcript.append(calls) + entries.append(calls) for invocation in invocations { - transcript.append(.toolOutput(invocation.output)) + let output = Transcript.Entry.toolOutput(invocation.output) + transcript.append(output) + entries.append(output) } } @@ -352,7 +362,7 @@ public struct GeminiLanguageModel: LanguageModel { return LanguageModelSession.Response( content: text as! Content, rawContent: GeneratedContent(text), - transcriptEntries: ArraySlice(transcript) + transcriptEntries: ArraySlice(entries) ) } @@ -361,7 +371,7 @@ public struct GeminiLanguageModel: LanguageModel { return LanguageModelSession.Response( content: content, rawContent: generatedContent, - transcriptEntries: ArraySlice(transcript) + transcriptEntries: ArraySlice(entries) ) } } diff --git a/Tests/AnyLanguageModelTests/GeminiThoughtSignatureTests.swift b/Tests/AnyLanguageModelTests/GeminiThoughtSignatureTests.swift index ba3e8f7b..c0f4c3ef 100644 --- a/Tests/AnyLanguageModelTests/GeminiThoughtSignatureTests.swift +++ b/Tests/AnyLanguageModelTests/GeminiThoughtSignatureTests.swift @@ -65,6 +65,12 @@ import Testing } } + /// The number of `contents` entries in a request body. + private func contentCount(in body: Data) throws -> Int { + let json = try JSONSerialization.jsonObject(with: body) as? [String: Any] + return (json?["contents"] as? [[String: Any]] ?? []).count + } + @Test("echoes the thought signature back with the function results") func echoesSignatureOnFollowUpRequest() async throws { StubURLProtocol.reset() @@ -101,6 +107,30 @@ import Testing #expect(!signatures.isEmpty) #expect(signatures.allSatisfy { $0 == Self.signature }) } + + @Test("does not replay the conversation history on later turns") + func doesNotReplayHistoryOnLaterTurns() async throws { + StubURLProtocol.reset() + StubURLProtocol.enqueue(json: functionCallResponse(signature: Self.signature)) + StubURLProtocol.enqueue(json: textResponse("It is sunny in Paris.")) + StubURLProtocol.enqueue(json: textResponse("You asked about the weather.")) + StubURLProtocol.enqueue(json: textResponse("Paris, specifically.")) + + let session = LanguageModelSession(model: makeModel(), tools: [WeatherTool()]) + _ = try await session.respond(to: "What is the weather in Paris?") + _ = try await session.respond(to: "What did I just ask about?") + _ = try await session.respond(to: "Which city was that?") + + let bodies = StubURLProtocol.recordedBodies + try #require(bodies.count == 4) + + // One entry per turn, plus the tool call and its output. A response that reported the + // whole transcript instead of its own entries would compound: 1, 3, 6, 14. + #expect(try bodies.map { try contentCount(in: $0) } == [1, 3, 5, 7]) + + // The tool call is replayed once, still signed, however many turns later. + #expect(try functionCallSignatures(in: bodies[3]) == [Self.signature]) + } } #endif