fix(client): cancel okhttp call when future cancelled

Author: stainless-app[bot]

openai-java-client-okhttp/build.gradle.kts

@@ -11,4 +11,5 @@ dependencies {
 
     testImplementation(kotlin("test"))
     testImplementation("org.assertj:assertj-core:3.25.3")
+    testImplementation("com.github.tomakehurst:wiremock-jre8:2.35.2")
 }

openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/OkHttpClient.kt

@@ -13,6 +13,7 @@ import java.io.IOException
 import java.io.InputStream
 import java.net.Proxy
 import java.time.Duration
+import java.util.concurrent.CancellationException
 import java.util.concurrent.CompletableFuture
 import javax.net.ssl.HostnameVerifier
 import javax.net.ssl.SSLSocketFactory
@@ -29,8 +30,8 @@ import okhttp3.Response
 import okhttp3.logging.HttpLoggingInterceptor
 import okio.BufferedSink
 
-class OkHttpClient private constructor(private val okHttpClient: okhttp3.OkHttpClient) :
-    HttpClient {
+class OkHttpClient
+private constructor(@JvmSynthetic internal val okHttpClient: okhttp3.OkHttpClient) : HttpClient {
 
     override fun execute(request: HttpRequest, requestOptions: RequestOptions): HttpResponse {
         val call = newCall(request, requestOptions)
@@ -50,20 +51,25 @@ class OkHttpClient private constructor(private val okHttpClient: okhttp3.OkHttpC
     ): CompletableFuture<HttpResponse> {
         val future = CompletableFuture<HttpResponse>()
 
-        request.body?.run { future.whenComplete { _, _ -> close() } }
-
-        newCall(request, requestOptions)
-            .enqueue(
-                object : Callback {
-                    override fun onResponse(call: Call, response: Response) {
-                        future.complete(response.toResponse())
-                    }
+        val call = newCall(request, requestOptions)
+        call.enqueue(
+            object : Callback {
+                override fun onResponse(call: Call, response: Response) {
+                    future.complete(response.toResponse())
+                }
 
-                    override fun onFailure(call: Call, e: IOException) {
-                        future.completeExceptionally(OpenAIIoException("Request failed", e))
-                    }
+                override fun onFailure(call: Call, e: IOException) {
+                    future.completeExceptionally(OpenAIIoException("Request failed", e))
                 }
-            )
+            }
+        )
+
+        future.whenComplete { _, e ->
+            if (e is CancellationException) {
+                call.cancel()
+            }
+            request.body?.close()
+        }
 
         return future
     }

openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/OkHttpClientTest.kt

@@ -0,0 +1,44 @@
+package com.openai.client.okhttp
+
+import com.github.tomakehurst.wiremock.client.WireMock.*
+import com.github.tomakehurst.wiremock.junit5.WireMockRuntimeInfo
+import com.github.tomakehurst.wiremock.junit5.WireMockTest
+import com.openai.core.http.HttpMethod
+import com.openai.core.http.HttpRequest
+import org.assertj.core.api.Assertions.assertThat
+import org.junit.jupiter.api.BeforeEach
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.parallel.ResourceLock
+
+@WireMockTest
+@ResourceLock("https://github.com/wiremock/wiremock/issues/169")
+internal class OkHttpClientTest {
+
+    private lateinit var baseUrl: String
+    private lateinit var httpClient: OkHttpClient
+
+    @BeforeEach
+    fun beforeEach(wmRuntimeInfo: WireMockRuntimeInfo) {
+        baseUrl = wmRuntimeInfo.httpBaseUrl
+        httpClient = OkHttpClient.builder().build()
+    }
+
+    @Test
+    fun executeAsync_whenFutureCancelled_cancelsUnderlyingCall() {
+        stubFor(post(urlPathEqualTo("/something")).willReturn(ok()))
+        val responseFuture =
+            httpClient.executeAsync(
+                HttpRequest.builder()
+                    .method(HttpMethod.POST)
+                    .baseUrl(baseUrl)
+                    .addPathSegment("something")
+                    .build()
+            )
+        val call = httpClient.okHttpClient.dispatcher.runningCalls().single()
+
+        responseFuture.cancel(false)
+
+        // Should have cancelled the underlying call
+        assertThat(call.isCanceled()).isTrue()
+    }
+}