From a1fbc527703e0c3789cf3448af94dfa45ee2321d Mon Sep 17 00:00:00 2001 From: Mattt Zmuda Date: Tue, 28 Oct 2025 05:06:40 -0700 Subject: [PATCH 1/7] Implement tracking for LanguageModelSession.isResponding --- Sources/AnyLanguageModel/LanguageModel.swift | 6 + .../LanguageModelSession.swift | 180 ++++++++++++------ .../Models/SystemLanguageModel.swift | 10 +- 3 files changed, 138 insertions(+), 58 deletions(-) diff --git a/Sources/AnyLanguageModel/LanguageModel.swift b/Sources/AnyLanguageModel/LanguageModel.swift index e25a7886..0b0a3a7d 100644 --- a/Sources/AnyLanguageModel/LanguageModel.swift +++ b/Sources/AnyLanguageModel/LanguageModel.swift @@ -5,6 +5,8 @@ public protocol LanguageModel: Sendable { var availability: Availability { get } + var isResponding: Bool { get } + func prewarm( for session: LanguageModelSession, promptPrefix: Prompt? @@ -45,6 +47,10 @@ extension LanguageModel { } } + public var isResponding: Bool { + return false + } + public func prewarm( for session: LanguageModelSession, promptPrefix: Prompt? = nil diff --git a/Sources/AnyLanguageModel/LanguageModelSession.swift b/Sources/AnyLanguageModel/LanguageModelSession.swift index fe52291f..af9bc747 100644 --- a/Sources/AnyLanguageModel/LanguageModelSession.swift +++ b/Sources/AnyLanguageModel/LanguageModelSession.swift @@ -2,7 +2,7 @@ import Foundation import Observation @Observable -public final class LanguageModelSession { +public final class LanguageModelSession: @unchecked Sendable { public private(set) var isResponding: Bool = false public private(set) var transcript: Transcript @@ -10,6 +10,8 @@ public final class LanguageModelSession { public let tools: [any Tool] public let instructions: Instructions? + @ObservationIgnored private let respondingState = RespondingState() + public convenience init( model: any LanguageModel, tools: [any Tool] = [], @@ -58,28 +60,78 @@ public final class LanguageModelSession { model.prewarm(for: self, promptPrefix: promptPrefix) } - public struct Response where Content: Generable { + nonisolated private func beginResponding() { + Task { + let count = await respondingState.increment() + let active = count > 0 + await MainActor.run { + self.isResponding = active + } + } + } + + nonisolated private func endResponding() { + Task { + let count = await respondingState.decrement() + let active = count > 0 + await MainActor.run { + self.isResponding = active + } + } + } + + nonisolated private func wrapRespond(_ operation: () async throws -> T) async throws -> T { + beginResponding() + defer { endResponding() } + return try await operation() + } + + nonisolated private consuming func wrapStream( + _ upstream: sending ResponseStream + ) -> ResponseStream where Content: Generable, Content.PartiallyGenerated: Sendable { + beginResponding() + let session = self + let relay = AsyncThrowingStream.Snapshot, any Error> { continuation in + let stream = upstream + Task { + do { + for try await snapshot in stream { + continuation.yield(snapshot) + } + continuation.finish() + } catch { + continuation.finish(throwing: error) + } + session.endResponding() + } + } + return ResponseStream(stream: relay) + } + + public struct Response: Sendable where Content: Generable, Content: Sendable { public let content: Content public let rawContent: GeneratedContent public let transcriptEntries: ArraySlice } @discardableResult - nonisolated public func respond( + public func respond( to prompt: Prompt, options: GenerationOptions = GenerationOptions() ) async throws -> Response { - try await model.respond( - within: self, - to: prompt, - generating: String.self, - includeSchemaInPrompt: true, - options: options - ) + try await wrapRespond { + try await model.respond( + within: self, + to: prompt, + generating: String.self, + includeSchemaInPrompt: true, + options: options + ) + } } @discardableResult - nonisolated public func respond( + public func respond( to prompt: String, options: GenerationOptions = GenerationOptions() ) async throws -> Response { @@ -87,7 +139,7 @@ public final class LanguageModelSession { } @discardableResult - nonisolated public func respond( + public func respond( options: GenerationOptions = GenerationOptions(), @PromptBuilder prompt: () throws -> Prompt ) async throws -> Response { @@ -95,23 +147,25 @@ public final class LanguageModelSession { } @discardableResult - nonisolated public func respond( + public func respond( to prompt: Prompt, schema: GenerationSchema, includeSchemaInPrompt: Bool = true, options: GenerationOptions = GenerationOptions() ) async throws -> Response { - try await model.respond( - within: self, - to: prompt, - generating: GeneratedContent.self, - includeSchemaInPrompt: includeSchemaInPrompt, - options: options - ) + try await wrapRespond { + try await model.respond( + within: self, + to: prompt, + generating: GeneratedContent.self, + includeSchemaInPrompt: includeSchemaInPrompt, + options: options + ) + } } @discardableResult - nonisolated public func respond( + public func respond( to prompt: String, schema: GenerationSchema, includeSchemaInPrompt: Bool = true, @@ -126,7 +180,7 @@ public final class LanguageModelSession { } @discardableResult - nonisolated public func respond( + public func respond( schema: GenerationSchema, includeSchemaInPrompt: Bool = true, options: GenerationOptions = GenerationOptions(), @@ -141,23 +195,25 @@ public final class LanguageModelSession { } @discardableResult - nonisolated public func respond( + public func respond( to prompt: Prompt, generating type: Content.Type = Content.self, includeSchemaInPrompt: Bool = true, options: GenerationOptions = GenerationOptions() ) async throws -> Response where Content: Generable { - try await model.respond( - within: self, - to: prompt, - generating: type, - includeSchemaInPrompt: includeSchemaInPrompt, - options: options - ) + try await wrapRespond { + try await model.respond( + within: self, + to: prompt, + generating: type, + includeSchemaInPrompt: includeSchemaInPrompt, + options: options + ) + } } @discardableResult - nonisolated public func respond( + public func respond( to prompt: String, generating type: Content.Type = Content.self, includeSchemaInPrompt: Bool = true, @@ -172,7 +228,7 @@ public final class LanguageModelSession { } @discardableResult - nonisolated public func respond( + public func respond( generating type: Content.Type = Content.self, includeSchemaInPrompt: Bool = true, options: GenerationOptions = GenerationOptions(), @@ -192,12 +248,14 @@ public final class LanguageModelSession { includeSchemaInPrompt: Bool = true, options: GenerationOptions = GenerationOptions() ) -> sending ResponseStream { - model.streamResponse( - within: self, - to: prompt, - generating: GeneratedContent.self, - includeSchemaInPrompt: includeSchemaInPrompt, - options: options + wrapStream( + model.streamResponse( + within: self, + to: prompt, + generating: GeneratedContent.self, + includeSchemaInPrompt: includeSchemaInPrompt, + options: options + ) ) } @@ -230,12 +288,14 @@ public final class LanguageModelSession { includeSchemaInPrompt: Bool = true, options: GenerationOptions = GenerationOptions() ) -> sending ResponseStream where Content: Generable { - model.streamResponse( - within: self, - to: prompt, - generating: type, - includeSchemaInPrompt: includeSchemaInPrompt, - options: options + wrapStream( + model.streamResponse( + within: self, + to: prompt, + generating: type, + includeSchemaInPrompt: includeSchemaInPrompt, + options: options + ) ) } @@ -271,12 +331,14 @@ public final class LanguageModelSession { to prompt: Prompt, options: GenerationOptions = GenerationOptions() ) -> sending ResponseStream { - model.streamResponse( - within: self, - to: prompt, - generating: String.self, - includeSchemaInPrompt: true, - options: options + wrapStream( + model.streamResponse( + within: self, + to: prompt, + generating: String.self, + includeSchemaInPrompt: true, + options: options + ) ) } @@ -309,7 +371,19 @@ public final class LanguageModelSession { } } -extension LanguageModelSession: @unchecked Sendable, Observable {} +private actor RespondingState { + private var count = 0 + + func increment() -> Int { + count += 1 + return count + } + + func decrement() -> Int { + count = max(0, count - 1) + return count + } +} extension LanguageModelSession { public enum GenerationError: Error, LocalizedError { @@ -401,7 +475,7 @@ extension LanguageModelSession { } extension LanguageModelSession { - public struct ResponseStream where Content: Generable { + public struct ResponseStream: Sendable where Content: Generable, Content.PartiallyGenerated: Sendable { private let content: Content private let rawContent: GeneratedContent private let streaming: AsyncThrowingStream? @@ -420,7 +494,7 @@ extension LanguageModelSession { self.streaming = stream } - public struct Snapshot { + public struct Snapshot: Sendable where Content.PartiallyGenerated: Sendable { public var content: Content.PartiallyGenerated public var rawContent: GeneratedContent } diff --git a/Sources/AnyLanguageModel/Models/SystemLanguageModel.swift b/Sources/AnyLanguageModel/Models/SystemLanguageModel.swift index 0f66f519..930fd274 100644 --- a/Sources/AnyLanguageModel/Models/SystemLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/SystemLanguageModel.swift @@ -12,7 +12,7 @@ /// let model = SystemLanguageModel() /// ``` @available(macOS 26.0, iOS 26.0, watchOS 26.0, tvOS 26.0, visionOS 26.0, *) - public struct SystemLanguageModel: LanguageModel { + public actor SystemLanguageModel: LanguageModel { /// The reason the model is unavailable. public typealias UnavailableReason = FoundationModels.SystemLanguageModel.Availability.UnavailableReason @@ -54,7 +54,7 @@ } /// The availability status for the system language model. - public var availability: Availability { + nonisolated public var availability: Availability { switch systemModel.availability { case .available: .available @@ -63,7 +63,7 @@ } } - public func respond( + nonisolated public func respond( within session: LanguageModelSession, to prompt: Prompt, generating type: Content.Type, @@ -100,7 +100,7 @@ } } - public func streamResponse( + nonisolated public func streamResponse( within session: LanguageModelSession, to prompt: Prompt, generating type: Content.Type, @@ -180,7 +180,7 @@ return LanguageModelSession.ResponseStream(stream: stream) } - public func logFeedbackAttachment( + nonisolated public func logFeedbackAttachment( within session: LanguageModelSession, sentiment: LanguageModelFeedback.Sentiment?, issues: [LanguageModelFeedback.Issue], From 99681cabc1bc79cb1c6f83ebdcb4118d37115ee7 Mon Sep 17 00:00:00 2001 From: Mattt Zmuda Date: Tue, 28 Oct 2025 05:13:45 -0700 Subject: [PATCH 2/7] Add test coverage for isResponding using MockLanguageModel --- .../MockLanguageModelTests.swift | 44 +++++++++++++++++++ .../Shared/MockLanguageModel.swift | 36 +++++++++++---- 2 files changed, 72 insertions(+), 8 deletions(-) diff --git a/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift b/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift index 3a8f4db5..d6dcdf0d 100644 --- a/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift +++ b/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift @@ -55,4 +55,48 @@ struct MockLanguageModelTests { #expect(model.availability == .unavailable(.custom("MockLanguageModel is unavailable"))) #expect(model.isAvailable == false) } + + @Test func isRespondingDuringAsyncResponse() async throws { + let model = MockLanguageModel { _, _ in + try await Task.sleep(for: .milliseconds(100)) + return "Response" + } + let session = LanguageModelSession(model: model) + + #expect(session.isResponding == false) + + let task = Task { + try await session.respond(to: "Test") + } + + try await Task.sleep(for: .milliseconds(50)) + #expect(session.isResponding == true) + + _ = try await task.value + try await Task.sleep(for: .milliseconds(10)) + #expect(session.isResponding == false) + } + + @Test func isRespondingDuringStreaming() async throws { + let model = MockLanguageModel.streamingMock() + let session = LanguageModelSession(model: model) + + #expect(session.isResponding == false) + + let stream = session.streamResponse(to: "Test") + + try await Task.sleep(for: .milliseconds(50)) + #expect(session.isResponding == true) + + var count = 0 + for try await _ in stream { + count += 1 + if count == 1 { + #expect(session.isResponding == true) + } + } + + try await Task.sleep(for: .milliseconds(10)) + #expect(session.isResponding == false) + } } diff --git a/Tests/AnyLanguageModelTests/Shared/MockLanguageModel.swift b/Tests/AnyLanguageModelTests/Shared/MockLanguageModel.swift index 8f65b212..9598a7f7 100644 --- a/Tests/AnyLanguageModelTests/Shared/MockLanguageModel.swift +++ b/Tests/AnyLanguageModelTests/Shared/MockLanguageModel.swift @@ -55,14 +55,27 @@ struct MockLanguageModel: LanguageModel { fatalError("MockLanguageModel only supports generating String content") } - // For MockLanguageModel, we'll simulate streaming by yielding the response immediately - // In a real implementation, this would stream the response as it's generated - // Since we can't make this function async, we'll need to handle this differently - // For now, we'll create a stream that yields immediately with a placeholder - let placeholderText = "Mock streaming response" - let generatedContent = GeneratedContent(placeholderText) - - return LanguageModelSession.ResponseStream(content: placeholderText as! Content, rawContent: generatedContent) + let promptWithInstructions = Prompt("Instructions: \(session.instructions?.description ?? "N/A")\n\(prompt))") + + let stream = AsyncThrowingStream.Snapshot, any Error> { + continuation in + Task { + do { + let text = try await responseProvider(promptWithInstructions, options) + let generatedContent = GeneratedContent(text) + let snapshot = LanguageModelSession.ResponseStream.Snapshot( + content: (text as! Content).asPartiallyGenerated(), + rawContent: generatedContent + ) + continuation.yield(snapshot) + continuation.finish() + } catch { + continuation.finish(throwing: error) + } + } + } + + return LanguageModelSession.ResponseStream(stream: stream) } } @@ -84,4 +97,11 @@ extension MockLanguageModel { model.availabilityProvider = { .unavailable(.custom("MockLanguageModel is unavailable")) } return model } + + static func streamingMock() -> Self { + MockLanguageModel { _, _ in + try await Task.sleep(for: .milliseconds(100)) + return "Streaming response" + } + } } From 046833dabd1076e0588ce2a6f5b9d1e1fdf9b89f Mon Sep 17 00:00:00 2001 From: Mattt Zmuda Date: Tue, 28 Oct 2025 05:17:26 -0700 Subject: [PATCH 3/7] Make structured concurrency usage more robust --- .../LanguageModelSession.swift | 42 ++++++++++--------- 1 file changed, 22 insertions(+), 20 deletions(-) diff --git a/Sources/AnyLanguageModel/LanguageModelSession.swift b/Sources/AnyLanguageModel/LanguageModelSession.swift index af9bc747..98470003 100644 --- a/Sources/AnyLanguageModel/LanguageModelSession.swift +++ b/Sources/AnyLanguageModel/LanguageModelSession.swift @@ -60,40 +60,42 @@ public final class LanguageModelSession: @unchecked Sendable { model.prewarm(for: self, promptPrefix: promptPrefix) } - nonisolated private func beginResponding() { - Task { - let count = await respondingState.increment() - let active = count > 0 - await MainActor.run { - self.isResponding = active - } + nonisolated private func beginResponding() async { + let count = await respondingState.increment() + let active = count > 0 + await MainActor.run { + self.isResponding = active } } - nonisolated private func endResponding() { - Task { - let count = await respondingState.decrement() - let active = count > 0 - await MainActor.run { - self.isResponding = active - } + nonisolated private func endResponding() async { + let count = await respondingState.decrement() + let active = count > 0 + await MainActor.run { + self.isResponding = active } } nonisolated private func wrapRespond(_ operation: () async throws -> T) async throws -> T { - beginResponding() - defer { endResponding() } - return try await operation() + await beginResponding() + do { + let result = try await operation() + await endResponding() + return result + } catch { + await endResponding() + throw error + } } - nonisolated private consuming func wrapStream( + nonisolated private func wrapStream( _ upstream: sending ResponseStream ) -> ResponseStream where Content: Generable, Content.PartiallyGenerated: Sendable { - beginResponding() let session = self let relay = AsyncThrowingStream.Snapshot, any Error> { continuation in let stream = upstream Task { + await session.beginResponding() do { for try await snapshot in stream { continuation.yield(snapshot) @@ -102,7 +104,7 @@ public final class LanguageModelSession: @unchecked Sendable { } catch { continuation.finish(throwing: error) } - session.endResponding() + await session.endResponding() } } return ResponseStream(stream: relay) From 404da13af24a345cbeb95a1dbfbd10386bff1103 Mon Sep 17 00:00:00 2001 From: Mattt Zmuda Date: Tue, 28 Oct 2025 05:20:16 -0700 Subject: [PATCH 4/7] Remove isResponding property from LanguageModel protocol --- Sources/AnyLanguageModel/LanguageModel.swift | 6 ------ 1 file changed, 6 deletions(-) diff --git a/Sources/AnyLanguageModel/LanguageModel.swift b/Sources/AnyLanguageModel/LanguageModel.swift index 0b0a3a7d..e25a7886 100644 --- a/Sources/AnyLanguageModel/LanguageModel.swift +++ b/Sources/AnyLanguageModel/LanguageModel.swift @@ -5,8 +5,6 @@ public protocol LanguageModel: Sendable { var availability: Availability { get } - var isResponding: Bool { get } - func prewarm( for session: LanguageModelSession, promptPrefix: Prompt? @@ -47,10 +45,6 @@ extension LanguageModel { } } - public var isResponding: Bool { - return false - } - public func prewarm( for session: LanguageModelSession, promptPrefix: Prompt? = nil From a2d3c2b524e6bdc37ece4916c365ec227c35bf74 Mon Sep 17 00:00:00 2001 From: Mattt Zmuda Date: Tue, 28 Oct 2025 05:31:50 -0700 Subject: [PATCH 5/7] Fix typo in MockLanguageModel --- Tests/AnyLanguageModelTests/Shared/MockLanguageModel.swift | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/Tests/AnyLanguageModelTests/Shared/MockLanguageModel.swift b/Tests/AnyLanguageModelTests/Shared/MockLanguageModel.swift index 9598a7f7..bcb2bd6f 100644 --- a/Tests/AnyLanguageModelTests/Shared/MockLanguageModel.swift +++ b/Tests/AnyLanguageModelTests/Shared/MockLanguageModel.swift @@ -33,7 +33,7 @@ struct MockLanguageModel: LanguageModel { fatalError("MockLanguageModel only supports generating String content") } - let promptWithInstructions = Prompt("Instructions: \(session.instructions?.description ?? "N/A")\n\(prompt))") + let promptWithInstructions = Prompt("Instructions: \(session.instructions?.description ?? "N/A")\n\(prompt)") let text = try await responseProvider(promptWithInstructions, options) return LanguageModelSession.Response( @@ -55,7 +55,7 @@ struct MockLanguageModel: LanguageModel { fatalError("MockLanguageModel only supports generating String content") } - let promptWithInstructions = Prompt("Instructions: \(session.instructions?.description ?? "N/A")\n\(prompt))") + let promptWithInstructions = Prompt("Instructions: \(session.instructions?.description ?? "N/A")\n\(prompt)") let stream = AsyncThrowingStream.Snapshot, any Error> { continuation in From 8379c756645711c7763409b7b33933b1d7459ba3 Mon Sep 17 00:00:00 2001 From: Mattt Zmuda Date: Tue, 28 Oct 2025 05:34:36 -0700 Subject: [PATCH 6/7] Re-annotate respond methods with nonisolated --- .../LanguageModelSession.swift | 28 +++++++++---------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/Sources/AnyLanguageModel/LanguageModelSession.swift b/Sources/AnyLanguageModel/LanguageModelSession.swift index 98470003..965a46f2 100644 --- a/Sources/AnyLanguageModel/LanguageModelSession.swift +++ b/Sources/AnyLanguageModel/LanguageModelSession.swift @@ -117,7 +117,7 @@ public final class LanguageModelSession: @unchecked Sendable { } @discardableResult - public func respond( + nonisolated public func respond( to prompt: Prompt, options: GenerationOptions = GenerationOptions() ) async throws -> Response { @@ -133,7 +133,7 @@ public final class LanguageModelSession: @unchecked Sendable { } @discardableResult - public func respond( + nonisolated public func respond( to prompt: String, options: GenerationOptions = GenerationOptions() ) async throws -> Response { @@ -141,7 +141,7 @@ public final class LanguageModelSession: @unchecked Sendable { } @discardableResult - public func respond( + nonisolated public func respond( options: GenerationOptions = GenerationOptions(), @PromptBuilder prompt: () throws -> Prompt ) async throws -> Response { @@ -149,7 +149,7 @@ public final class LanguageModelSession: @unchecked Sendable { } @discardableResult - public func respond( + nonisolated public func respond( to prompt: Prompt, schema: GenerationSchema, includeSchemaInPrompt: Bool = true, @@ -167,7 +167,7 @@ public final class LanguageModelSession: @unchecked Sendable { } @discardableResult - public func respond( + nonisolated public func respond( to prompt: String, schema: GenerationSchema, includeSchemaInPrompt: Bool = true, @@ -182,7 +182,7 @@ public final class LanguageModelSession: @unchecked Sendable { } @discardableResult - public func respond( + nonisolated public func respond( schema: GenerationSchema, includeSchemaInPrompt: Bool = true, options: GenerationOptions = GenerationOptions(), @@ -197,7 +197,7 @@ public final class LanguageModelSession: @unchecked Sendable { } @discardableResult - public func respond( + nonisolated public func respond( to prompt: Prompt, generating type: Content.Type = Content.self, includeSchemaInPrompt: Bool = true, @@ -215,7 +215,7 @@ public final class LanguageModelSession: @unchecked Sendable { } @discardableResult - public func respond( + nonisolated public func respond( to prompt: String, generating type: Content.Type = Content.self, includeSchemaInPrompt: Bool = true, @@ -230,7 +230,7 @@ public final class LanguageModelSession: @unchecked Sendable { } @discardableResult - public func respond( + nonisolated public func respond( generating type: Content.Type = Content.self, includeSchemaInPrompt: Bool = true, options: GenerationOptions = GenerationOptions(), @@ -244,7 +244,7 @@ public final class LanguageModelSession: @unchecked Sendable { ) } - public func streamResponse( + nonisolated public func streamResponse( to prompt: Prompt, schema: GenerationSchema, includeSchemaInPrompt: Bool = true, @@ -261,7 +261,7 @@ public final class LanguageModelSession: @unchecked Sendable { ) } - public func streamResponse( + nonisolated public func streamResponse( to prompt: String, schema: GenerationSchema, includeSchemaInPrompt: Bool = true, @@ -275,7 +275,7 @@ public final class LanguageModelSession: @unchecked Sendable { ) } - public func streamResponse( + nonisolated public func streamResponse( schema: GenerationSchema, includeSchemaInPrompt: Bool = true, options: GenerationOptions = GenerationOptions(), @@ -284,7 +284,7 @@ public final class LanguageModelSession: @unchecked Sendable { streamResponse(to: try prompt(), schema: schema, includeSchemaInPrompt: includeSchemaInPrompt, options: options) } - public func streamResponse( + nonisolated public func streamResponse( to prompt: Prompt, generating type: Content.Type = Content.self, includeSchemaInPrompt: Bool = true, @@ -301,7 +301,7 @@ public final class LanguageModelSession: @unchecked Sendable { ) } - public func streamResponse( + nonisolated public func streamResponse( to prompt: String, generating type: Content.Type = Content.self, includeSchemaInPrompt: Bool = true, From 4eb72935cb7577df3aef34720fa253450a070307 Mon Sep 17 00:00:00 2001 From: Mattt Zmuda Date: Tue, 28 Oct 2025 05:36:21 -0700 Subject: [PATCH 7/7] Fix race condition in mock language model tests --- .../MockLanguageModelTests.swift | 20 +++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift b/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift index d6dcdf0d..92fc8be5 100644 --- a/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift +++ b/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift @@ -84,18 +84,22 @@ struct MockLanguageModelTests { #expect(session.isResponding == false) let stream = session.streamResponse(to: "Test") + + // Start consuming the stream in a task + let task = Task { + for try await _ in stream { + // Just consume the stream + } + } + // Give the streaming task time to start and call beginResponding try await Task.sleep(for: .milliseconds(50)) #expect(session.isResponding == true) - var count = 0 - for try await _ in stream { - count += 1 - if count == 1 { - #expect(session.isResponding == true) - } - } - + // Wait for stream to complete + _ = try await task.value + + // Give time for endResponding to complete try await Task.sleep(for: .milliseconds(10)) #expect(session.isResponding == false) }