Preserve attachments in tool outputs

README.md

@@ -406,9 +406,10 @@ they throw an error when a prompt has an attachment.
 The system bridge requires local file URLs; other provider adapters retain URL
 images in their existing transcript representation.
 
-This first compatibility surface covers prompt `CGImage` and `imageURL`
-attachments. Tool outputs, instruction attachments, labels, `CIImage`, pixel
-buffers, and newer data attachments remain outside this change.
+This compatibility surface covers prompt `CGImage` and `imageURL` attachments,
+including attachments returned from ordinary `Tool` output values. Instruction
+attachments, labels, `CIImage`, pixel buffers, and newer data attachments remain
+outside this change.
 
 ### Tool Calling
 

Sources/AnyLanguageModel/Models/GeminiLanguageModel.swift

@@ -773,26 +773,47 @@ private func emptyResponseContent<Content: Generable>(
     }
 }
 
-private func toJSONValue(_ toolOutput: Transcript.ToolOutput) throws -> [String: JSONValue] {
-    var result: [String: JSONValue] = [:]
+private func toGeminiFunctionResponse(
+    _ toolOutput: Transcript.ToolOutput
+) throws -> GeminiFunctionResponse {
+    var resultValues: [String] = []
+    var parts: [GeminiPart] = []
 
     for segment in toolOutput.segments {
         switch segment {
         case .text(let text):
-            result["result"] = .string(text.content)
+            resultValues.append(text.content)
         case .structure(let structured):
-            // For structured segments, encode the content
             let data = try JSONEncoder().encode(structured.content)
             if let jsonString = String(data: data, encoding: .utf8) {
-                result["result"] = .string(jsonString)
+                resultValues.append(jsonString)
+            }
+        case .image(let image):
+            switch image.source {
+            case .data(let data, let mimeType):
+                parts.append(
+                    .inlineData(
+                        GeminiInlineData(
+                            mimeType: mimeType,
+                            data: data.base64EncodedString()
+                        )
+                    )
+                )
+            case .url(let url):
+                parts.append(.fileData(GeminiFileData(fileURI: url.absoluteString)))
             }
-        case .image:
-            // Ignore images in tool outputs for Gemini conversion
-            break
         }
     }
 
-    return result
+    var response: [String: JSONValue] = [:]
+    if !resultValues.isEmpty {
+        response["result"] = .string(resultValues.joined(separator: "\n"))
+    }
+    return GeminiFunctionResponse(
+        name: toolOutput.toolName,
+        response: response,
+        parts: parts.isEmpty ? nil : parts
+    )
 }
 
 // MARK: - Supporting Types
@@ -848,11 +869,13 @@ extension Transcript {
                 )
             case .toolOutput(let toolOutput):
                 // Add function response as a user message (Gemini API expects function responses from user role)
-                let response = try? toJSONValue(toolOutput)
-                let functionResponse = GeminiFunctionResponse(
-                    name: toolOutput.toolName,
-                    response: response ?? [:]
-                )
+                let functionResponse =
+                    (try? toGeminiFunctionResponse(toolOutput))
+                    ?? GeminiFunctionResponse(
+                        name: toolOutput.toolName,
+                        response: [:],
+                        parts: nil
+                    )
                 messages.append(
                     .init(
                         role: .user,
@@ -1120,6 +1143,7 @@ private struct GeminiFunctionCall: Codable, Sendable {
 private struct GeminiFunctionResponse: Codable, Sendable {
     let name: String
     let response: [String: JSONValue]
+    let parts: [GeminiPart]?
 }
 
 private struct GeminiGenerateContentResponse: Codable, Sendable {

Sources/AnyLanguageModel/Models/SystemLanguageModel.swift

@@ -577,7 +577,7 @@
     @available(macOS 26.0, iOS 26.0, watchOS 27.0, tvOS 26.0, visionOS 26.0, *)
     private struct AnyToolWrapper: FoundationModels.Tool {
         typealias Arguments = FoundationModels.GeneratedContent
-        typealias Output = String
+        typealias Output = FoundationModels.Prompt
 
         let name: String
         let description: String
@@ -596,7 +596,7 @@
 
         func call(arguments: FoundationModels.GeneratedContent) async throws -> Output {
             let output = try await wrappedTool.callFunction(arguments: arguments)
-            return output.promptRepresentation.description
+            return try output.promptRepresentation.toFoundationModels()
         }
     }
 

Sources/AnyLanguageModel/Tool.swift

@@ -110,7 +110,10 @@ extension Tool where Self.Arguments: Generable {
 // MARK: - Helpers
 
 extension Tool {
-    internal func makeOutputSegments(from arguments: GeneratedContent) async throws -> [Transcript.Segment] {
+    /// Compatibility integration SPI for provider adapters outside this module.
+    /// This is not part of Apple's public Tool contract.
+    @_spi(Compatibility)
+    public func makeOutputSegments(from arguments: GeneratedContent) async throws -> [Transcript.Segment] {
         let parsedArguments = try Arguments(arguments)
         let output = try await call(arguments: parsedArguments)
 
@@ -124,7 +127,6 @@ extension Tool {
             return [Transcript.Segment.text(.init(content: stringOutput))]
         }
 
-        let fallback = output.promptRepresentation.description
-        return [Transcript.Segment.text(.init(content: fallback))]
+        return try output.promptRepresentation.makeTranscriptSegments()
     }
 }

Tests/AnyLanguageModelTests/DynamicInstructionsTests.swift

@@ -1,7 +1,7 @@
 import Foundation
 import Testing
 
-@testable import AnyLanguageModel
+@_spi(Compatibility) @testable import AnyLanguageModel
 
 @Suite("Dynamic instructions")
 struct DynamicInstructionsTests {

Tests/AnyLanguageModelTests/GeminiThoughtSignatureTests.swift

@@ -4,11 +4,27 @@ import Testing
 @testable import AnyLanguageModel
 
 #if canImport(Darwin) && !canImport(AsyncHTTPClient)
+    import CoreGraphics
 
     @Suite("GeminiLanguageModel thought signatures", .serialized)
     struct GeminiThoughtSignatureTests {
         private static let signature = "CvsBAdHtim8n5xQK1pVX2H0lPQeXAMPLEsignature=="
 
+        @Generable
+        struct ImageArguments {
+            var request: String
+        }
+
+        private struct ImageOutputTool: Tool {
+            let name = "image"
+            let description = "Return an image"
+            let output: Prompt
+
+            func call(arguments: ImageArguments) async throws -> Prompt {
+                output
+            }
+        }
+
         private func makeModel() -> GeminiLanguageModel {
             GeminiLanguageModel(
                 apiKey: "test-key",
@@ -51,6 +67,21 @@ import Testing
             """
         }
 
+        private func image() throws -> CGImage {
+            let context = try #require(
+                CGContext(
+                    data: nil,
+                    width: 2,
+                    height: 2,
+                    bitsPerComponent: 8,
+                    bytesPerRow: 8,
+                    space: CGColorSpaceCreateDeviceRGB(),
+                    bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue
+                )
+            )
+            return try #require(context.makeImage())
+        }
+
         /// 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?] {
@@ -108,6 +139,46 @@ import Testing
             #expect(signatures.allSatisfy { $0 == Self.signature })
         }
 
+        @Test("preserves all text and image data in function responses")
+        func preservesSegmentedToolOutput() async throws {
+            StubURLProtocol.reset()
+            StubURLProtocol.enqueue(
+                json:
+                    #"{"candidates":[{"content":{"role":"model","parts":[{"functionCall":{"name":"image","args":{"request":"read"}}}]},"finishReason":"STOP"}]}"#
+            )
+            StubURLProtocol.enqueue(json: textResponse("OK"))
+
+            let output = try Prompt {
+                "before"
+                Attachment(try image())
+                "after"
+            }
+            let session = LanguageModelSession(
+                model: makeModel(),
+                tools: [ImageOutputTool(output: output)]
+            )
+            _ = try await session.respond(to: "read image")
+
+            let bodies = StubURLProtocol.recordedBodies
+            try #require(bodies.count == 2)
+            let json = try #require(
+                try JSONSerialization.jsonObject(with: bodies[1]) as? [String: Any]
+            )
+            let contents = try #require(json["contents"] as? [[String: Any]])
+            let outputContent = try #require(contents.last)
+            let outputParts = try #require(outputContent["parts"] as? [[String: Any]])
+            let functionResponse = try #require(
+                outputParts.first?["functionResponse"] as? [String: Any]
+            )
+            let response = try #require(functionResponse["response"] as? [String: Any])
+            #expect(response["result"] as? String == "before\nafter")
+
+            let mediaParts = try #require(functionResponse["parts"] as? [[String: Any]])
+            let inlineData = try #require(mediaParts.first?["inlineData"] as? [String: Any])
+            #expect(inlineData["mime_type"] as? String == "image/png")
+            #expect(Data(base64Encoded: try #require(inlineData["data"] as? String)) != nil)
+        }
+
         @Test("does not replay the conversation history on later turns")
         func doesNotReplayHistoryOnLaterTurns() async throws {
             StubURLProtocol.reset()

Tests/AnyLanguageModelTests/PromptAttachmentTests.swift

@@ -1,7 +1,7 @@
 import Foundation
 import Testing
 
-@testable import AnyLanguageModel
+@_spi(Compatibility) @testable import AnyLanguageModel
 
 #if canImport(CoreGraphics) && canImport(ImageIO)
     import CoreGraphics
@@ -122,6 +122,29 @@ struct PromptAttachmentTests {
         }
     }
 
+    @Test func ordinaryToolReturnsOrderedImageSegments() async throws {
+        let segments = try await ImageTool(output: prompt()).makeOutputSegments(
+            from: AnyLanguageModel.GeneratedContent("read")
+        )
+        #expect(segments.count == 3)
+        guard case .text(let before) = segments[0],
+            case .image(let image) = segments[1],
+            case .text(let after) = segments[2]
+        else {
+            Issue.record("Tool output lost its ordered image")
+            return
+        }
+        #expect(before.content == "before")
+        #expect(after.content == "after")
+        #expect(image.source == .url(imageURL))
+
+        let transcript = AnyLanguageModel.Transcript(entries: [
+            .toolOutput(.init(id: "image-call", toolName: "image", segments: segments))
+        ])
+        let encoded = try JSONEncoder().encode(transcript)
+        #expect(try JSONDecoder().decode(AnyLanguageModel.Transcript.self, from: encoded) == transcript)
+    }
+
     #if canImport(CoreGraphics) && canImport(ImageIO)
         private func image() throws -> CGImage {
             let context = try #require(
@@ -267,11 +290,41 @@ struct PromptAttachmentTests {
             #expect(second.content == "prompts:2 images:1 right:1")
         }
 
+        @available(macOS 27, iOS 27, visionOS 27, *)
+        @Test func nativeToolWrapperRetainsOrdinaryImageOutput() async throws {
+            let output = try AnyLanguageModel.Prompt {
+                "tool note"
+                AnyLanguageModel.Attachment(try image(), orientation: .right)
+            }
+            let tools: [any AnyLanguageModel.Tool] = [ImageTool(output: output)]
+            let session = FoundationModels.LanguageModelSession(
+                model: BridgeModel(callsTools: true),
+                tools: tools.toFoundationModels()
+            )
+            _ = try await session.respond(to: "read image")
+            let recorded = try #require(
+                session.transcript.compactMap { entry -> FoundationModels.Transcript.ToolOutput? in
+                    guard case .toolOutput(let output) = entry else { return nil }
+                    return output
+                }.first
+            )
+            #expect(recorded.id == "image-call")
+            #expect(recorded.segments.count == 2)
+            guard case .attachment(let segment) = recorded.segments[1],
+                case .image(let image) = segment.content
+            else {
+                Issue.record("Native Tool wrapper flattened image output")
+                return
+            }
+            #expect(image.orientation == .right)
+        }
+
         @available(macOS 27, iOS 27, visionOS 27, *)
         private struct BridgeModel: FoundationModels.LanguageModel {
             typealias Executor = BridgeExecutor
+            var callsTools = false
             var capabilities: FoundationModels.LanguageModelCapabilities {
-                .init([.vision])
+                .init([.vision, .toolCalling])
             }
             var executorConfiguration: String {
                 "bridge-fixture"
@@ -287,6 +340,27 @@ struct PromptAttachmentTests {
                 model: BridgeModel,
                 streamingInto channel: FoundationModels.LanguageModelExecutorGenerationChannel
             ) async throws {
+                if model.callsTools,
+                    !request.transcript.contains(where: {
+                        if case .toolOutput = $0 { return true }
+                        return false
+                    })
+                {
+                    await channel.send(
+                        .toolCalls(
+                            action: .toolCall(
+                                id: "image-call",
+                                name: "image",
+                                action: .appendArguments(#""read""#, tokenCount: 1)
+                            )
+                        )
+                    )
+                    return
+                }
+                if model.callsTools {
+                    await channel.send(.response(action: .appendText("OK", tokenCount: 1)))
+                    return
+                }
                 let prompts = request.transcript.compactMap { entry -> FoundationModels.Transcript.Prompt? in
                     guard case .prompt(let prompt) = entry else { return nil }
                     return prompt
@@ -313,3 +387,17 @@ struct PromptAttachmentTests {
         }
     #endif
 }
+
+private struct ImageTool: AnyLanguageModel.Tool {
+    let name = "image"
+    let description = "Return an image"
+    let output: AnyLanguageModel.Prompt
+
+    var parameters: AnyLanguageModel.GenerationSchema {
+        String.generationSchema
+    }
+
+    func call(arguments: String) async throws -> AnyLanguageModel.Prompt {
+        output
+    }
+}

Tests/AnyLanguageModelTests/ToolExecutionDelegateTests.swift

@@ -1,6 +1,6 @@
 import Testing
 
-@testable import AnyLanguageModel
+@_spi(Compatibility) @testable import AnyLanguageModel
 
 private actor ToolExecutionDelegateSpy: ToolExecutionDelegate {
     private(set) var generatedToolCalls: [Transcript.ToolCall] = []