feat: expose response headers on parsing and stream errors

buildSrc/src/main/kotlin/com/openai/gradle/CoreCompilationShards.kt

@@ -94,6 +94,7 @@ object CoreCompilationShards {
             "com/openai/core/Utils.kt",
             "com/openai/core/Values.kt",
             "com/openai/core/http/AsyncStreamResponse.kt",
+            "com/openai/core/http/Headers.kt",
             "com/openai/core/http/HttpMethod.kt",
             "com/openai/core/http/HttpRequestBody.kt",
             "com/openai/core/http/MultipartInputStreams.kt",

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

@@ -16,5 +16,6 @@ internal inline fun <reified T> jsonHandler(jsonMapper: JsonMapper): Handler<T>
                 jsonMapper.readValue(response.body(), jacksonTypeRef())
             } catch (e: Exception) {
                 throw OpenAIInvalidDataException("Error reading response", e)
+                    .withHeaders(response.headers())
             }
     }

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

@@ -13,6 +13,7 @@ import com.openai.core.http.SseMessage
 import com.openai.core.http.StreamResponse
 import com.openai.core.http.map
 import com.openai.core.http.markBodyComplete
+import com.openai.errors.OpenAIInvalidDataException
 import com.openai.errors.SseException
 import com.openai.models.ErrorObject
 
@@ -135,5 +136,11 @@ internal inline fun <reified T> Handler<StreamResponse<SseMessage>>.mapJson(
 ): Handler<StreamResponse<T>> =
     object : Handler<StreamResponse<T>> {
         override fun handle(response: HttpResponse): StreamResponse<T> =
-            this@mapJson.handle(response).map { it.json<T>(includeEventAndData) }
+            this@mapJson.handle(response).map {
+                try {
+                    it.json<T>(includeEventAndData)
+                } catch (e: OpenAIInvalidDataException) {
+                    throw e.withHeaders(response.headers())
+                }
+            }
     }

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

@@ -2,6 +2,7 @@
 
 package com.openai.core.handlers
 
+import com.openai.core.http.Headers
 import com.openai.core.http.HttpResponse
 import com.openai.core.http.HttpResponse.Handler
 import com.openai.core.http.PhantomReachableClosingStreamResponse
@@ -46,7 +47,7 @@ internal fun <T> streamHandler(
                                     // 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, isClosed),
+                                    IOExceptionWrappingSequence(lines, response.headers(), isClosed),
                                 )
                             }
                         }
