spring-ai/chat-ui-htmx/src/test/java/com/ankurm/chatui/OkHttpCloseReproTest.java

package com.ankurm.chatui;

import static org.assertj.core.api.Assertions.assertThat;

import com.ankurm.chatui.stub.OpenAiStubServer;
import okhttp3.MediaType;
import okhttp3.OkHttpClient;
import okhttp3.Request;
import okhttp3.RequestBody;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;

/**
 * No Spring AI, no SDK: plain OkHttp 4.12 against the stub. A reader thread is blocked reading the
 * SSE body; another thread then closes the response, or cancels the call. Which one closes the socket?
 */
class OkHttpCloseReproTest {

    @BeforeAll
    static void stub() {
        StubSupport.ensureStarted();
    }

    int run(boolean cancelInsteadOfClose, Transcript t) throws Exception {
        var http = new OkHttpClient();
        int closed = 0;
        for (int i = 0; i < 10; i++) {
            long[] before = OpenAiStubServer.counters();
            var req = new Request.Builder().url("http://localhost:" + StubSupport.PORT + "/v1/chat/completions")
                    .post(RequestBody.create("{\"model\":\"m\",\"stream\":true,\"messages\":[{\"role\":\"user\","
                            + "\"content\":\"count 60\"}]}", MediaType.get("application/json"))).build();
            var call = http.newCall(req);
            var response = call.execute();
            var source = response.body().source();
            var reader = new Thread(() -> {
                try {
                    while (source.readUtf8Line() != null) {
                        // keep reading
                    }
                } catch (Exception ignored) {
                    // closed under us
                }
            });
            reader.start();
            Thread.sleep(150);
            try {
                if (cancelInsteadOfClose) {
                    call.cancel();
                } else {
                    response.close();
                }
            } catch (RuntimeException e) {
                if (i == 0) {
                    t.line("  response.close() threw: " + e);
                }
            }
            Thread.sleep(500);
            if (OpenAiStubServer.counters()[3] - before[3] == 1) {
                closed++;
            }
            reader.join(2000);
        }
        return closed;
    }

    @Test
    void closeFromAnotherThreadLeaksCancelDoesNot() throws Exception {
        var t = new Transcript("05-okhttp-close-vs-cancel.txt");
        t.line("plain OkHttp 4.12.0, 10 runs each, stub streaming 60 chunks at 20 ms, reader thread blocked in a read");
        int viaClose = run(false, t);
        t.line("response.close() from another thread: upstream connection closed in " + viaClose + " of 10 runs");
        int viaCancel = run(true, t);
        t.line("call.cancel()     from another thread: upstream connection closed in " + viaCancel + " of 10 runs");
        assertThat(viaCancel).isEqualTo(10);
        assertThat(viaClose).isLessThanOrEqualTo(2);
    }
}