Incorporate feedback from review

Author: mattt

Sources/AnyLanguageModel/Models/MLXLanguageModel.swift

@@ -389,7 +389,7 @@ import Foundation
             func removeEntries(forModelKey modelKey: String) {
                 lock.withLock {
                     reapDeadSessionsLocked()
-                    for id in buckets.keys {
+                    for id in Array(buckets.keys) {
                         guard var bucket = buckets[id] else {
                             continue
                         }
@@ -410,7 +410,10 @@ import Foundation
             }
 
             private func reapDeadSessionsLocked() {
-                for (id, bucket) in buckets where bucket.sessionReference.session == nil {
+                let deadSessionIDs = buckets.compactMap { id, bucket in
+                    bucket.sessionReference.session == nil ? id : nil
+                }
+                for id in deadSessionIDs {
                     buckets[id] = nil
                 }
             }
@@ -442,7 +445,7 @@ import Foundation
             static let shared = GPUMemoryManager()
 
             private let lock = NSLock()
-            private var knownConfigs: Set<GPUMemoryConfiguration> = [.automatic]
+            private var knownConfigs: Set<GPUMemoryConfiguration> = []
             private var activeScopes: [UUID: GPUMemoryConfiguration] = [:]
 
             private init() {
@@ -500,7 +503,10 @@ import Foundation
             }
 
             private func shouldClearOnEviction() -> Bool {
-                knownConfigs.contains { $0.clearCacheOnEviction }
+                if knownConfigs.isEmpty {
+                    return GPUMemoryConfiguration.automatic.clearCacheOnEviction
+                }
+                return knownConfigs.contains { $0.clearCacheOnEviction }
             }
         }
 
@@ -852,17 +858,6 @@ import Foundation
                     }
                     previousToolCallSignature = signature
 
-                    if !assistantText.isEmpty {
-                        allEntries.append(
-                            .response(
-                                Transcript.Response(
-                                    assetIDs: [],
-                                    segments: [.text(.init(content: assistantText))]
-                                )
-                            )
-                        )
-                    }
-
                     let resolution = try await resolveToolCalls(collectedToolCalls, session: session)
                     switch resolution {
                     case .stop(let calls):

Tests/AnyLanguageModelTests/MLXLanguageModelTests.swift

@@ -78,6 +78,32 @@ import Testing
             #expect(!second.content.isEmpty)
         }
 
+        @Test func rejectsConcurrentRequestsForSameSession() async throws {
+            let session = LanguageModelSession(model: model)
+            let stream = session.streamResponse(
+                to: "Count from 1 to 400 with one number per line.",
+                options: .init(maximumResponseTokens: 256)
+            )
+
+            do {
+                _ = try await session.respond(to: "This concurrent request should fail.")
+                Issue.record("Expected concurrent request to throw.")
+            } catch let error as LanguageModelSession.GenerationError {
+                switch error {
+                case .concurrentRequests:
+                    break
+                default:
+                    Issue.record("Expected .concurrentRequests, got \(error)")
+                }
+            } catch {
+                Issue.record("Expected GenerationError.concurrentRequests, got \(error)")
+            }
+
+            for try await _ in stream {
+                break
+            }
+        }
+
         @Test func withGenerationOptions() async throws {
             let session = LanguageModelSession(model: model)