From a6ff6f714ed39f879030e63e3cfde77780f3cd46 Mon Sep 17 00:00:00 2001 From: Mattt Zmuda Date: Tue, 17 Feb 2026 03:06:43 -0800 Subject: [PATCH 1/4] Filter injected tool definitions by includesSchemaInInstructions --- .../AnyLanguageModel/LanguageModelSession.swift | 5 ++++- .../Models/SystemLanguageModel.swift | 14 ++++++++++++-- 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/Sources/AnyLanguageModel/LanguageModelSession.swift b/Sources/AnyLanguageModel/LanguageModelSession.swift index 1b569491..e25fa03a 100644 --- a/Sources/AnyLanguageModel/LanguageModelSession.swift +++ b/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) diff --git a/Sources/AnyLanguageModel/Models/SystemLanguageModel.swift b/Sources/AnyLanguageModel/Models/SystemLanguageModel.swift index 52c1db90..1502cd41 100644 --- a/Sources/AnyLanguageModel/Models/SystemLanguageModel.swift +++ b/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.Snapshot, Error> = From 93d11659d962e6b34c34e89cdc9d6560a5d0cbab Mon Sep 17 00:00:00 2001 From: Mattt Zmuda Date: Tue, 17 Feb 2026 03:07:11 -0800 Subject: [PATCH 2/4] Fixed synthesized Foundation Models transcript instructions to include tool definitions --- Sources/AnyLanguageModel/Models/SystemLanguageModel.swift | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/Sources/AnyLanguageModel/Models/SystemLanguageModel.swift b/Sources/AnyLanguageModel/Models/SystemLanguageModel.swift index 1502cd41..97fddb0b 100644 --- a/Sources/AnyLanguageModel/Models/SystemLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/SystemLanguageModel.swift @@ -652,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 @@ -666,7 +668,7 @@ if !hasInstructions { let fmInstructions = FoundationModels.Transcript.Instructions( segments: [.text(.init(content: instructions.description))], - toolDefinitions: [] + toolDefinitions: toolDefinitions.toFoundationModels() ) fmEntries.append(.instructions(fmInstructions)) } From 1a6539b3ad154cd646b777ede648a97515073d2c Mon Sep 17 00:00:00 2001 From: Mattt Zmuda Date: Tue, 17 Feb 2026 03:07:34 -0800 Subject: [PATCH 3/4] Add test coverage for tool definitions in instructions --- .../MockLanguageModelTests.swift | 102 ++++++++++++++++++ 1 file changed, 102 insertions(+) diff --git a/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift b/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift index baff85dd..5d5e927a 100644 --- a/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift +++ b/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift @@ -2,6 +2,22 @@ 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 + } +} + @Suite("MockLanguageModel") struct MockLanguageModelTests { @Test func fixedResponse() async throws { @@ -69,6 +85,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 From 97a8d24b9556e83a7aac0d3c0f20bfcbf429fe88 Mon Sep 17 00:00:00 2001 From: Mattt Zmuda Date: Tue, 17 Feb 2026 03:45:03 -0800 Subject: [PATCH 4/4] Fix race in response streaming tests --- .../MockLanguageModelTests.swift | 21 +++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift b/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift index 5d5e927a..ad8e1c06 100644 --- a/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift +++ b/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift @@ -18,6 +18,21 @@ private struct ToolSchemaOptOutTool: Tool { } } +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 { @@ -193,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)) @@ -214,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))