diff --git a/Sources/OpenAI/Private/Networking/URLSessionHTTPClientAdapter.swift b/Sources/OpenAI/Private/Networking/URLSessionHTTPClientAdapter.swift index b806a06..de1cc9a 100644 --- a/Sources/OpenAI/Private/Networking/URLSessionHTTPClientAdapter.swift +++ b/Sources/OpenAI/Private/Networking/URLSessionHTTPClientAdapter.swift @@ -49,16 +49,33 @@ public class URLSessionHTTPClientAdapter: HTTPClient { headers: convertHeaders(httpURLResponse.allHeaderFields)) let stream = AsyncThrowingStream { continuation in - Task { + let producer = Task { + defer { asyncBytes.task.cancel() } + var decoder = HTTPLineDecoder() + var chunk = [UInt8]() do { - for try await line in asyncBytes.lines { - continuation.yield(line) + for try await byte in asyncBytes { + try Task.checkCancellation() + chunk.append(byte) + if byte == 0x0A || byte == 0x0D || chunk.count >= 16_384 { + for line in decoder.append(chunk) { + if case .terminated = continuation.yield(line) { return } + } + chunk.removeAll(keepingCapacity: true) + } } + try Task.checkCancellation() + for line in decoder.append(chunk) { continuation.yield(line) } + if let tail = decoder.finish() { continuation.yield(tail) } continuation.finish() } catch { continuation.finish(throwing: error) } } + continuation.onTermination = { _ in + producer.cancel() + asyncBytes.task.cancel() + } } return (.lines(stream), response) diff --git a/Sources/OpenAI/Public/Service/OpenAIService.swift b/Sources/OpenAI/Public/Service/OpenAIService.swift index a0fdc6a..770048c 100644 --- a/Sources/OpenAI/Public/Service/OpenAIService.swift +++ b/Sources/OpenAI/Public/Service/OpenAIService.swift @@ -1290,19 +1290,7 @@ extension OpenAIService { } guard response.statusCode == 200 else { - var errorMessage = "status code \(response.statusCode)" - do { - // For error responses, we need to get the raw data instead of using the stream - // as error responses are regular JSON, not streaming data - let (errorData, _) = try await httpClient.data(for: httpRequest) - let error = try decoder.decode(OpenAIErrorResponse.self, from: errorData) - errorMessage = error.error.message ?? "NO ERROR MESSAGE PROVIDED" - } catch { - // If decoding fails, keep the original error message with status code - } - throw APIError.responseUnsuccessful( - description: errorMessage, - statusCode: response.statusCode) + throw try await streamResponseError(from: byteStream, statusCode: response.statusCode) } // Create a stream from the lines @@ -1312,58 +1300,47 @@ extension OpenAIService { return AsyncThrowingStream { continuation in let fetchTask = Task { + var eventData = [String]() + var firstLine = true + func emitEvent() throws -> Bool { + guard !eventData.isEmpty else { return false } + let payload = eventData.joined(separator: "\n") + eventData.removeAll(keepingCapacity: true) + if payload == "[DONE]" { return true } + let data = Data(payload.utf8) + #if DEBUG + if debugEnabled { + try print( + "DEBUG JSON STREAM EVENT = \(JSONSerialization.jsonObject(with: data, options: .allowFragments) as? [String: Any])") + } + #endif + let decoded = try self.decoder.decode(T.self, from: data) + if case .terminated = continuation.yield(decoded) { return true } + return false + } do { - for try await line in lineStream { - if - line.hasPrefix("data:"), line != "data: [DONE]", - let data = String(line.dropFirst(5)).data(using: .utf8) - { - #if DEBUG - if debugEnabled { - try print( - "DEBUG JSON STREAM LINE = \(JSONSerialization.jsonObject(with: data, options: .allowFragments) as? [String: Any])") - } - #endif - do { - let decoded = try self.decoder.decode(T.self, from: data) - continuation.yield(decoded) - } catch DecodingError.keyNotFound(let key, let context) { - let debug = "Key '\(key.stringValue)' not found: \(context.debugDescription)" - let codingPath = "codingPath: \(context.codingPath)" - let debugMessage = debug + codingPath - #if DEBUG - if debugEnabled { - print(debugMessage) - } - #endif - throw APIError.dataCouldNotBeReadMissingData(description: debugMessage) - } catch { - #if DEBUG - if debugEnabled { - debugPrint("CONTINUATION ERROR DECODING \(error.localizedDescription)") - } - #endif - continuation.finish(throwing: error) - } + for try await rawLine in lineStream { + try Task.checkCancellation() + let line = firstLine && rawLine.hasPrefix("\u{FEFF}") ? String(rawLine.dropFirst()) : rawLine + firstLine = false + if line.isEmpty { + if try emitEvent() { break } + } else if line == "data" { + eventData.append("") + } else if line.hasPrefix("data:") { + let value = line.dropFirst(5) + eventData.append(String(value.first == " " ? value.dropFirst() : value)) } } + try Task.checkCancellation() + // Preserve the existing behavior for a final event without a trailing blank line. + _ = try emitEvent() continuation.finish() } catch DecodingError.keyNotFound(let key, let context) { let debug = "Key '\(key.stringValue)' not found: \(context.debugDescription)" - let codingPath = "codingPath: \(context.codingPath)" - let debugMessage = debug + codingPath - #if DEBUG - if debugEnabled { - print(debugMessage) - } - #endif - throw APIError.dataCouldNotBeReadMissingData(description: debugMessage) + let message = debug + "codingPath: \(context.codingPath)" + continuation.finish(throwing: APIError.dataCouldNotBeReadMissingData(description: message)) } catch { - #if DEBUG - if debugEnabled { - print("CONTINUATION ERROR DECODING \(error.localizedDescription)") - } - #endif continuation.finish(throwing: error) } } @@ -1388,19 +1365,7 @@ extension OpenAIService { printHTTPResponse(response) guard response.statusCode == 200 else { - var errorMessage = "status code \(response.statusCode)" - do { - // For error responses, we need to get the raw data instead of using the stream - // as error responses are regular JSON, not streaming data - let (errorData, _) = try await httpClient.data(for: httpRequest) - let error = try decoder.decode(OpenAIErrorResponse.self, from: errorData) - errorMessage = error.error.message ?? "NO ERROR MESSAGE PROVIDED" - } catch { - // If decoding fails, keep the original error message with status code - } - throw APIError.responseUnsuccessful( - description: errorMessage, - statusCode: response.statusCode) + throw try await streamResponseError(from: byteStream, statusCode: response.statusCode) } // Create a stream from the lines @@ -1534,6 +1499,34 @@ extension OpenAIService { } } + private func streamResponseError(from stream: HTTPByteStream, statusCode: Int) async throws -> APIError { + var message = "status code \(statusCode)" + do { + // Consume the failed response, never send the same POST again to obtain its body. + var data = Data() + switch stream { + case .lines(let lines): + for try await line in lines { + try Task.checkCancellation() + data.append(contentsOf: line.utf8) + data.append(0x0A) + } + + case .bytes(let bytes): + for try await byte in bytes { + try Task.checkCancellation() + data.append(byte) + } + } + let error = try decoder.decode(OpenAIErrorResponse.self, from: data) + message = error.error.message ?? "NO ERROR MESSAGE PROVIDED" + } catch { + // Retain the HTTP status when the body cannot be read or decoded. + } + try Task.checkCancellation() + return .responseUnsuccessful(description: message, statusCode: statusCode) + } + // MARK: Debug Helpers private func prettyPrintJSON( diff --git a/Tests/OpenAITests/StreamingTransportTests.swift b/Tests/OpenAITests/StreamingTransportTests.swift new file mode 100644 index 0000000..1d47a17 --- /dev/null +++ b/Tests/OpenAITests/StreamingTransportTests.swift @@ -0,0 +1,172 @@ +#if compiler(>=6.0) +import Foundation +import SwiftOpenAI +import Testing + +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +// MARK: - StreamingTransportTests + +@Suite("Streaming transport regressions") +struct StreamingTransportTests { + enum Endpoint: CaseIterable { case responses, chat, assistants } + + @Test( + "HTTP failures read the original body without repeating the POST", + arguments: [400, 401, 403, 429, 500], + Endpoint.allCases) + func originalErrorResponse(status: Int, endpoint: Endpoint) async throws { + let client = RecordingStreamClient(status: status, lines: [ + "{", #" "error": {"message": "original failure"}"#, "}", + ]) + let service = OpenAIServiceFactory.service(apiKey: "test", httpClient: client) + do { + switch endpoint { + case .responses: + _ = try await service.responseCreateStream(.init(input: .string("hello"), model: .custom("test"))) + case .chat: + _ = try await service.startStreamedChat(parameters: .init(messages: [], model: .custom("test"))) + case .assistants: + var request = URLRequest(url: URL(string: "https://example.invalid/stream")!) + request.httpMethod = "POST" + _ = try await service.fetchAssistantStreamEvents(with: request, debugEnabled: false) + } + Issue.record("Expected an HTTP error") + } catch APIError.responseUnsuccessful(let description, let statusCode) { + #expect(statusCode == status) + #expect(description == "original failure") + } + #expect(await client.streamRequests == 1) + #expect(await client.dataRequests == 0) + } + + @Test("An SSE event joins data fields before decoding JSON") + func multilineEvent() async throws { + let values = try await collect([ + ": heartbeat", "event: message", "id: 1", "retry: 1000", + #"data: {"value":"#, + #"data: "café 😀"}"#, "", + #"data:{"value":"second"}"#, "", + ]) + #expect(values == ["café 😀", "second"]) + } + + @Test("HTTP error bodies support raw bytes and retain status for non-JSON bodies", arguments: [false, true]) + func errorBodyFormats(rawBytes: Bool) async throws { + for (body, message) in [(#"{"error":{"message":"original failure"}}"#, "original failure"), ("not JSON", "status code 500")] { + let client = RecordingStreamClient(status: 500, lines: [body], rawBytes: rawBytes) + let service = OpenAIServiceFactory.service(apiKey: "test", httpClient: client) + do { + _ = try await service.responseCreateStream(.init(input: .string("hello"), model: .custom("test"))) + Issue.record("Expected an HTTP error") + } catch APIError.responseUnsuccessful(let description, let statusCode) { + #expect(description == message) + #expect(statusCode == 500) + } + #expect(await client.streamRequests == 1) + #expect(await client.dataRequests == 0) + } + } + + @Test("Leading BOM and empty data fields preserve the SSE payload") + func eventFields() async throws { + #expect(try await collect(["\u{FEFF}: heartbeat", "data", #"data: {"value":"ok"}"#, ""]) == ["ok"]) + } + + @Test("Chat Completions still emits text and usage chunks") + func chatCompletions() async throws { + let client = RecordingStreamClient(status: 200, lines: [ + #"data: {"id":"c","object":"chat.completion.chunk","created":1,"model":"test","choices":[{"index":0,"delta":{"content":"hello"}}]}"#, + "", + #"data: {"id":"c","object":"chat.completion.chunk","created":1,"model":"test","choices":[],"usage":{"prompt_tokens":1,"completion_tokens":2,"total_tokens":3}}"#, + "", + "data: [DONE]", "", + ]) + let service = OpenAIServiceFactory.service(apiKey: "test", httpClient: client) + let stream = try await service.startStreamedChat(parameters: .init(messages: [], model: .custom("test"))) + var chunks = [ChatCompletionChunkObject]() + for try await chunk in stream { chunks.append(chunk) } + #expect(chunks.count == 2) + #expect(chunks.first?.choices?.first?.delta?.content == "hello") + #expect(chunks.last?.usage?.totalTokens == 3) + } + + @Test("DONE terminates consumption before trailing events") + func doneTerminates() async throws { + let values = try await collect([ + #"data: {"value":"first"}"#, "", "data: [DONE]", "", + #"data: {"value":"must not be delivered"}"#, "", + ]) + #expect(values == ["first"]) + } + + @Test("A final unterminated line keeps the existing EOF behavior") + func finalLineAtEOF() async throws { + #expect(try await collect([#"data: {"value":"last"}"#]) == ["last"]) + } + + @Test("Malformed known payloads fail the consumer") + func malformedPayload() async throws { + do { + _ = try await collect([#"data: {"value":17}"#, ""]) + Issue.record("Expected a decoding error") + } catch is DecodingError { } + } + + private struct Payload: Decodable { + let value: String + } + + private func collect(_ lines: [String]) async throws -> [String] { + let client = RecordingStreamClient(status: 200, lines: lines) + let service = OpenAIServiceFactory.service(apiKey: "test", httpClient: client) + var request = URLRequest(url: URL(string: "https://example.invalid/stream")!) + request.httpMethod = "POST" + let stream = try await service.fetchStream(debugEnabled: false, type: Payload.self, with: request) + var values = [String]() + for try await payload in stream { values.append(payload.value) } + return values + } +} + +// MARK: - RecordingStreamClient + +private actor RecordingStreamClient: HTTPClient { + init(status: Int, lines: [String], rawBytes: Bool = false) { + self.status = status + self.lines = lines + self.rawBytes = rawBytes + } + + private(set) var streamRequests = 0 + private(set) var dataRequests = 0 + + func data(for _: HTTPRequest) async throws -> (Data, HTTPResponse) { + dataRequests += 1 + return (Data(#"{"error":{"message":"second request"}}"#.utf8), .init(statusCode: status, headers: [:])) + } + + func bytes(for _: HTTPRequest) async throws -> (HTTPByteStream, HTTPResponse) { + streamRequests += 1 + if rawBytes { + let stream = AsyncThrowingStream { continuation in + for byte in lines.joined(separator: "\n").utf8 { continuation.yield(byte) } + continuation.finish() + } + return (.bytes(stream), .init(statusCode: status, headers: [:])) + } + let stream = AsyncThrowingStream { continuation in + for line in lines { continuation.yield(line) } + continuation.finish() + } + return (.lines(stream), .init(statusCode: status, headers: [:])) + } + + private let status: Int + private let lines: [String] + private let rawBytes: Bool + +} +#endif diff --git a/Tests/OpenAITests/URLSessionStreamingTests.swift b/Tests/OpenAITests/URLSessionStreamingTests.swift new file mode 100644 index 0000000..5ccd3bf --- /dev/null +++ b/Tests/OpenAITests/URLSessionStreamingTests.swift @@ -0,0 +1,197 @@ +#if os(macOS) && compiler(>=6.0) +import Foundation +import SwiftOpenAI +import Testing + +@Suite("URLSession streaming", .timeLimit(.minutes(1))) +struct URLSessionStreamingTests { + @Test("The Apple adapter preserves empty lines, CRLF, CR and split UTF-8") + func lineBoundaries() async throws { + let server = try await StreamingLoopbackServer() + defer { server.stop() } + let session = URLSession(configuration: .ephemeral) + defer { session.invalidateAndCancel() } + let adapter = URLSessionHTTPClientAdapter(urlSession: session) + let (body, _) = try await adapter.bytes(for: .init( + url: server.url.appendingPathComponent("lines"), + method: .get, + headers: [:])) + guard case .lines(let stream) = body else { + Issue.record("Expected lines") + return + } + var lines = [String]() + for try await line in stream { lines.append(line) } + #expect(lines == ["data: café 😀", "", "data: two", "", "tail"]) + } + + @Test("Responses decodes multiline events through the Apple adapter") + func multilineResponse() async throws { + let server = try await StreamingLoopbackServer() + defer { server.stop() } + let service = OpenAIServiceFactory.service(apiKey: "test", overrideBaseURL: server.url.absoluteString) + let stream = try await service.responseCreateStream(.init(input: .string("multiline"), model: .custom("test"))) + var text = "" + for try await event in stream { + if case .outputTextDelta(let delta) = event { text += delta.delta } + } + #expect(text == "café 😀") + } + + @Test("Cancelling one Responses consumer closes only its connection") + func cancellationIsPerRequest() async throws { + let server = try await StreamingLoopbackServer() + defer { server.stop() } + let session = URLSession(configuration: .ephemeral) + defer { session.invalidateAndCancel() } + let service = OpenAIServiceFactory.service( + apiKey: "test", overrideBaseURL: server.url.absoluteString, + httpClient: URLSessionHTTPClientAdapter(urlSession: session)) + func consume(_ name: String) async { + do { + let stream = try await service.responseCreateStream(.init(input: .string(name), model: .custom("test"))) + for try await _ in stream { } + } catch { } + } + let first = Task { await consume("first") } + let second = Task { await consume("second") } + defer { first.cancel() + second.cancel() + } + try #require(try await eventually { + let state = try await server.snapshot() + return state.frames["first", default: 0] >= 2 && state.frames["second", default: 0] >= 2 + }) + first.cancel() + await first.value + #expect(try await eventually { try await server.snapshot().closed["first"] == true }) + let before = try await server.snapshot() + #expect(before.closed["second"] != true) + #expect(try await eventually { + try await server.snapshot().frames["second", default: 0] > before.frames["second", default: 0] + 2 + }) + #expect(await session.allTasks.filter { $0.state == .running }.count == 1) + second.cancel() + await second.value + #expect(try await eventually { try await server.snapshot().closed["second"] == true }) + #expect(await eventually { await session.allTasks.isEmpty }) + let final = try await server.snapshot() + #expect(final.requests == ["first": 1, "second": 1]) + } + + private func eventually(_ condition: () async throws -> Bool) async rethrows -> Bool { + let deadline = ContinuousClock.now.advanced(by: .seconds(2)) + while ContinuousClock.now < deadline { + if try await condition() { return true } + try? await Task.sleep(for: .milliseconds(20)) + } + return try await condition() + } +} + +private final class StreamingLoopbackServer { + init() async throws { + directory = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString) + try FileManager.default.createDirectory(at: directory, withIntermediateDirectories: true) + let script = directory.appendingPathComponent("server.py") + try Self.source.write(to: script, atomically: true, encoding: .utf8) + process = Process() + process.executableURL = URL(fileURLWithPath: "/usr/bin/env") + process.arguments = ["python3", script.path] + process.standardOutput = FileHandle.nullDevice + process.standardError = FileHandle.nullDevice + do { + try process.run() + let portFile = directory.appendingPathComponent("port") + let deadline = ContinuousClock.now.advanced(by: .seconds(3)) + while ContinuousClock.now < deadline, process.isRunning { + if let port = try? String(contentsOf: portFile, encoding: .utf8), let number = Int(port) { + url = URL(string: "http://127.0.0.1:\(number)")! + return + } + try await Task.sleep(for: .milliseconds(20)) + } + throw URLError(.cannotConnectToHost) + } catch { + stop() + throw error + } + } + + struct Snapshot: Decodable { + let frames: [String: Int] + let closed: [String: Bool] + let requests: [String: Int] + } + + private(set) var url = URL(string: "http://127.0.0.1")! + + func snapshot() async throws -> Snapshot { + let (data, _) = try await URLSession.shared.data(from: url.appendingPathComponent("stats")) + return try JSONDecoder().decode(Snapshot.self, from: data) + } + + func stop() { + if process.isRunning { process.terminate() + process.waitUntilExit() + } + try? FileManager.default.removeItem(at: directory) + } + + private static let source = #""" + import json, threading, time + from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + from pathlib import Path + + state = {"frames": {}, "closed": {}, "requests": {}} + lock = threading.Lock() + + class Handler(BaseHTTPRequestHandler): + def log_message(self, *args): pass + + def do_GET(self): + if self.path == "/stats": + with lock: body = json.dumps(state).encode() + else: + body = "data: café 😀\r\n\r\ndata: two\r\rtail".encode() + self.send_response(200) + self.send_header("Content-Length", str(len(body))) + self.end_headers() + for byte in body: + self.wfile.write(bytes([byte])) + self.wfile.flush() + + def do_POST(self): + request = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) + name = request["input"] + with lock: state["requests"][name] = state["requests"].get(name, 0) + 1 + self.send_response(200) + self.send_header("Content-Type", "text/event-stream") + self.end_headers() + if name == "multiline": + body = ('data: {"type":"response.output_text.delta",\r\n' + 'data: "item_id":"m","output_index":0,"content_index":0,"delta":"café 😀"}\r\n\r\n' + 'data: [DONE]\r\n\r\n') + self.wfile.write(body.encode()) + return + try: + for _ in range(1000): + event = {"type": "response.output_text.delta", "item_id": "m", + "output_index": 0, "content_index": 0, "delta": "x"} + self.wfile.write(("data: " + json.dumps(event) + "\n\n").encode()) + self.wfile.flush() + with lock: state["frames"][name] = state["frames"].get(name, 0) + 1 + time.sleep(0.02) + except (BrokenPipeError, ConnectionResetError): + with lock: state["closed"][name] = True + + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + Path(__file__).with_name("port").write_text(str(server.server_port)) + server.serve_forever() + """# + + private let directory: URL + private let process: Process + +} +#endif