fix(client): translate streaming IOException into custom exception (#397)

Author: stainless-app[bot]

openai-java-core/src/main/kotlin/com/openai/core/handlers/StreamHandler.kt

@@ -6,6 +6,8 @@ import com.openai.core.http.HttpResponse
 import com.openai.core.http.HttpResponse.Handler
 import com.openai.core.http.PhantomReachableClosingStreamResponse
 import com.openai.core.http.StreamResponse
+import com.openai.errors.OpenAIIoException
+import java.io.IOException
 import java.util.stream.Stream
 import kotlin.streams.asStream
 
@@ -14,17 +16,30 @@ internal fun <T> streamHandler(
     block: suspend SequenceScope<T>.(response: HttpResponse, lines: Sequence<String>) -> Unit
 ): Handler<StreamResponse<T>> =
     object : Handler<StreamResponse<T>> {
+
         override fun handle(response: HttpResponse): StreamResponse<T> {
             val reader = response.body().bufferedReader()
             val sequence =
                 // Wrap in a `CloseableSequence` to avoid performing a read on the `reader`
                 // after it has been closed, which would throw an `IOException`.
                 CloseableSequence(
-                    sequence { reader.useLines { block(response, it) } }.constrainOnce()
+                    sequence {
+                            reader.useLines { lines ->
+                                block(
+                                    response,
+                                    // We wrap the `lines` instead of the top-level sequence because
+                                    // we only want to catch `IOException` from the reader; not from
+                                    // the user's own code.
+                                    IOExceptionWrappingSequence(lines),
+                                )
+                            }
+                        }
+                        .constrainOnce()
                 )
 
             return PhantomReachableClosingStreamResponse(
                 object : StreamResponse<T> {
+
                     override fun stream(): Stream<T> = sequence.asStream()
 
                     override fun close() {
@@ -37,18 +52,44 @@ internal fun <T> streamHandler(
         }
     }
 
+/** A sequence that catches, wraps, and rethrows [IOException] as [OpenAIIoException]. */
+private class IOExceptionWrappingSequence<T>(private val sequence: Sequence<T>) : Sequence<T> {
+
+    override fun iterator(): Iterator<T> {
+        val iterator = sequence.iterator()
+        return object : Iterator<T> {
+
+            override fun next(): T =
+                try {
+                    iterator.next()
+                } catch (e: IOException) {
+                    throw OpenAIIoException("Stream failed", e)
+                }
+
+            override fun hasNext(): Boolean =
+                try {
+                    iterator.hasNext()
+                } catch (e: IOException) {
+                    throw OpenAIIoException("Stream failed", e)
+                }
+        }
+    }
+}
+
 /**
  * A sequence that can be closed.
  *
  * Once [close] is called, it will not yield more elements. It will also no longer consult the
  * underlying [Iterator.hasNext] method.
  */
 private class CloseableSequence<T>(private val sequence: Sequence<T>) : Sequence<T> {
+
     private var isClosed: Boolean = false
 
     override fun iterator(): Iterator<T> {
         val iterator = sequence.iterator()
         return object : Iterator<T> {
+
             override fun next(): T = iterator.next()
 
             override fun hasNext(): Boolean = !isClosed && iterator.hasNext()

openai-java-core/src/test/kotlin/com/openai/core/handlers/StreamHandlerTest.kt

@@ -0,0 +1,64 @@
+package com.openai.core.handlers
+
+import com.openai.core.http.Headers
+import com.openai.core.http.HttpResponse
+import com.openai.errors.OpenAIIoException
+import java.io.IOException
+import java.io.InputStream
+import kotlin.test.Test
+import org.assertj.core.api.Assertions.assertThat
+import org.junit.jupiter.api.assertThrows
+
+internal class StreamHandlerTest {
+
+    @Test
+    fun streamHandler_whenReaderThrowsIOException_wrapsException() {
+        val handler = streamHandler<String> { _, lines -> lines.forEach {} }
+        val streamResponse = handler.handle(httpResponse("a\nb\nc\n".byteInputStream().throwing()))
+
+        val e = assertThrows<OpenAIIoException> { streamResponse.stream().forEach {} }
+        assertThat(e).hasMessage("Stream failed")
+        assertThat(e).hasCauseInstanceOf(IOException::class.java)
+    }
+
+    @Test
+    fun streamHandler_whenBlockThrowsIOException_doesNotWrapException() {
+        val ioException = IOException("BOOM!")
+        val handler =
+            streamHandler<String> { _, lines ->
+                lines.forEachIndexed { index, _ ->
+                    if (index == 2) {
+                        throw ioException
+                    }
+                }
+            }
+        val streamResponse = handler.handle(httpResponse("a\nb\nc\n".byteInputStream()))
+
+        val e = assertThrows<IOException> { streamResponse.stream().forEach {} }
+        assertThat(e).isSameAs(ioException)
+    }
+
+    private fun httpResponse(body: InputStream): HttpResponse =
+        object : HttpResponse {
+
+            override fun statusCode(): Int = 0
+
+            override fun headers(): Headers = Headers.builder().build()
+
+            override fun body(): InputStream = body
+
+            override fun close() {}
+        }
+
+    private fun InputStream.throwing(): InputStream =
+        object : InputStream() {
+
+            override fun read(): Int {
+                val byte = this@throwing.read()
+                if (byte == -1) {
+                    throw IOException("BOOM!")
+                }
+                return byte
+            }
+        }
+}