@@ -75,6 +76,7 @@ internal fun <T> streamHandler(
 /** A sequence that catches, wraps, and rethrows [IOException] as [OpenAIIoException]. */
 private class IOExceptionWrappingSequence<T>(
     private val sequence: Sequence<T>,
+    private val headers: Headers,
     private val isClosed: AtomicBoolean,
 ) : Sequence<T> {
 
@@ -86,15 +88,16 @@ private class IOExceptionWrappingSequence<T>(
                 try {
                     iterator.next()
                 } catch (e: IOException) {
-                    throw OpenAIIoException("Stream failed", e)
+                    throw OpenAIIoException("Stream failed", e).withHeaders(headers)
                 }
 
             override fun hasNext(): Boolean =
                 try {
                     iterator.hasNext()
                 } catch (e: IOException) {
                     // Closing the transport may abort an in-flight read instead of returning EOF.
-                    if (isClosed.get()) false else throw OpenAIIoException("Stream failed", e)
+                    if (isClosed.get()) false
+                    else throw OpenAIIoException("Stream failed", e).withHeaders(headers)
                 }
         }
     }

openai-java-core/src/main/kotlin/com/openai/errors/OpenAIInvalidDataException.kt

@@ -1,5 +1,20 @@
 package com.openai.errors
 
+import com.openai.core.http.Headers
+import java.util.Optional
+
 class OpenAIInvalidDataException
 @JvmOverloads
-constructor(message: String? = null, cause: Throwable? = null) : OpenAIException(message, cause)
+constructor(message: String? = null, cause: Throwable? = null) : OpenAIException(message, cause) {
+
+    private var responseHeaders: Headers? = null
+
+    /**
+     * Returns headers retained when reading or parsing the HTTP response failed.
+     *
+     * Empty for failures before a response is available or outside response handling.
+     */
+    fun headers(): Optional<Headers> = Optional.ofNullable(responseHeaders)
+
+    @JvmSynthetic internal fun withHeaders(headers: Headers) = apply { responseHeaders = headers }
+}

openai-java-core/src/main/kotlin/com/openai/errors/OpenAIIoException.kt

@@ -1,5 +1,20 @@
 package com.openai.errors
 
+import com.openai.core.http.Headers
+import java.util.Optional
+
 class OpenAIIoException
 @JvmOverloads
-constructor(message: String? = null, cause: Throwable? = null) : OpenAIException(message, cause)
+constructor(message: String? = null, cause: Throwable? = null) : OpenAIException(message, cause) {
+
+    private var responseHeaders: Headers? = null
+
+    /**
+     * Returns headers retained when reading or parsing the HTTP response failed.
+     *
+     * Empty for failures before a response is available or outside response handling.
+     */
+    fun headers(): Optional<Headers> = Optional.ofNullable(responseHeaders)
+
+    @JvmSynthetic internal fun withHeaders(headers: Headers) = apply { responseHeaders = headers }
+}

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

@@ -0,0 +1,120 @@
+package com.openai.core.handlers
+
+import com.fasterxml.jackson.core.JsonProcessingException
+import com.openai.core.http.Headers
+import com.openai.core.http.HttpResponse
+import com.openai.core.jsonMapper
+import com.openai.errors.OpenAIInvalidDataException
+import com.openai.errors.OpenAIIoException
+import java.io.InputStream
+import java.io.InterruptedIOException
+import org.assertj.core.api.Assertions.assertThat
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.assertThrows
+import org.junit.jupiter.params.ParameterizedTest
+import org.junit.jupiter.params.provider.ValueSource
+import org.mockito.kotlin.any
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.whenever
+
+internal class ResponseHeadersTest {
+    private val headers = Headers.builder().put("x-request-id", "req_test").build()
+
+    @Test
+    fun malformedJsonRetainsHeaders() {
+        val response = response("{".byteInputStream())
+        val error =
+            assertThrows<OpenAIInvalidDataException> {
+                jsonHandler<Map<String, String>>(jsonMapper()).handle(response)
+            }
+        assertThat(error.headers()).contains(headers)
+        assertThat(error.toString()).doesNotContain("req_test")
+        assertThat(error)
+            .hasMessage("Error reading response")
+            .hasCauseInstanceOf(JsonProcessingException::class.java)
+    }
+
+    @Test
+    fun bodyReadTimeoutRetainsHeadersAndCause() {
+        val timeout = InterruptedIOException("timeout")
+        val response = response(failingBody(timeout))
+        val error =
+            assertThrows<OpenAIInvalidDataException> {
+                jsonHandler<Map<String, String>>(jsonMapper()).handle(response)
+            }
+        assertThat(error.headers()).contains(headers)
+        assertThat(error.toString()).doesNotContain("req_test")
+        assertThat(error).hasMessage("Error reading response").hasCause(timeout)
+    }
+
+    @Test
+    fun streamReadTimeoutRetainsHeadersAndCause() {
+        val timeout = InterruptedIOException("timeout")
+        val response = response(failingBody(timeout))
+        val error =
+            assertThrows<OpenAIIoException> {
+                sseHandler(jsonMapper()).handle(response).use { it.stream().count() }
+            }
+        assertThat(error.headers()).contains(headers)
+        assertThat(error.toString()).doesNotContain("req_test")
+        assertThat(error).hasMessage("Stream failed").hasCause(timeout)
+    }
+
+    @ParameterizedTest
+    @ValueSource(strings = ["{", "[]"])
+    fun sseJsonFailuresRetainHeaders(data: String) {
+        val response = response("data: $data\n\n".byteInputStream())
+        val error =
+            assertThrows<OpenAIInvalidDataException> {
+                sseHandler(jsonMapper()).mapJson<Map<String, String>>().handle(response).use {
+                    it.stream().count()
+                }
+            }
+        assertThat(error.headers()).contains(headers)
+        assertThat(error.toString()).doesNotContain("req_test")
+        assertThat(error)
+            .hasMessage("Error reading response")
+            .hasCauseInstanceOf(JsonProcessingException::class.java)
+    }
+
+    @Test
+    fun consumerExceptionsAreNotEnriched() {
+        val expected = OpenAIInvalidDataException("consumer failed")
+        val response = response("data: {}\n\n".byteInputStream())
+        val error =
+            assertThrows<OpenAIInvalidDataException> {
+                sseHandler(jsonMapper()).mapJson<Map<String, String>>().handle(response).use {
+                    it.stream().forEach { throw expected }
+                }
+            }
+        assertThat(error).isSameAs(expected)
+        assertThat(error.headers()).isEmpty()
+    }
+
+    @Test
+    fun exceptionsWithoutResponsesHaveNoHeaders() {
+        val cause = InterruptedIOException("timeout")
+        val ioError = OpenAIIoException("Request failed", cause)
+        val dataError = OpenAIInvalidDataException("Invalid local data", cause)
+        assertThat(ioError.headers()).isEmpty()
+        assertThat(dataError.headers()).isEmpty()
+        assertThat(ioError).hasCause(cause)
+        assertThat(dataError).hasCause(cause)
+        assertThat(OpenAIIoException().headers()).isEmpty()
+        assertThat(OpenAIInvalidDataException().headers()).isEmpty()
+        assertThat(OpenAIIoException("message").headers()).isEmpty()
+        assertThat(OpenAIInvalidDataException("message").headers()).isEmpty()
+    }
+
+    private fun response(body: InputStream): HttpResponse = mock {
+        whenever(it.statusCode()).thenReturn(200)
+        whenever(it.headers()).thenReturn(headers)
+        whenever(it.body()).thenReturn(body)
+    }
+
+    private fun failingBody(failure: InterruptedIOException): InputStream = mock {
+        whenever(it.read(any<ByteArray>())).thenThrow(failure)
+        whenever(it.read(any<ByteArray>(), any(), any())).thenThrow(failure)
+        whenever(it.read()).thenThrow(failure)
+    }
+}

openai-java-core/src/test/kotlin/com/openai/core/http/HttpResponseForTest.kt

@@ -79,6 +79,7 @@ internal class HttpResponseForTest {
         val reading = CountDownLatch(1)
         val releaseRead = CountDownLatch(1)
         val response = mock<HttpResponse>()
+        whenever(response.headers()).thenReturn(Headers.builder().build())
         whenever(response.body())
             .thenReturn(
                 object : InputStream() {
@@ -131,6 +132,7 @@ internal class HttpResponseForTest {
         val reading = CountDownLatch(1)
         val releaseRead = CountDownLatch(1)
         val response = mock<HttpResponse>()
+        whenever(response.headers()).thenReturn(Headers.builder().build())
         whenever(response.body())
             .thenReturn(
                 object : InputStream() {