Fix duplicate streaming requests, cancellation and SSE framing

Sources/OpenAI/Private/Networking/URLSessionHTTPClientAdapter.swift

@@ -49,16 +49,33 @@ public class URLSessionHTTPClientAdapter: HTTPClient {
       headers: convertHeaders(httpURLResponse.allHeaderFields))
 
     let stream = AsyncThrowingStream<String, Error> { 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)

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(

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<UInt8, Error> { 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<String, Error> { 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