Skip to content

fix(http): close cancelled Java streams and preserve backpressure #2285

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
kyboi wants to merge 2 commits into
base: develop
Choose a base branch
from
+677 −83
Open

fix(http): close cancelled Java streams and preserve backpressure #2285

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
62fd729
fix(http): parse complete Java SSE data events
kyboi Sep 30, 2026
f072c1e
fix(http): close cancelled Java streams and preserve backpressure
kyboi Sep 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
154 changes: 71 additions & 83 deletions http-client/http-client-java/src/main/kotlin/ai/koog/http/client/java/JavaKoogHttpClient.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

Expand All @@ -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.
Expand Down Expand Up @@ -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: <content>"
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
}
}

Expand All @@ -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 <O : Any, B : AutoCloseable> ProducerScope<O>.readStreamingResponse(
request: HttpRequest,
bodyHandler: HttpResponse.BodyHandler<B>,
process: suspend ProducerScope<O>.(B) -> Unit,
) {
val body = AtomicReference<B?>()
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()
}
}

Expand Down
Loading