Actually stream the data (#31)

* Actually stream the data

* Apparently `Task { @mainactor in ... }` is non-blocking but `await MainActor.run` is blocking

* Revert last change, it's not the problem

* Should only need one while loop

* Use EventSource convenience method

* Add Linux specific implementation for request streaming

---------

Co-authored-by: Mattt Zmuda <[email protected]>

Author: nathanborror

Sources/AnyLanguageModel/Extensions/URLSession+Extensions.swift

@@ -142,30 +142,31 @@ extension URLSession {
                         request.addValue("application/json", forHTTPHeaderField: "Content-Type")
                     }
 
-                    let (asyncBytes, response) = try await self.data(for: request)
+                    #if canImport(FoundationNetworking)
+                        let (asyncBytes, response) = try await self.linuxBytes(for: request)
+                    #else
+                        let (asyncBytes, response) = try await self.bytes(for: request)
+                    #endif
 
                     guard let httpResponse = response as? HTTPURLResponse else {
                         throw URLSessionError.invalidResponse
                     }
 
                     guard (200 ..< 300).contains(httpResponse.statusCode) else {
-                        if let errorString = String(data: asyncBytes, encoding: .utf8) {
+                        var errorData = Data()
+                        for try await byte in asyncBytes {
+                            errorData.append(byte)
+                        }
+                        if let errorString = String(data: errorData, encoding: .utf8) {
                             throw URLSessionError.httpError(statusCode: httpResponse.statusCode, detail: errorString)
                         }
                         throw URLSessionError.httpError(statusCode: httpResponse.statusCode, detail: "Invalid response")
                     }
 
                     let decoder = JSONDecoder()
-                    let parser = EventSource.Parser()
-
-                    for byte in asyncBytes {
-                        await parser.consume(byte)
-                    }
-                    await parser.finish()
 
-                    while let event = await parser.getNextEvent() {
+                    for try await event in asyncBytes.events {
                         guard let data = event.data.data(using: .utf8) else { continue }
-
                         if let decoded = try? decoder.decode(T.self, from: data) {
                             continuation.yield(decoded)
                         }
@@ -184,6 +185,118 @@ extension URLSession {
     }
 }
 
+#if canImport(FoundationNetworking)
+    private extension URLSession {
+        func linuxBytes(for request: URLRequest) async throws -> (AsyncThrowingStream<UInt8, Error>, URLResponse) {
+            let delegate = LinuxBytesDelegate()
+            let delegateQueue = OperationQueue()
+            delegateQueue.maxConcurrentOperationCount = 1
+
+            let session = URLSession(
+                configuration: self.configuration,
+                delegate: delegate,
+                delegateQueue: delegateQueue
+            )
+
+            let byteStream = AsyncThrowingStream<UInt8, Error> { continuation in
+                delegate.attach(
+                    continuation,
+                    session: session
+                )
+            }
+
+            let response = try await delegate.start(
+                request: request,
+                session: session
+            )
+
+            return (byteStream, response)
+        }
+    }
+
+    private final class LinuxBytesDelegate: NSObject, URLSessionDataDelegate, @unchecked Sendable {
+        private var responseContinuation: CheckedContinuation<URLResponse, Error>?
+        private var byteContinuation: AsyncThrowingStream<UInt8, Error>.Continuation?
+        private weak var task: URLSessionDataTask?
+        private weak var session: URLSession?
+
+        func attach(
+            _ continuation: AsyncThrowingStream<UInt8, Error>.Continuation,
+            session: URLSession
+        ) {
+            byteContinuation = continuation
+            self.session = session
+            continuation.onTermination = { [weak self] _ in
+                guard let self else { return }
+                self.task?.cancel()
+                self.session?.invalidateAndCancel()
+            }
+        }
+
+        func start(
+            request: URLRequest,
+            session: URLSession
+        ) async throws -> URLResponse {
+            try await withCheckedThrowingContinuation { continuation in
+                responseContinuation = continuation
+                let task = session.dataTask(with: request)
+                self.task = task
+                task.resume()
+            }
+        }
+
+        func urlSession(
+            _ session: URLSession,
+            dataTask: URLSessionDataTask,
+            didReceive response: URLResponse,
+            completionHandler: @escaping @Sendable (URLSession.ResponseDisposition) -> Void
+        ) {
+            if let continuation = responseContinuation {
+                continuation.resume(returning: response)
+                responseContinuation = nil
+            }
+            completionHandler(.allow)
+        }
+
+        func urlSession(
+            _ session: URLSession,
+            dataTask: URLSessionDataTask,
+            didReceive data: Data
+        ) {
+            guard let continuation = byteContinuation else { return }
+            for byte in data {
+                continuation.yield(byte)
+            }
+        }
+
+        func urlSession(
+            _ session: URLSession,
+            task: URLSessionTask,
+            didCompleteWithError error: (any Error)?
+        ) {
+            if let continuation = responseContinuation {
+                if let error {
+                    continuation.resume(throwing: error)
+                } else if let response = task.response {
+                    continuation.resume(returning: response)
+                } else {
+                    continuation.resume(throwing: URLSessionError.invalidResponse)
+                }
+                responseContinuation = nil
+            }
+
+            if let error {
+                byteContinuation?.finish(throwing: error)
+            } else {
+                byteContinuation?.finish()
+            }
+            byteContinuation = nil
+
+            session.invalidateAndCancel()
+        }
+    }
+#endif
+
 enum URLSessionError: Error, CustomStringConvertible {
     case invalidResponse
     case httpError(statusCode: Int, detail: String)