Make tool schema inclusion behavior consistent with Foundation Models

Sources/AnyLanguageModel/LanguageModelSession.swift

@@ -77,7 +77,10 @@ public final class LanguageModelSession: @unchecked Sendable {
                 let instructionsEntry = Transcript.Entry.instructions(
                     Transcript.Instructions(
                         segments: [.text(.init(content: instructions.description))],
-                        toolDefinitions: tools.map { Transcript.ToolDefinition(tool: $0) }
+                        toolDefinitions:
+                            tools
+                            .filter(\.includesSchemaInInstructions)
+                            .map { Transcript.ToolDefinition(tool: $0) }
                     )
                 )
                 finalTranscript.append(instructionsEntry)

Sources/AnyLanguageModel/Models/SystemLanguageModel.swift

@@ -78,7 +78,12 @@
             let fmSession = FoundationModels.LanguageModelSession(
                 model: systemModel,
                 tools: session.tools.toFoundationModels(),
-                transcript: session.transcript.toFoundationModels(instructions: session.instructions)
+                transcript: session.transcript.toFoundationModels(
+                    instructions: session.instructions,
+                    toolDefinitions: session.tools
+                        .filter(\.includesSchemaInInstructions)
+                        .map { Transcript.ToolDefinition(tool: $0) }
+                )
             )
 
             if type == String.self {
@@ -154,7 +159,12 @@
             let fmSession = FoundationModels.LanguageModelSession(
                 model: systemModel,
                 tools: session.tools.toFoundationModels(),
-                transcript: session.transcript.toFoundationModels(instructions: session.instructions)
+                transcript: session.transcript.toFoundationModels(
+                    instructions: session.instructions,
+                    toolDefinitions: session.tools
+                        .filter(\.includesSchemaInInstructions)
+                        .map { Transcript.ToolDefinition(tool: $0) }
+                )
             )
 
             let stream: AsyncThrowingStream<LanguageModelSession.ResponseStream<Content>.Snapshot, Error> =
@@ -642,8 +652,10 @@
 
     @available(macOS 26.0, iOS 26.0, watchOS 26.0, tvOS 26.0, visionOS 26.0, *)
     extension Transcript {
-        fileprivate func toFoundationModels(instructions: AnyLanguageModel.Instructions?) -> FoundationModels.Transcript
-        {
+        fileprivate func toFoundationModels(
+            instructions: AnyLanguageModel.Instructions?,
+            toolDefinitions: [Transcript.ToolDefinition]
+        ) -> FoundationModels.Transcript {
             var fmEntries: [FoundationModels.Transcript.Entry] = []
 
             // Add instructions entry if provided and not already in transcript
@@ -656,7 +668,7 @@
                 if !hasInstructions {
                     let fmInstructions = FoundationModels.Transcript.Instructions(
                         segments: [.text(.init(content: instructions.description))],
-                        toolDefinitions: []
+                        toolDefinitions: toolDefinitions.toFoundationModels()
                     )
                     fmEntries.append(.instructions(fmInstructions))
                 }

Tests/AnyLanguageModelTests/MockLanguageModelTests.swift

@@ -2,6 +2,37 @@ import Testing
 
 @testable import AnyLanguageModel
 
+private struct ToolSchemaOptOutTool: Tool {
+    let name = "schemaOptOutTool"
+    let description = "A tool that opts out of schema injection."
+    let includesSchemaInInstructions = false
+
+    @Generable
+    struct Arguments {
+        @Guide(description: "A value to echo")
+        var value: String
+    }
+
+    func call(arguments: Arguments) async throws -> String {
+        arguments.value
+    }
+}
+
+private func waitUntil(
+    timeout: Duration = .milliseconds(500),
+    pollInterval: Duration = .milliseconds(5),
+    _ condition: @autoclosure () -> Bool
+) async throws -> Bool {
+    let deadline = ContinuousClock.now + timeout
+    while !condition() {
+        if ContinuousClock.now >= deadline {
+            return false
+        }
+        try await Task.sleep(for: pollInterval)
+    }
+    return true
+}
+
 @Suite("MockLanguageModel")
 struct MockLanguageModelTests {
     @Test func fixedResponse() async throws {
@@ -69,6 +100,92 @@ struct MockLanguageModelTests {
         }
     }
 
+    @Test func instructionsOnlyIncludeToolsThatOptInToSchemaInjection() {
+        let model = MockLanguageModel.fixed("ok")
+        let tools: [any Tool] = [WeatherTool(), ToolSchemaOptOutTool()]
+        let session = LanguageModelSession(
+            model: model,
+            tools: tools,
+            instructions: "Use tools when appropriate."
+        )
+
+        #expect(session.transcript.count == 1)
+        guard case let .instructions(instructionsEntry) = session.transcript[0] else {
+            Issue.record("First entry should be instructions")
+            return
+        }
+
+        #expect(instructionsEntry.toolDefinitions.count == 1)
+        #expect(instructionsEntry.toolDefinitions.first?.name == "getWeather")
+    }
+
+    @Test func transcriptAndInstructionsInitHaveEquivalentToolDefinitionBehavior() {
+        let model = MockLanguageModel.fixed("ok")
+        let tools: [any Tool] = [WeatherTool(), ToolSchemaOptOutTool()]
+        let expectedToolDefinitions =
+            tools
+            .filter(\.includesSchemaInInstructions)
+            .map { Transcript.ToolDefinition(tool: $0) }
+
+        let instructionsSession = LanguageModelSession(
+            model: model,
+            tools: tools,
+            instructions: "Use tools when appropriate."
+        )
+
+        let explicitTranscript = Transcript(entries: [
+            .instructions(
+                Transcript.Instructions(
+                    segments: [.text(.init(content: "Use tools when appropriate."))],
+                    toolDefinitions: expectedToolDefinitions
+                )
+            )
+        ])
+        let transcriptSession = LanguageModelSession(
+            model: model,
+            tools: tools,
+            transcript: explicitTranscript
+        )
+
+        guard case let .instructions(instructionsEntry) = instructionsSession.transcript[0] else {
+            Issue.record("First instructions-based entry should be instructions")
+            return
+        }
+        guard case let .instructions(transcriptEntry) = transcriptSession.transcript[0] else {
+            Issue.record("First transcript-based entry should be instructions")
+            return
+        }
+
+        #expect(instructionsEntry.toolDefinitions == transcriptEntry.toolDefinitions)
+    }
+
+    @Test func explicitTranscriptInstructionsArePreservedWhenToolsOptOut() {
+        let model = MockLanguageModel.fixed("ok")
+        let tools: [any Tool] = [WeatherTool(), ToolSchemaOptOutTool()]
+        let explicitToolDefinitions = tools.map { Transcript.ToolDefinition(tool: $0) }
+        let transcript = Transcript(entries: [
+            .instructions(
+                Transcript.Instructions(
+                    segments: [.text(.init(content: "Custom instructions."))],
+                    toolDefinitions: explicitToolDefinitions
+                )
+            )
+        ])
+
+        let session = LanguageModelSession(
+            model: model,
+            tools: tools,
+            transcript: transcript
+        )
+
+        guard case let .instructions(entry) = session.transcript[0] else {
+            Issue.record("First entry should be instructions")
+            return
+        }
+
+        #expect(entry.toolDefinitions == explicitToolDefinitions)
+    }
+
     @Test func unavailable() async throws {
         let model = MockLanguageModel.unavailable
 
@@ -91,8 +208,7 @@ struct MockLanguageModelTests {
             try await asyncSession.respond(to: "Async test")
         }
 
-        try await Task.sleep(for: .milliseconds(50))
-        #expect(asyncSession.isResponding == true)
+        #expect(try await waitUntil(asyncSession.isResponding == true))
 
         _ = try await asyncTask.value
         try await Task.sleep(for: .milliseconds(10))
@@ -112,8 +228,7 @@ struct MockLanguageModelTests {
             for try await _ in stream {}
         }
 
-        try await Task.sleep(for: .milliseconds(50))
-        #expect(streamSession.isResponding == true)
+        #expect(try await waitUntil(streamSession.isResponding == true))
 
         _ = try await streamTask.value
         try await Task.sleep(for: .milliseconds(10))