feat(core): expose request IDs on streaming responses

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

@@ -8,6 +8,7 @@ import com.openai.core.http.PhantomReachableClosingStreamResponse
 import com.openai.core.http.StreamResponse
 import com.openai.errors.OpenAIIoException
 import java.io.IOException
+import java.util.Optional
 import java.util.stream.Stream
 import kotlin.streams.asStream
 
@@ -18,6 +19,7 @@ internal fun <T> streamHandler(
     object : Handler<StreamResponse<T>> {
 
         override fun handle(response: HttpResponse): StreamResponse<T> {
+            val requestId = response.requestId()
             val reader = response.body().bufferedReader()
             val sequence =
                 // Wrap in a `CloseableSequence` to avoid performing a read on the `reader`
@@ -40,6 +42,8 @@ internal fun <T> streamHandler(
             return PhantomReachableClosingStreamResponse(
                 object : StreamResponse<T> {
 
+                    override fun requestId(): Optional<String> = requestId
+
                     override fun stream(): Stream<T> = sequence.asStream()
 
                     override fun close() {

openai-java-core/src/main/kotlin/com/openai/core/http/AsyncStreamResponse.kt

@@ -1,8 +1,11 @@
 package com.openai.core.http
 
 import com.openai.core.http.AsyncStreamResponse.Handler
+import com.openai.errors.OpenAIServiceException
 import java.util.Optional
+import java.util.concurrent.CancellationException
 import java.util.concurrent.CompletableFuture
+import java.util.concurrent.CompletionException
 import java.util.concurrent.Executor
 import java.util.concurrent.atomic.AtomicReference
 
@@ -13,6 +16,14 @@ import java.util.concurrent.atomic.AtomicReference
  */
 interface AsyncStreamResponse<T> {
 
+    /**
+     * Returns the value of the `x-request-id` header, or an empty [Optional] if it's unavailable.
+     *
+     * This method does not wait for response headers. Its result is empty until those headers are
+     * available, and may remain empty if the response has no `x-request-id` header.
+     */
+    fun requestId(): Optional<String> = Optional.empty()
+
     /**
      * Registers [handler] to be called for events of this stream.
      *
@@ -78,6 +89,19 @@ internal fun <T> CompletableFuture<StreamResponse<T>>.toAsync(streamHandlerExecu
                 }
             }
 
+            override fun requestId(): Optional<String> {
+                val streamResponse =
+                    try {
+                        this@toAsync.getNow(null)
+                    } catch (error: CompletionException) {
+                        return requestIdFromError(error)
+                    } catch (_: CancellationException) {
+                        return Optional.empty()
+                    }
+
+                return streamResponse?.requestId() ?: Optional.empty()
+            }
+
             override fun subscribe(handler: Handler<T>): AsyncStreamResponse<T> =
                 subscribe(handler, streamHandlerExecutor)
 
@@ -150,6 +174,14 @@ internal fun <T> CompletableFuture<StreamResponse<T>>.toAsync(streamHandlerExecu
         }
     )
 
+private tailrec fun requestIdFromError(error: Throwable): Optional<String> =
+    when {
+        error is CompletionException && error.cause != null -> requestIdFromError(error.cause!!)
+        error is OpenAIServiceException ->
+            Optional.ofNullable(error.headers().values("x-request-id").firstOrNull())
+        else -> Optional.empty()
+    }
+
 private enum class State {
     NEW,
     SUBSCRIBED,

openai-java-core/src/main/kotlin/com/openai/core/http/PhantomReachableClosingAsyncStreamResponse.kt

@@ -25,6 +25,8 @@ internal class PhantomReachableClosingAsyncStreamResponse<T>(
         closeWhenPhantomReachable(reachabilityTracker, asyncStreamResponse::close)
     }
 
+    override fun requestId(): Optional<String> = asyncStreamResponse.requestId()
+
     override fun subscribe(handler: Handler<T>): AsyncStreamResponse<T> = apply {
         asyncStreamResponse.subscribe(TrackedHandler(handler, reachabilityTracker))
     }

openai-java-core/src/main/kotlin/com/openai/core/http/PhantomReachableClosingStreamResponse.kt

@@ -1,6 +1,7 @@
 package com.openai.core.http
 
 import com.openai.core.closeWhenPhantomReachable
+import java.util.Optional
 import java.util.stream.Stream
 
 /**
@@ -15,6 +16,8 @@ internal class PhantomReachableClosingStreamResponse<T>(
         closeWhenPhantomReachable(this, streamResponse)
     }
 
+    override fun requestId(): Optional<String> = streamResponse.requestId()
+
     override fun stream(): Stream<T> = streamResponse.stream()
 
     override fun close() = streamResponse.close()

openai-java-core/src/main/kotlin/com/openai/core/http/StreamResponse.kt

@@ -1,9 +1,15 @@
 package com.openai.core.http
 
+import java.util.Optional
 import java.util.stream.Stream
 
 interface StreamResponse<T> : AutoCloseable {
 
+    /**
+     * Returns the value of the `x-request-id` header, or an empty [Optional] if it's unavailable.
+     */
+    fun requestId(): Optional<String> = Optional.empty()
+
     fun stream(): Stream<T>
 
     /** Overridden from [AutoCloseable] to not have a checked exception in its signature. */
@@ -13,6 +19,8 @@ interface StreamResponse<T> : AutoCloseable {
 @JvmSynthetic
 internal fun <T, R> StreamResponse<T>.map(transform: (T) -> R): StreamResponse<R> =
     object : StreamResponse<R> {
+        override fun requestId(): Optional<String> = this@map.requestId()
+
         override fun stream(): Stream<R> = this@map.stream().map(transform)
 
         override fun close() = this@map.close()

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

@@ -22,6 +22,22 @@ internal class StreamHandlerTest {
         assertThat(lines).containsExactly("a", "bb", "ccc", "dddd")
     }
 
+    @Test
+    fun streamHandler_exposesRequestIdAfterClose() {
+        val handler = streamHandler { _, lines -> yieldAll(lines) }
+        val streamResponse =
+            handler.handle(
+                httpResponse(
+                    "a\n".byteInputStream(),
+                    Headers.builder().put("x-request-id", "req_123").build(),
+                )
+            )
+
+        assertThat(streamResponse.requestId()).contains("req_123")
+        streamResponse.close()
+        assertThat(streamResponse.requestId()).contains("req_123")
+    }
+
     @Test
     fun streamHandler_whenClosedEarly_stopsYielding() {
         val handler = streamHandler { _, lines -> yieldAll(lines) }
@@ -68,12 +84,15 @@ internal class StreamHandlerTest {
         assertThat(e).isSameAs(ioException)
     }
 
-    private fun httpResponse(body: InputStream): HttpResponse =
+    private fun httpResponse(
+        body: InputStream,
+        headers: Headers = Headers.builder().build(),
+    ): HttpResponse =
         object : HttpResponse {
 
             override fun statusCode(): Int = 0
 
-            override fun headers(): Headers = Headers.builder().build()
+            override fun headers(): Headers = headers
 
             override fun body(): InputStream = body
 

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

@@ -1,7 +1,9 @@
 package com.openai.core.http
 
+import com.openai.errors.BadRequestException
 import java.util.*
 import java.util.concurrent.CompletableFuture
+import java.util.concurrent.CompletionException
 import java.util.concurrent.Executor
 import java.util.stream.Stream
 import kotlin.streams.asStream
@@ -45,6 +47,78 @@ internal class AsyncStreamResponseTest {
         }
     private val handler = mock<AsyncStreamResponse.Handler<String>>()
 
+    @Test
+    fun requestId_whenStreamFutureIsIncomplete_isEmpty() {
+        val asyncStreamResponse = CompletableFuture<StreamResponse<String>>().toAsync(executor)
+
+        assertThat(asyncStreamResponse.requestId()).isEmpty
+    }
+
+    @Test
+    fun requestId_whenStreamFutureCompletes_isAvailable() {
+        val future = CompletableFuture<StreamResponse<String>>()
+        val asyncStreamResponse = future.toAsync(executor)
+        doReturn(Optional.of("req_123")).whenever(streamResponse).requestId()
+
+        future.complete(streamResponse)
+
+        assertThat(asyncStreamResponse.requestId()).contains("req_123")
+    }
+
+    @Test
+    fun requestId_whenHandlerRuns_isAvailable() {
+        val future = CompletableFuture<StreamResponse<String>>()
+        val asyncStreamResponse = future.toAsync(executor)
+        val requestIds = mutableListOf<Optional<String>>()
+        doReturn(Optional.of("req_123")).whenever(streamResponse).requestId()
+        asyncStreamResponse.subscribe { requestIds.add(asyncStreamResponse.requestId()) }
+
+        future.complete(streamResponse)
+
+        assertThat(requestIds).containsOnly(Optional.of("req_123"))
+    }
+
+    @Test
+    fun requestId_whenStreamFutureFailsWithServiceException_isAvailableInExceptionally() {
+        val future = CompletableFuture<StreamResponse<String>>()
+        val asyncStreamResponse = future.toAsync(executor)
+        val requestIds = mutableListOf<Optional<String>>()
+        val error =
+            BadRequestException.builder()
+                .headers(Headers.builder().put("x-request-id", "req_error").build())
+                .build()
+        val assertion =
+            asyncStreamResponse.onCompleteFuture().exceptionally {
+                requestIds.add(asyncStreamResponse.requestId())
+                null
+            }
+
+        future.completeExceptionally(CompletionException(error))
+
+        assertion.join()
+        assertThat(requestIds).containsExactly(Optional.of("req_error"))
+    }
+
+    @Test
+    fun requestId_whenStreamFutureFailsBeforeResponse_isEmpty() {
+        val future = CompletableFuture<StreamResponse<String>>()
+        val asyncStreamResponse = future.toAsync(executor)
+
+        future.completeExceptionally(ERROR)
+
+        assertThat(asyncStreamResponse.requestId()).isEmpty
+    }
+
+    @Test
+    fun requestId_whenStreamFutureIsCancelled_isEmpty() {
+        val future = CompletableFuture<StreamResponse<String>>()
+        val asyncStreamResponse = future.toAsync(executor)
+
+        future.cancel(false)
+
+        assertThat(asyncStreamResponse.requestId()).isEmpty
+    }
+
     @Test
     fun subscribe_whenAlreadySubscribed_throws() {
         val asyncStreamResponse = CompletableFuture<StreamResponse<String>>().toAsync(executor)

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

@@ -0,0 +1,36 @@
+package com.openai.core.http
+
+import java.util.Optional
+import java.util.stream.Stream
+import org.assertj.core.api.Assertions.assertThat
+import org.junit.jupiter.api.Test
+
+internal class StreamResponseTest {
+
+    @Test
+    fun requestId_byDefault_isEmpty() {
+        val response = streamResponse()
+
+        assertThat(response.requestId()).isEmpty
+    }
+
+    @Test
+    fun map_preservesRequestId() {
+        val response = streamResponse(Optional.of("req_123"))
+
+        val mappedResponse = response.map(String::length).map(Int::toString)
+
+        assertThat(mappedResponse.requestId()).contains("req_123")
+    }
+
+    private fun streamResponse(
+        requestId: Optional<String> = Optional.empty()
+    ): StreamResponse<String> =
+        object : StreamResponse<String> {
+            override fun requestId(): Optional<String> = requestId
+
+            override fun stream(): Stream<String> = Stream.empty()
+
+            override fun close() {}
+        }
+}