diff --git a/Sources/AnyLanguageModel/LanguageModelSession.swift b/Sources/AnyLanguageModel/LanguageModelSession.swift index 1b569491..1c00de75 100644 --- a/Sources/AnyLanguageModel/LanguageModelSession.swift +++ b/Sources/AnyLanguageModel/LanguageModelSession.swift @@ -3,8 +3,17 @@ import Observation @Observable public final class LanguageModelSession: @unchecked Sendable { - public private(set) var isResponding: Bool = false - public private(set) var transcript: Transcript + public var isResponding: Bool { + access(keyPath: \.isResponding) + return state.access { $0.isResponding } + } + + public var transcript: Transcript { + access(keyPath: \.transcript) + return state.access { $0.transcript } + } + + @ObservationIgnored private let state: Locked private let model: any LanguageModel public let tools: [any Tool] @@ -20,8 +29,6 @@ public final class LanguageModelSession: @unchecked Sendable { /// with the Foundation Models framework. @ObservationIgnored public var toolExecutionDelegate: (any ToolExecutionDelegate)? - @ObservationIgnored private let respondingState = RespondingState() - public convenience init( model: any LanguageModel, tools: [any Tool] = [], @@ -84,37 +91,33 @@ public final class LanguageModelSession: @unchecked Sendable { } } - self.transcript = finalTranscript + self.state = .init(.init(finalTranscript)) } public func prewarm(promptPrefix: Prompt? = nil) { model.prewarm(for: self, promptPrefix: promptPrefix) } - nonisolated private func beginResponding() async { - let count = await respondingState.increment() - let active = count > 0 - await MainActor.run { - self.isResponding = active + nonisolated private func beginResponding() { + withMutation(keyPath: \.isResponding) { + state.access { $0.beginResponding() } } } - nonisolated private func endResponding() async { - let count = await respondingState.decrement() - let active = count > 0 - await MainActor.run { - self.isResponding = active + nonisolated private func endResponding() { + withMutation(keyPath: \.isResponding) { + state.access { $0.endResponding() } } } nonisolated private func wrapRespond(_ operation: () async throws -> T) async throws -> T { - await beginResponding() + beginResponding() do { let result = try await operation() - await endResponding() + endResponding() return result } catch { - await endResponding() + endResponding() throw error } } @@ -127,7 +130,7 @@ public final class LanguageModelSession: @unchecked Sendable { let relay = AsyncThrowingStream.Snapshot, any Error> { continuation in let stream = upstream Task { - await session.beginResponding() + session.beginResponding() var lastSnapshot: ResponseStream.Snapshot? do { for try await snapshot in stream { @@ -152,14 +155,14 @@ public final class LanguageModelSession: @unchecked Sendable { segments: [.text(.init(content: textContent))] ) ) - await MainActor.run { - session.transcript.append(responseEntry) + session.withMutation(keyPath: \.transcript) { + session.state.access { $0.transcript.append(responseEntry) } } } } catch { continuation.finish(throwing: error) } - await session.endResponding() + session.endResponding() } } return ResponseStream(stream: relay) @@ -202,8 +205,8 @@ public final class LanguageModelSession: @unchecked Sendable { responseFormat: nil ) ) - await MainActor.run { - self.transcript.append(promptEntry) + withMutation(keyPath: \.transcript) { + state.access { $0.transcript.append(promptEntry) } } let response = try await model.respond( @@ -230,9 +233,11 @@ public final class LanguageModelSession: @unchecked Sendable { ) // Add tool entries and response to transcript - await MainActor.run { - self.transcript.append(contentsOf: response.transcriptEntries) - self.transcript.append(responseEntry) + withMutation(keyPath: \.transcript) { + state.access { state in + state.transcript.append(contentsOf: response.transcriptEntries) + state.transcript.append(responseEntry) + } } return response @@ -253,7 +258,9 @@ public final class LanguageModelSession: @unchecked Sendable { responseFormat: nil ) ) - transcript.append(promptEntry) + withMutation(keyPath: \.transcript) { + state.access { $0.transcript.append(promptEntry) } + } return wrapStream( model.streamResponse( @@ -547,8 +554,8 @@ extension LanguageModelSession { responseFormat: nil ) ) - await MainActor.run { - self.transcript.append(promptEntry) + withMutation(keyPath: \.transcript) { + state.access { $0.transcript.append(promptEntry) } } // Extract text content for the Prompt parameter @@ -578,9 +585,11 @@ extension LanguageModelSession { ) // Add tool entries and response to transcript - await MainActor.run { - self.transcript.append(contentsOf: response.transcriptEntries) - self.transcript.append(responseEntry) + withMutation(keyPath: \.transcript) { + state.access { state in + state.transcript.append(contentsOf: response.transcriptEntries) + state.transcript.append(responseEntry) + } } return response @@ -651,7 +660,9 @@ extension LanguageModelSession { responseFormat: nil ) ) - transcript.append(promptEntry) + withMutation(keyPath: \.transcript) { + state.access { $0.transcript.append(promptEntry) } + } // Extract text content for the Prompt parameter let textPrompt = Prompt(prompt) @@ -902,16 +913,21 @@ private enum ResponseStreamError: Error { // MARK: - -private actor RespondingState { +private struct State: Equatable, Sendable { + var transcript: Transcript + + var isResponding: Bool { count > 0 } private var count = 0 - func increment() -> Int { + init(_ transcript: Transcript) { + self.transcript = transcript + } + + mutating func beginResponding() { count += 1 - return count } - func decrement() -> Int { + mutating func endResponding() { count = max(0, count - 1) - return count } } diff --git a/Sources/AnyLanguageModel/Locked.swift b/Sources/AnyLanguageModel/Locked.swift new file mode 100644 index 00000000..25d93d46 --- /dev/null +++ b/Sources/AnyLanguageModel/Locked.swift @@ -0,0 +1,16 @@ +import Foundation + +final class Locked { + private let lock = NSLock() + private var state: State + + init(_ state: State) { + self.state = state + } + + func access(_ block: (inout State) throws -> T) rethrows -> T { + try lock.withLock { try block(&self.state) } + } +} + +extension Locked: @unchecked Sendable where State: Sendable {} diff --git a/Tests/AnyLanguageModelTests/LockedTests.swift b/Tests/AnyLanguageModelTests/LockedTests.swift new file mode 100644 index 00000000..b81eda9e --- /dev/null +++ b/Tests/AnyLanguageModelTests/LockedTests.swift @@ -0,0 +1,152 @@ +import Foundation +import Testing + +@testable import AnyLanguageModel + +@Suite("Locked") +struct LockedTests { + @Test("Read access returns the initial value") + func readAccess() { + let locked = Locked(42) + let value = locked.access { $0 } + #expect(value == 42) + } + + @Test("Write access mutates the state") + func writeAccess() { + let locked = Locked(0) + locked.access { $0 = 99 } + let value = locked.access { $0 } + #expect(value == 99) + } + + @Test("Access returns the value from the closure") + func returnValue() { + let locked = Locked("hello") + let result = locked.access { state -> Int in + state += " world" + return state.count + } + #expect(result == 11) + #expect(locked.access { $0 } == "hello world") + } + + @Test("Access propagates thrown errors") + func throwsPropagation() { + struct TestError: Error, Equatable {} + + let locked = Locked(0) + #expect(throws: TestError.self) { + try locked.access { _ in throw TestError() } + } + } + + @Test("Works with complex state types") + func complexState() { + struct State { + var name: String + var count: Int + var tags: [String] + } + + let locked = Locked(State(name: "initial", count: 0, tags: [])) + locked.access { state in + state.name = "updated" + state.count = 5 + state.tags.append("a") + state.tags.append("b") + } + + let snapshot = locked.access { $0 } + #expect(snapshot.name == "updated") + #expect(snapshot.count == 5) + #expect(snapshot.tags == ["a", "b"]) + } + + @Test("Concurrent increments produce the correct final count") + func concurrentCounter() async { + let locked = Locked(0) + let iterations = 100_000 + + await withTaskGroup(of: Void.self) { group in + for _ in 0 ..< iterations { + group.addTask { locked.access { $0 += 1 } } + } + } + + let finalValue = locked.access { $0 } + #expect(finalValue == iterations) + } + + @Test("Concurrent increments with mixed task priorities") + func concurrentCounterWithMixedPriorities() async { + let locked = Locked(0) + let iterations = 100_000 + + await withTaskGroup(of: Void.self) { group in + for i in 0 ..< iterations { + let priority: TaskPriority = i.isMultiple(of: 2) ? .high : .background + group.addTask(priority: priority) { + locked.access { $0 += 1 } + } + } + } + + let finalValue = locked.access { $0 } + #expect(finalValue == iterations) + } + + @Test("Concurrent appends to an array") + func concurrentArrayAppend() async { + let locked = Locked<[Int]>([]) + let iterations = 10_000 + + await withTaskGroup(of: Void.self) { group in + for i in 0 ..< iterations { + group.addTask { locked.access { $0.append(i) } } + } + } + + let finalArray = locked.access { $0 } + #expect(finalArray.count == iterations) + } + + @Test("Multiple independent locked values mutated concurrently") + func multipleLockedValues() async { + let lockedA = Locked(0) + let lockedB = Locked(0) + let iterations = 50_000 + + await withTaskGroup(of: Void.self) { group in + for _ in 0 ..< iterations { + group.addTask { lockedA.access { $0 += 1 } } + group.addTask { lockedB.access { $0 += 1 } } + } + } + + #expect(lockedA.access { $0 } == iterations) + #expect(lockedB.access { $0 } == iterations) + } + + @Test("Can wrap a non-Sendable type") + func nonSendableElement() { + class Box { + var value: Int + init(_ value: Int) { self.value = value } + } + + let locked = Locked(Box(10)) + locked.access { $0.value += 5 } + let result = locked.access { $0.value } + #expect(result == 15) + } + + @Test("Copies share the same underlying storage") + func copySharesStorage() { + let original = Locked(0) + let copy = original + original.access { $0 = 42 } + let value = copy.access { $0 } + #expect(value == 42) + } +} diff --git a/Tests/AnyLanguageModelTests/ObservationTests.swift b/Tests/AnyLanguageModelTests/ObservationTests.swift new file mode 100644 index 00000000..470a65ac --- /dev/null +++ b/Tests/AnyLanguageModelTests/ObservationTests.swift @@ -0,0 +1,95 @@ +import Observation +import Testing + +@testable import AnyLanguageModel + +@Suite("Observation") +struct ObservationTests { + @Test("Tracking transcript fires onChange when respond mutates it") + func transcriptObservationTriggeredByRespond() async throws { + let session = LanguageModelSession(model: MockLanguageModel.fixed("Hello")) + let changed = Locked(false) + + withObservationTracking { + _ = session.transcript + } onChange: { + changed.access { $0 = true } + } + + try await session.respond(to: "Hi") + #expect(changed.access { $0 } == true) + } + + @Test("Tracking transcript fires onChange when streamResponse mutates it") + func transcriptObservationTriggeredByStreamResponse() async throws { + let session = LanguageModelSession(model: MockLanguageModel.streamingMock()) + let changed = Locked(false) + + withObservationTracking { + _ = session.transcript + } onChange: { + changed.access { $0 = true } + } + + let stream = session.streamResponse(to: "Hi") + for try await _ in stream {} + try await Task.sleep(for: .milliseconds(10)) + + #expect(changed.access { $0 } == true) + } + + @Test("Tracking isResponding fires onChange when respond mutates it") + func isRespondingObservationTriggeredByRespond() async throws { + let session = LanguageModelSession(model: MockLanguageModel.fixed("Hello")) + let changed = Locked(false) + + withObservationTracking { + _ = session.isResponding + } onChange: { + changed.access { $0 = true } + } + + try await session.respond(to: "Hi") + #expect(changed.access { $0 } == true) + } + + @Test("No onChange fires when no properties are tracked") + func noObservationWithoutTrackedRead() async throws { + let session = LanguageModelSession(model: MockLanguageModel.fixed("Hello")) + let changed = Locked(false) + + withObservationTracking { + // Intentionally do not read any session properties + } onChange: { + changed.access { $0 = true } + } + + try await session.respond(to: "Hi") + #expect(changed.access { $0 } == false) + } + + @Test("Re-registering observation tracks subsequent changes") + func observationReregistersAfterChange() async throws { + let session = LanguageModelSession(model: MockLanguageModel.fixed("Hello")) + let changeCount = Locked(0) + + withObservationTracking { + _ = session.transcript + } onChange: { + changeCount.access { $0 += 1 } + } + + try await session.respond(to: "First") + #expect(changeCount.access { $0 } == 1) + + // Re-register after the first change fires + withObservationTracking { + _ = session.transcript + } onChange: { + changeCount.access { $0 += 1 } + } + + try await session.respond(to: "Second") + #expect(changeCount.access { $0 } == 2) + } +}