Carry Gemini thought signatures on the transcript

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)
                 )
             }
         }
@@ -613,7 +623,8 @@ private func resolveFunctionCalls(
             Transcript.ToolCall(
                 id: callID,
                 toolName: call.name,
-                arguments: args
+                arguments: args,
+                providerMetadata: call.thoughtSignature.map { [thoughtSignatureMetadataKey: $0] }
             )
         )
     }
@@ -786,7 +797,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 +912,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 +940,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 +994,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

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
         }
     }
 

Tests/AnyLanguageModelTests/GeminiThoughtSignatureTests.swift

@@ -0,0 +1,136 @@
+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)
+                }
+            }
+        }
+
+        /// 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()
+            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 })
+        }
+
+        @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

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