Cancelling the completable future of an OpenAI call doesn't cancel the underlying OkHttp request

[SoftMemes]

I have a few scenarios where we may opportunistically issue a request to an LLM but then cancel it, sometimes quickly so potentially before it's even been processed by the backend and charged for. It would seem natural that using the cancel method of an async call would do that, but it doesn't appear that this propagates to the underlying OkHttp layer.

This method appears to be the bridge between the SDK interfaces and the underlying HTTP engine:

    override fun executeAsync(
        request: HttpRequest,
        requestOptions: RequestOptions,
    ): 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())
                    }

                    override fun onFailure(call: Call, e: IOException) {
                        future.completeExceptionally(OpenAIIoException("Request failed", e))
                    }
                }
            )

        return future
    }

The problem appears to be that the completable future here doesn't link to the HttpRequest. The Call interface from OkHttp has a cancel method, but this is not called as a result of the future being cancelled.

Instead what appears to happen is that the request is orphaned if the outer future is cancelled, leading to a resource leak.

[TomerAberbach]

Hmmm, good catch! I think we need something like this:

    override fun executeAsync(
        request: HttpRequest,
        requestOptions: RequestOptions,
    ): CompletableFuture<HttpResponse> {
        val future = CompletableFuture<HttpResponse>()

        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))
                }
            }
        )

        request.body?.run {
            future.whenComplete { _, e ->
                if (e is CancellationException) {
                    call.cancel()
                }
                close()
            }
        }

        return future
    }

I'll go write a test for this and confirm it works

[SoftMemes]

Actually, digging into this a bit further, it seems that I had not appreciated how the Java CompletableFuture works. When I cancel the one returned from the SDK at the top level, this (by design) only cancels that future and nothing below it. It just stops listening to the result, but Java does not by default propagate this down to the futures constructed via the methods used to compose futures.

It looks like there are two options here to properly support cancellation:

• Return a custom implementation of CompletableFuture that knows how to actually cancel the real underlying operation, using some custom signalling like a cancellation token
• Add a new argument to these methods that accepts a cancellation token passed in from the outside

Both would rely on a cancellation token being passed all the way in to the http client ultimately, where it would have to be listened to to cleanly abort the operation.

[SoftMemes]

As a sketch, I have opened a PR that demonstrates how this could be done end-to-end, see #657. This only covers a couple of cases but seem to work as a pattern.