diff --git a/http-client/http-client-java/src/main/kotlin/ai/koog/http/client/java/JavaKoogHttpClient.kt b/http-client/http-client-java/src/main/kotlin/ai/koog/http/client/java/JavaKoogHttpClient.kt index ff5c5fa694..4ea1b6a148 100644 --- a/http-client/http-client-java/src/main/kotlin/ai/koog/http/client/java/JavaKoogHttpClient.kt +++ b/http-client/http-client-java/src/main/kotlin/ai/koog/http/client/java/JavaKoogHttpClient.kt @@ -8,10 +8,15 @@ import io.github.oshai.kotlinlogging.KLogger import io.github.oshai.kotlinlogging.KotlinLogging import kotlinx.coroutines.CancellationException import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.cancel +import kotlinx.coroutines.channels.ProducerScope import kotlinx.coroutines.channels.awaitClose +import kotlinx.coroutines.ensureActive import kotlinx.coroutines.flow.Flow import kotlinx.coroutines.flow.callbackFlow +import kotlinx.coroutines.isActive import kotlinx.coroutines.launch +import kotlinx.coroutines.runInterruptible import kotlinx.coroutines.withContext import kotlinx.serialization.json.Json import kotlinx.serialization.serializer @@ -21,6 +26,7 @@ import java.net.http.HttpClient import java.net.http.HttpRequest import java.net.http.HttpResponse import java.time.Duration +import java.util.concurrent.atomic.AtomicReference import kotlin.reflect.KClass import kotlin.time.Duration.Companion.minutes @@ -31,6 +37,10 @@ import kotlin.time.Duration.Companion.minutes * This client provides enhanced logging, flexible request and response handling, and supports * configurability for underlying Java HttpClient instances. * + * SSE dispatches complete UTF-8 data events, preserving multiline payloads and ignoring event metadata. + * Streaming delivery suspends when collectors are slow. Cancellation interrupts pending requests and + * closes response bodies, including idle streams. + * * @property clientName The name of the client, used for logging and traceability. * @property logger A logging instance of type KLogger for recording client-related events and errors. * @property httpClient The configured Java HttpClient instance used for making HTTP requests. @@ -188,69 +198,34 @@ public class JavaKoogHttpClient internal constructor( // Note: "Connection" header is restricted in Java HttpClient and managed automatically .build() - try { - val response = httpClient.send(httpRequest, HttpResponse.BodyHandlers.ofLines()) - - if (response.statusCode() !in 200..299) { - close( - KoogHttpClientException( - clientName = clientName, - statusCode = response.statusCode(), - ) - ) - return@callbackFlow - } - + readStreamingResponse(httpRequest, HttpResponse.BodyHandlers.ofInputStream()) { body -> + val reader = body.bufferedReader(Charsets.UTF_8) logger.debug { "SSE connection opened for $clientName" } - - // Process the stream of lines - response.body().forEach { line -> - try { - val dataPrefix = "data: " - // SSE format: "data: " - val data = if (line.startsWith(dataPrefix)) { - line.substring(dataPrefix.length) - } else if (line.isNotEmpty() && !line.startsWith(":")) { - line - } else { - null + val data = StringBuilder() + var firstLine = true + while (isActive) { + val rawLine = reader.readLine() ?: break + val line = if (firstLine) rawLine.removePrefix("\uFEFF") else rawLine + firstLine = false + if (line.isEmpty()) { + if (data.isNotEmpty()) { + val payload = data.dropLast(1).toString() + data.setLength(0) + if (dataFilter(payload)) { + processStreamingChunk(decodeStreamingResponse(payload))?.let { send(it) } + } } - - if (data != null && dataFilter(data)) { - data.trim() - .let(decodeStreamingResponse) - .let(processStreamingChunk) - ?.let { trySend(it) } + } else if (!line.startsWith(":")) { + val colon = line.indexOf(':') + val field = if (colon < 0) line else line.substring(0, colon) + if (field == "data") { + val value = if (colon < 0) "" else line.substring(colon + 1).removePrefix(" ") + data.append(value).append('\n') } - } catch (e: CancellationException) { - throw e - } catch (e: Exception) { - close( - KoogHttpClientException( - clientName = clientName, - message = "Error processing SSE event: ${e.message}", - cause = e - ) - ) } } - + // Pending data without a terminating empty line belongs to an incomplete event. logger.debug { "SSE connection closed for $clientName" } - close() - } catch (e: CancellationException) { - throw e - } catch (e: Exception) { - close( - KoogHttpClientException( - clientName = clientName, - message = "Exception during streaming: ${e.message}", - cause = e - ) - ) - } - - awaitClose { - // Cleanup if needed } } @@ -276,51 +251,64 @@ public class JavaKoogHttpClient internal constructor( .POST(HttpRequest.BodyPublishers.ofString(preparedRequestBody.body)) .build() - val responseFuture = httpClient.sendAsync(httpRequest, HttpResponse.BodyHandlers.ofLines()) + readStreamingResponse(httpRequest, HttpResponse.BodyHandlers.ofLines()) { lines -> + logger.debug { "Lines flow opened for $clientName" } + val iterator = lines.iterator() + while (isActive && iterator.hasNext()) { + val line = iterator.next() + if (line.isBlank()) continue + send(line) + } + logger.debug { "Lines flow closed for $clientName" } + } + } + + private suspend fun ProducerScope.readStreamingResponse( + request: HttpRequest, + bodyHandler: HttpResponse.BodyHandler, + process: suspend ProducerScope.(B) -> Unit, + ) { + val body = AtomicReference() val readerJob = launch(Dispatchers.SuitableForIO) { try { - val response = responseFuture.get() + // Keep ownership if cancellation races with receiving the response headers. + val response = runInterruptible { + httpClient.send(request, bodyHandler).also { + body.set(it.body()) + } + } if (response.statusCode() !in 200..299) { - close( - KoogHttpClientException( - clientName = clientName, - statusCode = response.statusCode(), - ) + throw KoogHttpClientException( + clientName = clientName, + statusCode = response.statusCode(), ) - return@launch - } - - logger.debug { "Lines flow opened for $clientName" } - response.body().use { lines -> - val iterator = lines.iterator() - while (iterator.hasNext()) { - val line = iterator.next() - if (line.isBlank()) continue - if (trySend(line).isClosed) { - responseFuture.cancel(true) - return@launch - } - } } - - logger.debug { "Lines flow closed for $clientName" } + process(response.body()) close() } catch (e: CancellationException) { + // Callback cancellation must also terminate the producer waiting in awaitClose. + this@readStreamingResponse.cancel(e) throw e + } catch (e: KoogHttpClientException) { + coroutineContext.ensureActive() + close(e) } catch (e: Exception) { + coroutineContext.ensureActive() close( KoogHttpClientException( clientName = clientName, message = "Exception during streaming: ${e.message}", - cause = e + cause = e, ) ) + } finally { + body.getAndSet(null)?.close() } } - awaitClose { - responseFuture.cancel(true) readerJob.cancel() + // JDK body reads can ignore interruption, so closing the body releases an idle reader. + body.getAndSet(null)?.close() } } diff --git a/http-client/http-client-java/src/test/kotlin/ai/koog/http/client/java/JavaKoogHttpClientBodyClosureTest.kt b/http-client/http-client-java/src/test/kotlin/ai/koog/http/client/java/JavaKoogHttpClientBodyClosureTest.kt new file mode 100644 index 0000000000..ae7b109679 --- /dev/null +++ b/http-client/http-client-java/src/test/kotlin/ai/koog/http/client/java/JavaKoogHttpClientBodyClosureTest.kt @@ -0,0 +1,189 @@ +package ai.koog.http.client.java + +import ai.koog.http.client.KoogHttpClient +import ai.koog.http.client.KoogHttpClientException +import io.github.oshai.kotlinlogging.KotlinLogging +import kotlinx.coroutines.CompletableDeferred +import kotlinx.coroutines.CoroutineStart +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.Job +import kotlinx.coroutines.flow.flowOn +import kotlinx.coroutines.flow.toList +import kotlinx.coroutines.launch +import kotlinx.coroutines.test.runTest +import kotlinx.coroutines.withContext +import kotlinx.coroutines.withTimeout +import java.io.FilterInputStream +import java.io.InputStream +import java.net.Authenticator +import java.net.CookieHandler +import java.net.ProxySelector +import java.net.URI +import java.net.http.HttpClient +import java.net.http.HttpHeaders +import java.net.http.HttpRequest +import java.net.http.HttpResponse +import java.nio.ByteBuffer +import java.time.Duration +import java.util.Optional +import java.util.concurrent.CompletableFuture +import java.util.concurrent.Executor +import java.util.concurrent.Flow +import java.util.concurrent.atomic.AtomicInteger +import java.util.stream.Stream +import javax.net.ssl.SSLContext +import javax.net.ssl.SSLParameters +import javax.net.ssl.SSLSession +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith + +class JavaKoogHttpClientBodyClosureTest { + @Test + fun testSseClosesBodyAfterEof() = runTest { assertClosure(sse = true, status = 200) } + + @Test + fun testLinesCloseBodyAfterEof() = runTest { assertClosure(sse = false, status = 200) } + + @Test + fun testSseClosesBodyAfterHttpFailure() = runTest { assertClosure(sse = true, status = 429) } + + @Test + fun testLinesCloseBodyAfterHttpFailure() = runTest { assertClosure(sse = false, status = 429) } + + @Test + fun testSseClosesBodyAfterCallbackFailure() = runTest { + val transport = BodyClient(200, "data: bad\n\n") + val client = KoogHttpClient.fromJavaHttpClient("body-closure", KotlinLogging.logger {}, transport) + assertFailsWith { + client.sse( + path = "http://localhost/stream", + requestBody = "{}", + requestBodyType = String::class, + decodeStreamingResponse = { throw IllegalArgumentException("bad") }, + processStreamingChunk = { value: String -> value }, + ).flowOn(Dispatchers.IO).toList() + } + assertEquals(1, transport.closures.get()) + } + + @Test + fun testCancellationWhenHeadersArriveClosesReceivedBody() = runTest { + withContext(Dispatchers.Default) { + lateinit var collection: Job + val responseReady = CompletableDeferred() + val transport = BodyClient(200, "data: payload\n\n") { + responseReady.complete(Unit) + collection.cancel() + } + val client = KoogHttpClient.fromJavaHttpClient("body-race", KotlinLogging.logger {}, transport) + collection = launch(Dispatchers.IO, start = CoroutineStart.LAZY) { + client.sse( + path = "http://localhost/stream", + requestBody = "{}", + requestBodyType = String::class, + decodeStreamingResponse = { it }, + processStreamingChunk = { it }, + ).toList() + } + collection.start() + withTimeout(5_000) { + responseReady.await() + collection.join() + } + assertEquals(1, transport.closures.get()) + } + } + + private suspend fun assertClosure(sse: Boolean, status: Int) { + val transport = BodyClient(status, if (sse) "data: payload\n\n" else "payload\n") + val client = KoogHttpClient.fromJavaHttpClient("body-closure", KotlinLogging.logger {}, transport) + val stream = if (sse) { + client.sse( + path = "http://localhost/stream", + requestBody = "{}", + requestBodyType = String::class, + decodeStreamingResponse = { it }, + processStreamingChunk = { it }, + ) + } else { + client.lines(path = "http://localhost/stream", requestBody = "{}", requestBodyType = String::class) + } + if (status == 200) { + assertEquals(listOf("payload"), stream.flowOn(Dispatchers.IO).toList()) + } else { + val error = assertFailsWith { stream.flowOn(Dispatchers.IO).toList() } + assertEquals(status, error.statusCode) + } + assertEquals(1, transport.closures.get(), "The owned body must be closed exactly once") + } + + /** Supplies bytes through the actual JDK body handler, observing response-body closure. */ + private class BodyClient( + private val status: Int, + private val content: String, + private val onResponseReady: () -> Unit = {}, + ) : HttpClient() { + val closures = AtomicInteger() + private val responseHeaders = HttpHeaders.of(emptyMap()) { _, _ -> true } + + override fun send(request: HttpRequest, handler: HttpResponse.BodyHandler): HttpResponse { + val subscriber = handler.apply(object : HttpResponse.ResponseInfo { + override fun statusCode(): Int = status + override fun headers(): HttpHeaders = responseHeaders + override fun version(): Version = Version.HTTP_1_1 + }) + subscriber.onSubscribe(object : Flow.Subscription { + override fun request(n: Long) {} + override fun cancel() {} + }) + subscriber.onNext(listOf(ByteBuffer.wrap(content.toByteArray(Charsets.UTF_8)))) + subscriber.onComplete() + val original = subscriber.body.toCompletableFuture().get() + val observed = when (original) { + is InputStream -> object : FilterInputStream(original) { + override fun close() { + closures.incrementAndGet() + super.close() + } + } + is Stream<*> -> original.onClose { closures.incrementAndGet() } + else -> error("Expected a streaming response body") + } + + // These tests request only InputStream or Stream; wrappers preserve that body type. + @Suppress("UNCHECKED_CAST") + val body = observed as T + onResponseReady() + return object : HttpResponse { + override fun statusCode(): Int = status + override fun request(): HttpRequest = request + override fun previousResponse(): Optional> = Optional.empty() + override fun headers(): HttpHeaders = responseHeaders + override fun body(): T = body + override fun sslSession(): Optional = Optional.empty() + override fun uri(): URI = request.uri() + override fun version(): Version = Version.HTTP_1_1 + } + } + + override fun sendAsync(request: HttpRequest, handler: HttpResponse.BodyHandler): CompletableFuture> = + CompletableFuture.completedFuture(send(request, handler)) + + override fun sendAsync( + request: HttpRequest, + handler: HttpResponse.BodyHandler, + pushPromiseHandler: HttpResponse.PushPromiseHandler, + ): CompletableFuture> = sendAsync(request, handler) + + override fun cookieHandler(): Optional = Optional.empty() + override fun connectTimeout(): Optional = Optional.empty() + override fun followRedirects(): Redirect = Redirect.NEVER + override fun proxy(): Optional = Optional.empty() + override fun sslContext(): SSLContext = SSLContext.getDefault() + override fun sslParameters(): SSLParameters = SSLParameters() + override fun authenticator(): Optional = Optional.empty() + override fun version(): Version = Version.HTTP_1_1 + override fun executor(): Optional = Optional.empty() + } +} diff --git a/http-client/http-client-java/src/test/kotlin/ai/koog/http/client/java/JavaKoogHttpClientSseFramingTest.kt b/http-client/http-client-java/src/test/kotlin/ai/koog/http/client/java/JavaKoogHttpClientSseFramingTest.kt new file mode 100644 index 0000000000..8c10640f0a --- /dev/null +++ b/http-client/http-client-java/src/test/kotlin/ai/koog/http/client/java/JavaKoogHttpClientSseFramingTest.kt @@ -0,0 +1,97 @@ +package ai.koog.http.client.java + +import ai.koog.http.client.KoogHttpClient +import ai.koog.http.client.test.MockWebServer +import io.ktor.http.ContentType +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.flow.flowOn +import kotlinx.coroutines.flow.toList +import kotlinx.coroutines.test.runTest +import kotlin.test.Test +import kotlin.test.assertEquals + +class JavaKoogHttpClientSseFramingTest { + @Test + fun testMultilineDataDispatchesAsOneEvent() = runTest { + assertEquals(listOf("first\nsecond"), collect("data: first\ndata:second\n\n")) + } + + @Test + fun testMetadataAndCommentsAreIgnored() = runTest { + assertEquals(listOf("payload"), collect(": heartbeat\nevent: message\nid: 7\nretry: 100\nfuture: ignored\nraw text\ndata: payload\n\n")) + } + + @Test + fun testLeadingBomAndAllLineEndings() = runTest { + assertEquals(listOf("one", "two", "three"), collect("\uFEFFdata: one\r\rdata: two\r\n\r\ndata: three\n\n")) + } + + @Test + fun testDataWhitespaceIsPreserved() = runTest { + assertEquals(listOf(" leading \n\ttab "), collect("data: leading \ndata:\ttab \n\n")) + } + + @Test + fun testEmptyDataAndColonlessDataDispatch() = runTest { + assertEquals(listOf("", ""), collect("\nevent: ignored\n\ndata:\n\ndata\n\n")) + } + + @Test + fun testIncompleteEventIsDiscardedAtEof() = runTest { + assertEquals(listOf("complete"), collect("data: complete\n\ndata: incomplete\n")) + } + + @Test + fun testFilterAndDecoderReceiveCompleteDataEvents() = runTest { + val filtered = mutableListOf() + withResponse("data: first\ndata: second\n\ndata: skip\n\ndata: keep\n\n") { client -> + val result = client.sse( + path = "/stream", + requestBody = "{}", + requestBodyType = String::class, + dataFilter = { + filtered.add(it) + it != "skip" + }, + decodeStreamingResponse = { it.uppercase() }, + processStreamingChunk = { it.takeUnless { it == "FIRST\nSECOND" } }, + ).flowOn(Dispatchers.IO).toList() + assertEquals(listOf("first\nsecond", "skip", "keep"), filtered) + assertEquals(listOf("KEEP"), result) + } + } + + private suspend fun collect(response: String): List { + var result = emptyList() + withResponse(response) { client -> + result = client.sse( + path = "/stream", + requestBody = "{}", + requestBodyType = String::class, + decodeStreamingResponse = { it }, + processStreamingChunk = { it }, + ).flowOn(Dispatchers.IO).toList() + } + return result + } + + private suspend fun withResponse(response: String, block: suspend (KoogHttpClient) -> Unit) { + val server = MockWebServer() + server.start( + postEndpoints = listOf( + MockWebServer.PostEndpointConfig( + path = "/stream", + responseBody = response, + contentType = ContentType.Text.EventStream, + ) + ) + ) + val client = JavaKoogHttpClient.Factory().create(clientName = "java-sse", baseUrl = server.url("")) + try { + block(client) + } finally { + client.close() + server.stop() + } + } +} diff --git a/http-client/http-client-java/src/test/kotlin/ai/koog/http/client/java/JavaKoogHttpClientStreamingTest.kt b/http-client/http-client-java/src/test/kotlin/ai/koog/http/client/java/JavaKoogHttpClientStreamingTest.kt new file mode 100644 index 0000000000..5af023197c --- /dev/null +++ b/http-client/http-client-java/src/test/kotlin/ai/koog/http/client/java/JavaKoogHttpClientStreamingTest.kt @@ -0,0 +1,320 @@ +package ai.koog.http.client.java + +import ai.koog.http.client.KoogHttpClient +import ai.koog.http.client.KoogHttpClientException +import com.sun.net.httpserver.HttpExchange +import com.sun.net.httpserver.HttpServer +import kotlinx.coroutines.CancellationException +import kotlinx.coroutines.CompletableDeferred +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.async +import kotlinx.coroutines.cancelAndJoin +import kotlinx.coroutines.flow.Flow +import kotlinx.coroutines.flow.buffer +import kotlinx.coroutines.flow.first +import kotlinx.coroutines.flow.flowOn +import kotlinx.coroutines.flow.toList +import kotlinx.coroutines.launch +import kotlinx.coroutines.test.runTest +import kotlinx.coroutines.withContext +import kotlinx.coroutines.withTimeout +import java.net.InetSocketAddress +import java.util.concurrent.CountDownLatch +import java.util.concurrent.Executors +import java.util.concurrent.TimeUnit +import java.util.concurrent.atomic.AtomicInteger +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertSame +import kotlin.test.assertTrue + +class JavaKoogHttpClientStreamingTest { + @Test + fun testSsePreservesAllEventsWithRendezvousBackpressure() = runTest { + assertDelivery(sse = true) + } + + @Test + fun testLinesPreserveAllChunksWithRendezvousBackpressure() = runTest { + assertDelivery(sse = false) + } + + @Test + fun testLinesRespectTheDeclaredResponseCharset() = runTest { + withServer({ + responseHeaders.add("Content-Type", "text/plain; charset=ISO-8859-1") + val bytes = "café\n".toByteArray(Charsets.ISO_8859_1) + sendResponseHeaders(200, bytes.size.toLong()) + responseBody.write(bytes) + }) { client -> + assertEquals(listOf("café"), client.stream(sse = false).flowOn(Dispatchers.IO).toList()) + } + } + + @Test + fun testSseBackpressureBoundsDecodedEvents() = runTest { + withContext(Dispatchers.Default) { + val expected = (0..199).map(Int::toString) + withServer({ respond(expected.joinToString("") { "data: $it\n\n" }) }) { client -> + val decoded = AtomicInteger() + val firstCollected = CompletableDeferred() + val secondDecoded = CompletableDeferred() + val releaseCollector = CompletableDeferred() + val received = mutableListOf() + val collection = launch(Dispatchers.IO) { + client.sse( + path = "/stream", + requestBody = "{}", + requestBodyType = String::class, + decodeStreamingResponse = { + if (decoded.incrementAndGet() == 2) secondDecoded.complete(Unit) + it + }, + processStreamingChunk = { it }, + ).buffer(0).collect { + received.add(it) + if (received.size == 1) { + firstCollected.complete(Unit) + releaseCollector.await() + } + } + } + try { + withTimeout(5_000) { + firstCollected.await() + secondDecoded.await() + } + assertEquals(2, decoded.get()) + releaseCollector.complete(Unit) + withTimeout(5_000) { collection.join() } + assertEquals(expected, received) + } finally { + releaseCollector.complete(Unit) + collection.cancelAndJoin() + } + } + } + } + + @Test + fun testSseCancelsWhileWaitingForHeaders() = runTest { + assertCancellation(sse = true, headers = true) + } + + @Test + fun testLinesCancelWhileWaitingForHeaders() = runTest { + assertCancellation(sse = false, headers = true) + } + + @Test + fun testSseCancelsWhileWaitingForBody() = runTest { + assertCancellation(sse = true, headers = false) + } + + @Test + fun testLinesCancelWhileWaitingForBody() = runTest { + assertCancellation(sse = false, headers = false) + } + + @Test + fun testSseFirstEventClosesIdleBody() = runTest { + assertFirst(sse = true) + } + + @Test + fun testLinesFirstChunkClosesIdleBody() = runTest { + assertFirst(sse = false) + } + + @Test + fun testSseDownstreamFailureClosesIdleBody() = runTest { + assertDownstreamFailure(sse = true) + } + + @Test + fun testLinesDownstreamFailureClosesIdleBody() = runTest { + assertDownstreamFailure(sse = false) + } + + @Test + fun testSseCallbackCancellationRetainsItsIdentity() = runTest { + withContext(Dispatchers.Default) { + for (stage in listOf("filter", "decoder", "processor")) { + withServer({ respond("data: first\n\n") }) { client -> + val cancellation = CallbackCancellation() + val error = assertFailsWith { + withTimeout(5_000) { + client.sse( + path = "/stream", + requestBody = "{}", + requestBodyType = String::class, + dataFilter = { + if (stage == "filter") throw cancellation + true + }, + decodeStreamingResponse = { + if (stage == "decoder") throw cancellation + it + }, + processStreamingChunk = { + if (stage == "processor") throw cancellation + it + }, + ).flowOn(Dispatchers.IO).toList() + } + } + assertSame(cancellation, error) + } + } + } + } + + @Test + fun testSseDecoderFailureStopsProcessingLaterEvents() = runTest { + withServer({ respond("data: bad\n\ndata: later\n\n") }) { client -> + var calls = 0 + val cause = IllegalArgumentException("invalid event") + val error = assertFailsWith { + client.sse( + path = "/stream", + requestBody = "{}", + requestBodyType = String::class, + decodeStreamingResponse = { + calls++ + throw cause + }, + processStreamingChunk = { value: String -> value }, + ).flowOn(Dispatchers.IO).toList() + } + assertEquals(1, calls) + assertSame(cause, error.cause) + } + } + + private class CallbackCancellation : CancellationException("callback cancelled") { + // An extra field prevents coroutine stack-trace recovery from copying this exception. + val marker = Any() + } + + private suspend fun assertDelivery(sse: Boolean) { + val expected = (0..199).map(Int::toString) + val body = expected.joinToString("") { if (sse) "data: $it\n\n" else "$it\n" } + withServer({ respond(body) }) { client -> + assertEquals(expected, client.stream(sse).buffer(0).flowOn(Dispatchers.IO).toList()) + } + } + + private suspend fun assertCancellation(sse: Boolean, headers: Boolean) = withContext(Dispatchers.Default) { + val entered = CompletableDeferred() + val release = CountDownLatch(1) + withServer({ + if (!headers) { + sendResponseHeaders(200, 0) + responseBody.write(if (sse) "data: preview\n\n".toByteArray() else "preview\n".toByteArray()) + responseBody.flush() + } + entered.complete(Unit) + check(release.await(10, TimeUnit.SECONDS)) + if (headers) respond("late") + }) { client -> + val received = CompletableDeferred() + val collection = launch(Dispatchers.IO) { client.stream(sse).collect { received.complete(Unit) } } + try { + withTimeout(5_000) { entered.await() } + if (!headers) withTimeout(5_000) { received.await() } + withTimeout(2_000) { collection.cancelAndJoin() } + assertTrue(collection.isCancelled) + assertEquals(1L, release.count) + } finally { + release.countDown() + collection.cancelAndJoin() + } + } + } + + private suspend fun assertFirst(sse: Boolean) = withIdleBody(sse) { client, release -> + val first = async(Dispatchers.IO) { client.stream(sse).first() } + try { + assertEquals("preview", withTimeout(5_000) { first.await() }) + assertEquals(1L, release.count) + } finally { + release.countDown() + first.cancelAndJoin() + } + } + + private suspend fun assertDownstreamFailure(sse: Boolean) = withIdleBody(sse) { client, release -> + val cause = ConsumerFailure() + val result = async(Dispatchers.IO) { + assertFailsWith { client.stream(sse).collect { throw cause } } + } + try { + assertSame(cause, withTimeout(5_000) { result.await() }) + assertEquals(1L, release.count) + } finally { + release.countDown() + result.cancelAndJoin() + } + } + + private class ConsumerFailure : IllegalStateException("consumer failed") { + val marker = Any() + } + + private suspend fun withIdleBody( + sse: Boolean, + block: suspend kotlinx.coroutines.CoroutineScope.(KoogHttpClient, CountDownLatch) -> Unit, + ) = withContext(Dispatchers.Default) { + val release = CountDownLatch(1) + withServer({ + sendResponseHeaders(200, 0) + responseBody.write(if (sse) "data: preview\n\n".toByteArray() else "preview\n".toByteArray()) + responseBody.flush() + check(release.await(10, TimeUnit.SECONDS)) + }) { client -> + try { + block(client, release) + } finally { + release.countDown() + } + } + } + + private suspend fun withServer(handler: HttpExchange.() -> Unit, block: suspend (KoogHttpClient) -> Unit) { + val executor = Executors.newCachedThreadPool() + val server = HttpServer.create(InetSocketAddress("127.0.0.1", 0), 0) + server.executor = executor + server.createContext("/stream") { exchange -> exchange.use { it.handler() } } + server.start() + val client = JavaKoogHttpClient.Factory().create( + clientName = "java-streaming", + baseUrl = "http://127.0.0.1:${server.address.port}", + ) + try { + block(client) + } finally { + client.close() + server.stop(0) + executor.shutdownNow() + } + } + + private fun HttpExchange.respond(body: String) { + val bytes = body.toByteArray(Charsets.UTF_8) + sendResponseHeaders(200, bytes.size.toLong()) + responseBody.write(bytes) + } + + private fun KoogHttpClient.stream(sse: Boolean): Flow = if (sse) { + sse( + path = "/stream", + requestBody = "{}", + requestBodyType = String::class, + decodeStreamingResponse = { it }, + processStreamingChunk = { it }, + ) + } else { + lines(path = "/stream", requestBody = "{}", requestBodyType = String::class) + } +}