From 95383b062489d8773509eaa1377627d0491645e5 Mon Sep 17 00:00:00 2001 From: qflen Date: Thu, 23 Apr 2026 00:49:23 +0200 Subject: [PATCH 01/10] feat(prompt): expose HTTP response headers on KoogHttpClientException for rate-limit retries Lets retry-after extractors honor structured rate-limit hints (retry-after, x-ratelimit-reset-*) from HTTP response headers instead of parsing them out of unstructured error bodies. - KoogHttpClientException gains a Map> of response headers. Keys are normalized to lowercase; added as a trailing optional constructor parameter (via JvmOverloads) to preserve source/binary compatibility. - Ktor, OkHttp, and Java HTTP clients populate headers at every throw site. - RetryAfterExtractor gains an extract(KoogHttpClientException) default overload that falls through to the existing extract(message) path, keeping fun interface SAM compatibility. - StandardHeaderRetryAfterExtractor parses retry-after (seconds / IMF-fixdate) and x-ratelimit-reset-* (OpenAI-style durations); CompositeRetryAfterExtractor stacks extractors. RetryConfig now defaults to Composite(Standard, Default) so existing users get header-aware retries for free when they become available. - RetryingLLMClient.calculateDelay dispatches to the header-aware overload when the caught throwable is a KoogHttpClientException (directly or one level deep via .cause, to cover wrappers like LLMClientException). Tests cover header capture end-to-end for all three HTTP clients, the new extractor's parsing rules, and virtual-time delay assertions through the retry layer. closes #1354 (cherry picked from commit a6bc613e2086320450c2c61e9f6bb633200ec6f2) --- docs/docs/prompts/handling-failures.md | 51 +++++ .../kotlin/ai/koog/http/client/Exceptions.kt | 17 +- .../kotlin/ai/koog/http/client/Headers.kt | 21 +++ .../http/client/java/JavaKoogHttpClient.kt | 3 + .../client/java/JavaKoogHttpClientTestBase.kt | 8 + .../http/client/ktor/KtorKoogHttpClient.kt | 8 +- .../client/ktor/KtorKoogHttpClientTestBase.kt | 8 + .../client/okhttp/OkHttpKoogHttpClient.kt | 5 +- .../okhttp/OkHttpKoogHttpClientTestBase.kt | 8 + .../client/test/BaseKoogHttpClientTest.kt | 32 ++++ .../ai/koog/http/client/test/MockWebServer.kt | 12 +- .../prompt-executor-clients/build.gradle.kts | 1 + .../executor/clients/retry/RetryConfig.kt | 176 +++++++++++++++++- .../clients/retry/RetryingLLMClient.kt | 19 +- .../clients/retry/RetryingLLMClientTest.kt | 141 +++++++++++++- .../StandardHeaderRetryAfterExtractorTest.kt | 165 ++++++++++++++++ 16 files changed, 654 insertions(+), 21 deletions(-) create mode 100644 http-client/http-client-core/src/commonMain/kotlin/ai/koog/http/client/Headers.kt create mode 100644 prompt/prompt-executor/prompt-executor-clients/src/commonTest/kotlin/ai/koog/prompt/executor/clients/retry/StandardHeaderRetryAfterExtractorTest.kt diff --git a/docs/docs/prompts/handling-failures.md b/docs/docs/prompts/handling-failures.md index 1ccdaaf13d..f16e444a7f 100644 --- a/docs/docs/prompts/handling-failures.md +++ b/docs/docs/prompts/handling-failures.md @@ -220,6 +220,57 @@ val config = RetryConfig( ``` +### Header-aware retry-after hints + +Many providers return rate-limit metadata in HTTP response headers rather than (or in addition +to) the error body. OpenAI in particular uses `retry-after`, `x-ratelimit-reset-requests`, and +`x-ratelimit-reset-tokens` to signal when a client should back off. + +When the underlying HTTP client raises a `KoogHttpClientException`, the exception exposes the +response headers on its `headers: Map>` property. Keys are normalized to +lowercase, so extractors can look them up without re-casing. + +The default `RetryConfig.retryAfterExtractor` is a composite that consults headers first and +falls back to parsing the error message: + +```kotlin +public val DEFAULT_RETRY_AFTER_EXTRACTOR: RetryAfterExtractor = CompositeRetryAfterExtractor( + StandardHeaderRetryAfterExtractor.DEFAULT, + DefaultRetryAfterExtractor +) +``` + +So out of the box, a 429 carrying `retry-after: 5` is honored without any extra configuration. +If you want to supply a custom header-aware extractor, implement `RetryAfterExtractor` and +override `extract(error: KoogHttpClientException)`: + +```kotlin +val config = RetryConfig( + retryAfterExtractor = RetryAfterExtractor { message -> + // message-based fallback + null + }.let { fallback -> + CompositeRetryAfterExtractor( + object : RetryAfterExtractor { + override fun extract(message: String): Duration? = null + override fun extract(error: KoogHttpClientException): Duration? = + error.headers["x-my-custom-retry-after"]?.firstOrNull()?.toLongOrNull()?.seconds + }, + fallback + ) + } +) +``` + +`StandardHeaderRetryAfterExtractor` understands: + +- `retry-after` as either delta-seconds or an IMF-fixdate (RFC 9110 §10.2.3); +- `x-ratelimit-reset-requests` / `x-ratelimit-reset-tokens` as OpenAI-style durations like + `1s`, `6m0s`, `100ms`. + +When multiple hints are present it returns the smallest strictly-positive delay so the client +doesn't sleep longer than any of the individual rate-limit buckets require. + ### Streaming with retry Streaming operations can optionally be retried. This feature is disabled by default. diff --git a/http-client/http-client-core/src/commonMain/kotlin/ai/koog/http/client/Exceptions.kt b/http-client/http-client-core/src/commonMain/kotlin/ai/koog/http/client/Exceptions.kt index 751ce63de0..1d39322ac5 100644 --- a/http-client/http-client-core/src/commonMain/kotlin/ai/koog/http/client/Exceptions.kt +++ b/http-client/http-client-core/src/commonMain/kotlin/ai/koog/http/client/Exceptions.kt @@ -1,14 +1,25 @@ package ai.koog.http.client +import kotlin.jvm.JvmOverloads + /** - * Base exception class for HTTP clients in koog + * Base exception class for HTTP clients in koog. + * + * @property clientName Name of the HTTP client that produced the error, for log attribution. + * @property statusCode HTTP status code returned by the server, or `null` if no response was received. + * @property errorBody Raw body of the failed response, or `null` if it could not be read. + * @property headers HTTP response headers captured from the failed response. + * Keys are normalized to lowercase so consumers can look headers up without re-casing + * (see [lowercaseHeaderKeys]); values preserve the order and formatting from the server. + * Defaults to an empty map when no response was available (e.g., connection errors). */ -public class KoogHttpClientException( +public class KoogHttpClientException @JvmOverloads constructor( public val clientName: String? = null, public val statusCode: Int? = null, public val errorBody: String? = null, message: String? = null, - cause: Throwable? = null + cause: Throwable? = null, + public val headers: Map> = emptyMap() ) : Exception( buildString { appendLine("Error from client: ${clientName ?: "unknown client"}") diff --git a/http-client/http-client-core/src/commonMain/kotlin/ai/koog/http/client/Headers.kt b/http-client/http-client-core/src/commonMain/kotlin/ai/koog/http/client/Headers.kt new file mode 100644 index 0000000000..8f2a205d07 --- /dev/null +++ b/http-client/http-client-core/src/commonMain/kotlin/ai/koog/http/client/Headers.kt @@ -0,0 +1,21 @@ +package ai.koog.http.client + +/** + * Returns a copy of this multi-valued header map with all keys converted to lowercase. + * + * Values for keys that differ only in case are concatenated into a single list so that + * a response containing e.g. both `Set-Cookie` and `set-cookie` becomes a single lowercase + * entry with both values. Empty maps are returned as [emptyMap] to avoid allocation. + * + * This matches the contract documented on [KoogHttpClientException.headers] and lets + * individual HTTP client implementations adapt their native header types without each + * reinventing the normalization. + */ +public fun Map>.lowercaseHeaderKeys(): Map> { + if (isEmpty()) return emptyMap() + val normalized = LinkedHashMap>(size) + for ((key, values) in this) { + normalized.getOrPut(key.lowercase()) { mutableListOf() }.addAll(values) + } + return normalized +} 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..716c4fc6a3 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 @@ -2,6 +2,7 @@ package ai.koog.http.client.java import ai.koog.http.client.KoogHttpClient import ai.koog.http.client.KoogHttpClientException +import ai.koog.http.client.lowercaseHeaderKeys import ai.koog.http.client.mergeHeaders import ai.koog.utils.io.SuitableForIO import io.github.oshai.kotlinlogging.KLogger @@ -71,6 +72,7 @@ public class JavaKoogHttpClient internal constructor( clientName = clientName, statusCode = response.statusCode(), errorBody = response.body(), + headers = response.headers().map().lowercaseHeaderKeys(), ) } @@ -196,6 +198,7 @@ public class JavaKoogHttpClient internal constructor( KoogHttpClientException( clientName = clientName, statusCode = response.statusCode(), + headers = response.headers().map().lowercaseHeaderKeys(), ) ) return@callbackFlow diff --git a/http-client/http-client-java/src/test/kotlin/ai/koog/http/client/java/JavaKoogHttpClientTestBase.kt b/http-client/http-client-java/src/test/kotlin/ai/koog/http/client/java/JavaKoogHttpClientTestBase.kt index 2156bcee73..a48ed63348 100644 --- a/http-client/http-client-java/src/test/kotlin/ai/koog/http/client/java/JavaKoogHttpClientTestBase.kt +++ b/http-client/http-client-java/src/test/kotlin/ai/koog/http/client/java/JavaKoogHttpClientTestBase.kt @@ -32,10 +32,18 @@ abstract class JavaKoogHttpClientTestBase : BaseKoogHttpClientTest() { override fun `test handle on non-success status`() = super.`test handle on non-success status`() + @Test + override fun testCaptureHeadersOnNonSuccess() = + super.testCaptureHeadersOnNonSuccess() + @Test override fun `test get SSE flow and collect events`() = super.`test get SSE flow and collect events`() + @Test + override fun testCaptureHeadersOnSseError() = + super.testCaptureHeadersOnSseError() + @Test override fun `test filter SSE events`() = super.`test filter SSE events`() diff --git a/http-client/http-client-ktor/src/commonMain/kotlin/ai/koog/http/client/ktor/KtorKoogHttpClient.kt b/http-client/http-client-ktor/src/commonMain/kotlin/ai/koog/http/client/ktor/KtorKoogHttpClient.kt index 074e01e89b..a5937deb76 100644 --- a/http-client/http-client-ktor/src/commonMain/kotlin/ai/koog/http/client/ktor/KtorKoogHttpClient.kt +++ b/http-client/http-client-ktor/src/commonMain/kotlin/ai/koog/http/client/ktor/KtorKoogHttpClient.kt @@ -2,6 +2,7 @@ package ai.koog.http.client.ktor import ai.koog.http.client.KoogHttpClient import ai.koog.http.client.KoogHttpClientException +import ai.koog.http.client.lowercaseHeaderKeys import ai.koog.http.client.mergeHeaders import ai.koog.utils.io.SuitableForIO import io.github.oshai.kotlinlogging.KLogger @@ -27,6 +28,7 @@ import io.ktor.client.statement.HttpResponse import io.ktor.client.statement.bodyAsChannel import io.ktor.client.statement.bodyAsText import io.ktor.http.ContentType +import io.ktor.http.Headers import io.ktor.http.HttpHeaders import io.ktor.http.HttpMethod import io.ktor.http.URLBuilder @@ -95,6 +97,7 @@ public class KtorKoogHttpClient internal constructor( clientName = clientName, statusCode = response.status.value, errorBody = response.bodyAsText(), + headers = response.headers.toResponseHeaderMap(), ) } @@ -110,6 +113,8 @@ public class KtorKoogHttpClient internal constructor( } } } + private fun Headers.toResponseHeaderMap(): Map> = + entries().associate { it.key to it.value }.lowercaseHeaderKeys() override suspend fun get( path: String, @@ -210,7 +215,8 @@ public class KtorKoogHttpClient internal constructor( statusCode = e.response?.status?.value, errorBody = errorBody, message = e.message, - cause = e + cause = e, + headers = e.response?.headers?.toResponseHeaderMap() ?: emptyMap() ) } catch (e: CancellationException) { throw e diff --git a/http-client/http-client-ktor/src/jvmTest/kotlin/ai/koog/http/client/ktor/KtorKoogHttpClientTestBase.kt b/http-client/http-client-ktor/src/jvmTest/kotlin/ai/koog/http/client/ktor/KtorKoogHttpClientTestBase.kt index 96fc8e83fe..8b3c213b12 100644 --- a/http-client/http-client-ktor/src/jvmTest/kotlin/ai/koog/http/client/ktor/KtorKoogHttpClientTestBase.kt +++ b/http-client/http-client-ktor/src/jvmTest/kotlin/ai/koog/http/client/ktor/KtorKoogHttpClientTestBase.kt @@ -46,10 +46,18 @@ abstract class KtorKoogHttpClientTestBase : BaseKoogHttpClientTest() { override fun `test handle on non-success status`() = super.`test handle on non-success status`() + @Test + override fun testCaptureHeadersOnNonSuccess() = + super.testCaptureHeadersOnNonSuccess() + @Test override fun `test get SSE flow and collect events`() = super.`test get SSE flow and collect events`() + @Test + override fun testCaptureHeadersOnSseError() = + super.testCaptureHeadersOnSseError() + @Test override fun `test filter SSE events`() = super.`test filter SSE events`() diff --git a/http-client/http-client-okhttp/src/main/kotlin/ai/koog/http/client/okhttp/OkHttpKoogHttpClient.kt b/http-client/http-client-okhttp/src/main/kotlin/ai/koog/http/client/okhttp/OkHttpKoogHttpClient.kt index 742f2a89de..fac6fd2f5c 100644 --- a/http-client/http-client-okhttp/src/main/kotlin/ai/koog/http/client/okhttp/OkHttpKoogHttpClient.kt +++ b/http-client/http-client-okhttp/src/main/kotlin/ai/koog/http/client/okhttp/OkHttpKoogHttpClient.kt @@ -2,6 +2,7 @@ package ai.koog.http.client.okhttp import ai.koog.http.client.KoogHttpClient import ai.koog.http.client.KoogHttpClientException +import ai.koog.http.client.lowercaseHeaderKeys import ai.koog.http.client.mergeHeaders import ai.koog.utils.io.SuitableForIO import io.github.oshai.kotlinlogging.KLogger @@ -72,6 +73,7 @@ public class OkHttpKoogHttpClient internal constructor( clientName = clientName, statusCode = response.code, errorBody = response.body.string(), + headers = response.headers.toMultimap().lowercaseHeaderKeys(), ) } @@ -209,7 +211,8 @@ public class OkHttpKoogHttpClient internal constructor( statusCode = response?.code, errorBody = response?.body?.string(), message = t?.message, - cause = t + cause = t, + headers = response?.headers?.toMultimap()?.lowercaseHeaderKeys() ?: emptyMap() ) logger.error(exception) { exception.message } close(exception) diff --git a/http-client/http-client-okhttp/src/test/kotlin/ai/koog/http/client/okhttp/OkHttpKoogHttpClientTestBase.kt b/http-client/http-client-okhttp/src/test/kotlin/ai/koog/http/client/okhttp/OkHttpKoogHttpClientTestBase.kt index abd3b562f0..0a835eca40 100644 --- a/http-client/http-client-okhttp/src/test/kotlin/ai/koog/http/client/okhttp/OkHttpKoogHttpClientTestBase.kt +++ b/http-client/http-client-okhttp/src/test/kotlin/ai/koog/http/client/okhttp/OkHttpKoogHttpClientTestBase.kt @@ -32,10 +32,18 @@ abstract class OkHttpKoogHttpClientTestBase : BaseKoogHttpClientTest() { override fun `test handle on non-success status`() = super.`test handle on non-success status`() + @Test + override fun testCaptureHeadersOnNonSuccess() = + super.testCaptureHeadersOnNonSuccess() + @Test override fun `test get SSE flow and collect events`() = super.`test get SSE flow and collect events`() + @Test + override fun testCaptureHeadersOnSseError() = + super.testCaptureHeadersOnSseError() + @Test override fun `test filter SSE events`() = super.`test filter SSE events`() diff --git a/http-client/http-client-test/src/main/kotlin/ai/koog/http/client/test/BaseKoogHttpClientTest.kt b/http-client/http-client-test/src/main/kotlin/ai/koog/http/client/test/BaseKoogHttpClientTest.kt index 2377ea4795..951bdb105f 100644 --- a/http-client/http-client-test/src/main/kotlin/ai/koog/http/client/test/BaseKoogHttpClientTest.kt +++ b/http-client/http-client-test/src/main/kotlin/ai/koog/http/client/test/BaseKoogHttpClientTest.kt @@ -185,6 +185,38 @@ abstract class BaseKoogHttpClientTest { } } + open fun testCaptureHeadersOnNonSuccess(): Unit = runTest { + mockServer.start( + postEndpoints = listOf( + MockWebServer.PostEndpointConfig( + path = "/rate-limited", + responseBody = "rate limit exceeded", + statusCode = HttpStatusCode.TooManyRequests, + contentType = ContentType.Text.Plain, + responseHeaders = mapOf( + "Retry-After" to "5", + "X-RateLimit-Reset-Tokens" to "6m0s" + ) + ) + ) + ) + + val client = createClient() + + try { + client.post( + path = mockServer.url("/rate-limited"), + request = "PAYLOAD", + ) + fail("Expected a KoogHttpClientException for 429") + } catch (e: KoogHttpClientException) { + assertEquals(429, e.statusCode) + // Header keys are normalized to lowercase regardless of how the server sent them. + assertEquals(listOf("5"), e.headers["retry-after"]) + assertEquals(listOf("6m0s"), e.headers["x-ratelimit-reset-tokens"]) + } + } + @Suppress("FunctionName") open fun `test get SSE flow and collect events`(): Unit = runTest { val events = listOf("event1", "event2", "event3") diff --git a/http-client/http-client-test/src/main/kotlin/ai/koog/http/client/test/MockWebServer.kt b/http-client/http-client-test/src/main/kotlin/ai/koog/http/client/test/MockWebServer.kt index 7393d738b9..491890327b 100644 --- a/http-client/http-client-test/src/main/kotlin/ai/koog/http/client/test/MockWebServer.kt +++ b/http-client/http-client-test/src/main/kotlin/ai/koog/http/client/test/MockWebServer.kt @@ -41,7 +41,8 @@ class MockWebServer { val responseBody: String, val statusCode: HttpStatusCode = HttpStatusCode.OK, val contentType: ContentType = ContentType.Text.Plain, - val expectedParameters: Map = emptyMap() + val expectedParameters: Map = emptyMap(), + val responseHeaders: Map = emptyMap() ) /** @@ -53,7 +54,8 @@ class MockWebServer { val statusCode: HttpStatusCode = HttpStatusCode.OK, val contentType: ContentType = ContentType.Text.Plain, val expectedParameters: Map = emptyMap(), - val expectedHeaders: Map = emptyMap() + val expectedHeaders: Map = emptyMap(), + val responseHeaders: Map = emptyMap() ) /** @@ -116,6 +118,9 @@ class MockWebServer { } } + config.responseHeaders.forEach { (name, value) -> + call.response.header(name, value) + } call.respondText( text = config.responseBody, contentType = config.contentType, @@ -153,6 +158,9 @@ class MockWebServer { } } + config.responseHeaders.forEach { (name, value) -> + call.response.header(name, value) + } call.respondText( text = config.responseBody, contentType = config.contentType, diff --git a/prompt/prompt-executor/prompt-executor-clients/build.gradle.kts b/prompt/prompt-executor/prompt-executor-clients/build.gradle.kts index 10c16b7282..10ab1cfe32 100644 --- a/prompt/prompt-executor/prompt-executor-clients/build.gradle.kts +++ b/prompt/prompt-executor/prompt-executor-clients/build.gradle.kts @@ -10,6 +10,7 @@ kotlin { dependencies { api(project(":prompt:prompt-model")) api(project(":agents:agents-tools")) + api(project(":http-client:http-client-core")) api(libs.kotlinx.coroutines.core) api(project(":prompt:prompt-structure")) implementation(libs.oshai.kotlin.logging) diff --git a/prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/retry/RetryConfig.kt b/prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/retry/RetryConfig.kt index f57c83fe7f..8b740fdf29 100644 --- a/prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/retry/RetryConfig.kt +++ b/prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/retry/RetryConfig.kt @@ -1,10 +1,13 @@ package ai.koog.prompt.executor.clients.retry +import ai.koog.http.client.KoogHttpClientException import kotlin.jvm.JvmField import kotlin.jvm.JvmOverloads +import kotlin.time.Clock import kotlin.time.Duration import kotlin.time.Duration.Companion.milliseconds import kotlin.time.Duration.Companion.seconds +import kotlin.time.Instant /** * Configuration for retry behavior in LLM client operations. @@ -15,7 +18,9 @@ import kotlin.time.Duration.Companion.seconds * @property backoffMultiplier Multiplier for exponential backoff * @property jitterFactor Random jitter factor (0.0 to 1.0) * @property retryablePatterns Patterns to identify retryable errors - * @property retryAfterExtractor Optional extractor for retry-after hints + * @property retryAfterExtractor Optional extractor for retry-after hints. Defaults to + * [DEFAULT_RETRY_AFTER_EXTRACTOR], which prefers HTTP response headers and falls back to + * parsing the error message when no usable header is present. */ public data class RetryConfig @JvmOverloads constructor( val maxAttempts: Int = 3, @@ -24,7 +29,7 @@ public data class RetryConfig @JvmOverloads constructor( val backoffMultiplier: Double = 2.0, val jitterFactor: Double = 0.1, val retryablePatterns: List = DEFAULT_PATTERNS, - val retryAfterExtractor: RetryAfterExtractor? = DefaultRetryAfterExtractor + val retryAfterExtractor: RetryAfterExtractor? = DEFAULT_RETRY_AFTER_EXTRACTOR ) { init { require(maxAttempts >= 1) { "maxAttempts must be at least 1" } @@ -68,6 +73,17 @@ public data class RetryConfig @JvmOverloads constructor( RetryablePattern.Keyword("service unavailable") ) + /** + * Default retry-after extractor: prefers structured HTTP response headers + * ([StandardHeaderRetryAfterExtractor]) and falls back to parsing the error message + * ([DefaultRetryAfterExtractor]) when no usable header is present. + */ + @JvmField + public val DEFAULT_RETRY_AFTER_EXTRACTOR: RetryAfterExtractor = CompositeRetryAfterExtractor( + StandardHeaderRetryAfterExtractor.DEFAULT, + DefaultRetryAfterExtractor + ) + /** * Conservative configuration - fewer retries, longer delays. */ @@ -168,7 +184,12 @@ public sealed class RetryablePattern { } /** - * Extracts retry-after hints from error messages. + * Extracts retry-after hints from transient errors. + * + * Implementations can inspect either the error message (for backward compatibility with + * lambda-style callers) or a [KoogHttpClientException] carrying structured HTTP response + * metadata. The default [extract] overload bridges the two: implementations that only parse + * messages remain callable on [KoogHttpClientException] without changes. */ public fun interface RetryAfterExtractor { /** @@ -178,10 +199,23 @@ public fun interface RetryAfterExtractor { * @return The extracted retry-after duration, or null if no valid duration could be determined. */ public fun extract(message: String): Duration? + + /** + * Extracts a retry-after duration from a [KoogHttpClientException], which exposes + * HTTP response headers, the status code, and the error body. + * + * The default implementation delegates to [extract] with the exception's message so that + * existing SAM lambdas continue to work. Header-aware extractors should override this. + * + * @param error The HTTP client exception carrying the response metadata. + * @return The extracted retry-after duration, or null if no valid hint could be determined. + */ + public fun extract(error: KoogHttpClientException): Duration? = + extract(error.message ?: "") } /** - * Default implementation that extracts common retry-after patterns. + * Default implementation that extracts common retry-after patterns from free-form error text. */ public object DefaultRetryAfterExtractor : RetryAfterExtractor { private val patterns = listOf( @@ -202,3 +236,137 @@ public object DefaultRetryAfterExtractor : RetryAfterExtractor { return null } } + +/** + * Extractor that reads structured rate-limit hints from HTTP response headers attached to a + * [KoogHttpClientException]. + * + * Recognized headers (case-insensitive; keys on [KoogHttpClientException.headers] are + * already normalized to lowercase): + * + * - `retry-after` — integer delta-seconds or an IMF-fixdate (RFC 9110 §10.2.3). Non-integer, + * non-date values are ignored. + * - `x-ratelimit-reset-requests`, `x-ratelimit-reset-tokens` — OpenAI-style durations such as + * `1s`, `6m0s`, or `100ms` (RFC 9110 does not standardize these; the format is OpenAI's). + * + * When multiple hints are present, the smallest strictly-positive delay wins; otherwise the + * extractor returns `null` so the caller can fall back to exponential backoff or a secondary + * extractor (see [CompositeRetryAfterExtractor]). + * + * This extractor only consults headers; [extract] with a bare message always returns `null`. + * Compose it with [DefaultRetryAfterExtractor] to retain message-based fallback behavior. + * + * @property clock Clock used to evaluate HTTP-date values relative to "now". Override in tests. + */ +public class StandardHeaderRetryAfterExtractor( + private val clock: Clock = Clock.System +) : RetryAfterExtractor { + + override fun extract(message: String): Duration? = null + + override fun extract(error: KoogHttpClientException): Duration? { + val headers = error.headers + if (headers.isEmpty()) return null + + val candidates = buildList { + headers.firstValue("retry-after")?.let { parseRetryAfter(it)?.let(::add) } + headers.firstValue("x-ratelimit-reset-requests")?.let { parseOpenAIDuration(it)?.let(::add) } + headers.firstValue("x-ratelimit-reset-tokens")?.let { parseOpenAIDuration(it)?.let(::add) } + } + + return candidates.filter { it > Duration.ZERO }.minOrNull() + } + + private fun Map>.firstValue(name: String): String? = + this[name]?.firstOrNull()?.trim()?.takeIf { it.isNotEmpty() } + + private fun parseRetryAfter(value: String): Duration? { + val trimmed = value.trim() + trimmed.toLongOrNull()?.let { return if (it >= 0) it.seconds else null } + val target = parseHttpDate(trimmed) ?: return null + val now = clock.now() + val delta = target - now + return if (delta > Duration.ZERO) delta else Duration.ZERO + } + + /** + * Parses an IMF-fixdate (`"Sun, 06 Nov 1994 08:49:37 GMT"`) as an [Instant]. RFC 9110 §5.6.7 + * requires new clients to generate this form; obsolete RFC 850 and asctime formats must be + * accepted but are rare in practice and intentionally unsupported here. + */ + private fun parseHttpDate(value: String): Instant? { + val match = IMF_FIXDATE.matchEntire(value) ?: return null + val day = match.groupValues[2].padStart(2, '0') + val month = MONTH_NUMBERS[match.groupValues[3]] ?: return null + val year = match.groupValues[4] + val time = "${match.groupValues[5]}:${match.groupValues[6]}:${match.groupValues[7]}" + return runCatching { Instant.parse("$year-$month-${day}T${time}Z") }.getOrNull() + } + + private fun parseOpenAIDuration(value: String): Duration? { + val trimmed = value.trim() + if (trimmed.isEmpty()) return null + val tokens = DURATION_TOKEN.findAll(trimmed).toList() + if (tokens.isEmpty()) return null + if (tokens.sumOf { it.value.length } != trimmed.length) return null + + var totalMillis = 0.0 + for (token in tokens) { + val amount = token.groupValues[1].toDoubleOrNull() ?: return null + val unitMillis = when (token.groupValues[2].lowercase()) { + "ms" -> 1.0 + "s" -> 1000.0 + "m" -> 60_000.0 + "h" -> 3_600_000.0 + else -> return null + } + totalMillis += amount * unitMillis + } + return totalMillis.toLong().milliseconds + } + + /** + * Companion exposing a shared default instance backed by [Clock.System]. + */ + public companion object { + /** Shared instance using the system clock. */ + @JvmField + public val DEFAULT: StandardHeaderRetryAfterExtractor = StandardHeaderRetryAfterExtractor() + + private val IMF_FIXDATE = Regex( + "(Mon|Tue|Wed|Thu|Fri|Sat|Sun), " + + "(\\d{2}) " + + "(Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec) " + + "(\\d{4}) " + + "(\\d{2}):(\\d{2}):(\\d{2}) GMT" + ) + + private val MONTH_NUMBERS = mapOf( + "Jan" to "01", "Feb" to "02", "Mar" to "03", "Apr" to "04", + "May" to "05", "Jun" to "06", "Jul" to "07", "Aug" to "08", + "Sep" to "09", "Oct" to "10", "Nov" to "11", "Dec" to "12" + ) + + private val DURATION_TOKEN = Regex("(\\d+(?:\\.\\d+)?)(ms|h|m|s)", RegexOption.IGNORE_CASE) + } +} + +/** + * Runs a list of [RetryAfterExtractor]s in order and returns the first non-null result. Useful + * for stacking a header-aware extractor over a message-based fallback — the ordering models + * the idea that structured hints are preferred when present. + * + * @param extractors Extractors to consult, in order. + */ +public class CompositeRetryAfterExtractor( + private val extractors: List +) : RetryAfterExtractor { + /** Convenience varargs constructor. */ + public constructor(vararg extractors: RetryAfterExtractor) : this(extractors.toList()) + + override fun extract(message: String): Duration? = + extractors.firstNotNullOfOrNull { it.extract(message) } + + override fun extract(error: KoogHttpClientException): Duration? = + extractors.firstNotNullOfOrNull { it.extract(error) } +} diff --git a/prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClient.kt b/prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClient.kt index c940690cce..58d9992664 100644 --- a/prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClient.kt +++ b/prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClient.kt @@ -1,6 +1,7 @@ package ai.koog.prompt.executor.clients.retry import ai.koog.agents.core.tools.ToolDescriptor +import ai.koog.http.client.KoogHttpClientException import ai.koog.prompt.Prompt import ai.koog.prompt.dsl.ModerationResult import ai.koog.prompt.executor.clients.LLMClient @@ -192,11 +193,21 @@ public class RetryingLLMClient @JvmOverloads constructor( } private fun calculateDelay(attempt: Int, error: Throwable? = null): Duration { - // Check for retry-after hint in error message - error?.message?.let { message -> - config.retryAfterExtractor?.extract(message)?.let { retryAfter -> - return retryAfter + // Check for retry-after hint. Prefer the header-aware overload when the caught + // throwable is a KoogHttpClientException (directly or wrapped one level deep — some + // clients re-throw as LLMClientException with the original exception as cause) so + // structured response metadata is available to the extractor; otherwise fall back + // to the message-based overload. + val extractor = config.retryAfterExtractor + if (extractor != null && error != null) { + val koogException = error as? KoogHttpClientException + ?: error.cause as? KoogHttpClientException + val retryAfter = if (koogException != null) { + extractor.extract(koogException) + } else { + error.message?.let { extractor.extract(it) } } + if (retryAfter != null) return retryAfter } // Exponential backoff with jitter diff --git a/prompt/prompt-executor/prompt-executor-clients/src/commonTest/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClientTest.kt b/prompt/prompt-executor/prompt-executor-clients/src/commonTest/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClientTest.kt index daffef7aa9..3d740a8ad1 100644 --- a/prompt/prompt-executor/prompt-executor-clients/src/commonTest/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClientTest.kt +++ b/prompt/prompt-executor/prompt-executor-clients/src/commonTest/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClientTest.kt @@ -1,6 +1,7 @@ package ai.koog.prompt.executor.clients.retry import ai.koog.agents.core.tools.ToolDescriptor +import ai.koog.http.client.KoogHttpClientException import ai.koog.prompt.Prompt import ai.koog.prompt.dsl.ModerationResult import ai.koog.prompt.dsl.prompt @@ -31,6 +32,7 @@ import kotlin.test.assertEquals import kotlin.test.assertFailsWith import kotlin.test.assertSame import kotlin.time.Duration.Companion.milliseconds +import kotlin.time.Duration.Companion.seconds class RetryingLLMClientTest { @@ -53,6 +55,131 @@ class RetryingLLMClientTest { metaInfo = testMetaInfo ) + @Test + fun testRetryAfterHeaderOverridesExponentialBackoff() = runTest { + // Error message alone doesn't carry a retry-after hint; only the structured header + // does. Using virtual time we verify the sleep was exactly 7s, proving calculateDelay + // dispatched to the KoogHttpClientException overload rather than exponential backoff. + val retryDelayHeader = KoogHttpClientException( + clientName = "TestClient", + statusCode = 429, + errorBody = "rate limited", + message = "rate limited", + headers = mapOf("retry-after" to listOf("7")) + ) + + val mockClient = MockLLMClient( + executeResponse = testResponse, + failuresBeforeSuccess = 1, + failure = retryDelayHeader + ) + + val retryingClient = RetryingLLMClient( + mockClient, + RetryConfig( + maxAttempts = 2, + initialDelay = 500.milliseconds, + maxDelay = 60.seconds + ) + ) + + val startTime = testScheduler.currentTime + val result = retryingClient.execute(testPrompt, testModel, emptyList()) + val elapsed = testScheduler.currentTime - startTime + + assertEquals(testResponse, result) + assertEquals(2, mockClient.executeCalls) + assertEquals(7_000L, elapsed) + } + + @Test + fun testHeaderHintWinsOverParseableMessage() = runTest { + // Header says 3s, message says 99s. The composite extractor tries headers first, so + // the header wins and the client sleeps 3s instead of 99s. + val error = KoogHttpClientException( + clientName = "TestClient", + statusCode = 429, + errorBody = "retry-after: 99", + message = "retry-after: 99", + headers = mapOf("retry-after" to listOf("3")) + ) + + val mockClient = MockLLMClient( + executeResponse = testResponse, + failuresBeforeSuccess = 1, + failure = error + ) + + val retryingClient = RetryingLLMClient( + mockClient, + RetryConfig(maxAttempts = 2, initialDelay = 10.milliseconds) + ) + + val startTime = testScheduler.currentTime + retryingClient.execute(testPrompt, testModel, emptyList()) + val elapsed = testScheduler.currentTime - startTime + + assertEquals(3_000L, elapsed) + } + + @Test + fun testUnwrapsKoogHttpClientExceptionFromCause() = runTest { + // Some upstream LLM clients re-throw as a domain exception that carries the + // KoogHttpClientException as cause (Ollama does this today). The retry layer should + // still see the headers and dispatch to the header-aware extractor overload. + val httpError = KoogHttpClientException( + clientName = "TestClient", + statusCode = 429, + errorBody = "rate limited", + message = "rate limited", + headers = mapOf("retry-after" to listOf("4")) + ) + // Wrapping message includes the status so the existing shouldRetry patterns match. + val wrapped = RuntimeException("LLM call failed with status 429", httpError) + + val mockClient = MockLLMClient( + executeResponse = testResponse, + failuresBeforeSuccess = 1, + failure = wrapped + ) + + val retryingClient = RetryingLLMClient( + mockClient, + RetryConfig(maxAttempts = 2, initialDelay = 10.milliseconds) + ) + + val startTime = testScheduler.currentTime + retryingClient.execute(testPrompt, testModel, emptyList()) + val elapsed = testScheduler.currentTime - startTime + + assertEquals(4_000L, elapsed) + } + + @Test + fun testMessageBasedExtractorStillUsedForNonKoogExceptions() = runTest { + // When the throwable isn't a KoogHttpClientException, calculateDelay must continue to + // consult the message-based extract(message) overload. The default extractor picks up + // "retry-after: 3" from the message text. + val mockClient = MockLLMClient( + executeResponse = testResponse, + failuresBeforeSuccess = 1, + failureMessage = "Error 429 retry-after: 3" + ) + + val retryingClient = RetryingLLMClient( + mockClient, + RetryConfig(maxAttempts = 2, initialDelay = 10.milliseconds) + ) + + val startTime = testScheduler.currentTime + val result = retryingClient.execute(testPrompt, testModel, emptyList()) + val elapsed = testScheduler.currentTime - startTime + + assertEquals(testResponse, result) + assertEquals(2, mockClient.executeCalls) + assertEquals(3_000L, elapsed) + } + @Test fun testSucceedOnFirstAttempt() = runTest { val mockClient = MockLLMClient( @@ -544,9 +671,11 @@ class RetryingLLMClientTest { private var failuresBeforeSuccess: Int = 0, private var streamFailuresBeforeSuccess: Int = 0, private val failureMessage: String = "Mock failure", + private val failure: Throwable? = null, private val throwCancellation: Boolean = false, private val llmProvider: LLMProvider = LLMProvider.OpenAI, ) : LLMClient() { + private fun failureToThrow(): Throwable = failure ?: RuntimeException(failureMessage) val basicJsonSchemaGeneratorDefault = object : BasicJsonSchemaGenerator() {} val standardJsonSchemaGeneratorDefault = object : StandardJsonSchemaGenerator() {} @@ -580,7 +709,7 @@ class RetryingLLMClientTest { if (executeFailures < failuresBeforeSuccess) { executeFailures++ - throw RuntimeException(failureMessage) + throw failureToThrow() } return executeResponse!! @@ -595,7 +724,7 @@ class RetryingLLMClientTest { if (streamFailures < streamFailuresBeforeSuccess) { streamFailures++ - throw RuntimeException(failureMessage) + throw failureToThrow() } streamResponse.collect { emit(it) } @@ -610,7 +739,7 @@ class RetryingLLMClientTest { if (multipleChoicesFailures < failuresBeforeSuccess) { multipleChoicesFailures++ - throw RuntimeException(failureMessage) + throw failureToThrow() } return multipleChoicesResponse!! @@ -621,7 +750,7 @@ class RetryingLLMClientTest { if (moderateFailures < failuresBeforeSuccess) { moderateFailures++ - throw RuntimeException(failureMessage) + throw failureToThrow() } return moderateResponse @@ -635,7 +764,7 @@ class RetryingLLMClientTest { if (embedFailures < failuresBeforeSuccess) { embedFailures++ - throw RuntimeException(failureMessage) + throw failureToThrow() } return embedResponse @@ -649,7 +778,7 @@ class RetryingLLMClientTest { if (batchEmbedFailures < failuresBeforeSuccess) { batchEmbedFailures++ - throw RuntimeException(failureMessage) + throw failureToThrow() } return batchEmbedResponse diff --git a/prompt/prompt-executor/prompt-executor-clients/src/commonTest/kotlin/ai/koog/prompt/executor/clients/retry/StandardHeaderRetryAfterExtractorTest.kt b/prompt/prompt-executor/prompt-executor-clients/src/commonTest/kotlin/ai/koog/prompt/executor/clients/retry/StandardHeaderRetryAfterExtractorTest.kt new file mode 100644 index 0000000000..fc12a9f405 --- /dev/null +++ b/prompt/prompt-executor/prompt-executor-clients/src/commonTest/kotlin/ai/koog/prompt/executor/clients/retry/StandardHeaderRetryAfterExtractorTest.kt @@ -0,0 +1,165 @@ +package ai.koog.prompt.executor.clients.retry + +import ai.koog.http.client.KoogHttpClientException +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertNull +import kotlin.time.Clock +import kotlin.time.Duration.Companion.milliseconds +import kotlin.time.Duration.Companion.seconds +import kotlin.time.Instant + +class StandardHeaderRetryAfterExtractorTest { + + private val extractor = StandardHeaderRetryAfterExtractor() + + private fun error(headers: Map>): KoogHttpClientException = + KoogHttpClientException( + clientName = "TestClient", + statusCode = 429, + errorBody = null, + message = null, + cause = null, + headers = headers + ) + + @Test + fun testReturnsNullWhenNoHeaders() { + assertNull(extractor.extract(error(emptyMap()))) + } + + @Test + fun testReturnsNullWhenNoRelevantHeaders() { + val e = error(mapOf("x-request-id" to listOf("abc-123"), "content-type" to listOf("application/json"))) + assertNull(extractor.extract(e)) + } + + @Test + fun testExtractsNumericRetryAfter() { + assertEquals(5.seconds, extractor.extract(error(mapOf("retry-after" to listOf("5"))))) + assertEquals(120.seconds, extractor.extract(error(mapOf("retry-after" to listOf("120"))))) + } + + @Test + fun testIgnoresRetryAfterZeroAsNonCandidate() { + // Zero is not a "nonzero" candidate per the extractor contract; no other hints → null. + assertNull(extractor.extract(error(mapOf("retry-after" to listOf("0"))))) + } + + @Test + fun testIgnoresNegativeRetryAfter() { + assertNull(extractor.extract(error(mapOf("retry-after" to listOf("-3"))))) + } + + @Test + fun testIgnoresUnparseableRetryAfter() { + assertNull(extractor.extract(error(mapOf("retry-after" to listOf("soon"))))) + assertNull(extractor.extract(error(mapOf("retry-after" to listOf(""))))) + } + + @Test + fun testExtractsHttpDateRetryAfter() { + // Reference "now" = 2026-01-01T00:00:00Z; server says retry at 2026-01-01T00:00:42Z. + val now = Instant.parse("2026-01-01T00:00:00Z") + val fixedExtractor = StandardHeaderRetryAfterExtractor(clock = fixedClockAt(now)) + + val e = error(mapOf("retry-after" to listOf("Thu, 01 Jan 2026 00:00:42 GMT"))) + assertEquals(42.seconds, fixedExtractor.extract(e)) + } + + @Test + fun testHttpDateRetryAfterInPastClampsToZeroAndIsIgnored() { + val now = Instant.parse("2026-01-01T00:01:00Z") + val fixedExtractor = StandardHeaderRetryAfterExtractor(clock = fixedClockAt(now)) + + val e = error(mapOf("retry-after" to listOf("Thu, 01 Jan 2026 00:00:00 GMT"))) + // Past date → Duration.ZERO → filtered out as non-positive → null. + assertNull(fixedExtractor.extract(e)) + } + + @Test + fun testExtractsOpenAIStyleResetRequests() { + assertEquals(1.seconds, extractor.extract(error(mapOf("x-ratelimit-reset-requests" to listOf("1s"))))) + assertEquals(100.milliseconds, extractor.extract(error(mapOf("x-ratelimit-reset-requests" to listOf("100ms"))))) + } + + @Test + fun testExtractsOpenAIStyleResetTokens() { + assertEquals( + (6 * 60).seconds, + extractor.extract(error(mapOf("x-ratelimit-reset-tokens" to listOf("6m0s")))) + ) + } + + @Test + fun testParsesCompoundOpenAIDuration() { + val e = error(mapOf("x-ratelimit-reset-tokens" to listOf("1h2m3s"))) + assertEquals((1 * 3600 + 2 * 60 + 3).seconds, extractor.extract(e)) + } + + @Test + fun testRejectsMalformedOpenAIDuration() { + // Extra characters that don't match any token should reject the whole value. + assertNull(extractor.extract(error(mapOf("x-ratelimit-reset-tokens" to listOf("6minutes"))))) + assertNull(extractor.extract(error(mapOf("x-ratelimit-reset-tokens" to listOf("1s foo"))))) + assertNull(extractor.extract(error(mapOf("x-ratelimit-reset-tokens" to listOf(""))))) + } + + @Test + fun testReturnsSmallestNonzeroWhenMultipleHeadersPresent() { + val e = error( + mapOf( + "retry-after" to listOf("10"), + "x-ratelimit-reset-requests" to listOf("500ms"), + "x-ratelimit-reset-tokens" to listOf("6m0s") + ) + ) + assertEquals(500.milliseconds, extractor.extract(e)) + } + + @Test + fun testIgnoresHeadersWithUppercaseKeysFromCaller() { + // Extractor looks headers up by lowercase key. KoogHttpClientException.headers is + // contractually lowercase; if a caller builds an exception with uppercase keys by + // mistake, the extractor simply finds nothing. + val e = error(mapOf("Retry-After" to listOf("5"))) + assertNull(extractor.extract(e)) + } + + @Test + fun testExtractOnMessageAlwaysReturnsNull() { + // StandardHeaderRetryAfterExtractor intentionally does not parse free-form messages; + // compose with DefaultRetryAfterExtractor for that. + assertNull(extractor.extract("Retry after 5 seconds")) + } + + @Test + fun testCompositeExtractorPrefersHeadersThenMessage() { + val composite = CompositeRetryAfterExtractor( + StandardHeaderRetryAfterExtractor.DEFAULT, + DefaultRetryAfterExtractor + ) + + val withHeader = KoogHttpClientException( + clientName = "TestClient", + statusCode = 429, + errorBody = "Retry after 99 seconds", + message = "Retry after 99 seconds", + headers = mapOf("retry-after" to listOf("3")) + ) + assertEquals(3.seconds, composite.extract(withHeader)) + + val withoutHeader = KoogHttpClientException( + clientName = "TestClient", + statusCode = 429, + errorBody = "Retry after 99 seconds", + message = "Retry after 99 seconds", + headers = emptyMap() + ) + assertEquals(99.seconds, composite.extract(withoutHeader)) + } + + private fun fixedClockAt(instant: Instant): Clock = object : Clock { + override fun now(): Instant = instant + } +} From da5b42bf3e1717cf25bbecc9a758a609c482ad32 Mon Sep 17 00:00:00 2001 From: qflen Date: Thu, 23 Apr 2026 02:22:43 +0200 Subject: [PATCH 02/10] fix(prompt): walk wrapped cause in RetryingLLMClient.shouldRetry shouldRetry only inspected error.message, so a wrapper whose own message lacked retry tokens (e.g. "LLM call failed") short-circuited to false before calculateDelay - which does walk .cause - ever ran. Factor the unwrap into Throwable.unwrapHttpException() and use it from both, then extract the retry-after dispatch into retryAfterHint() so calculateDelay reads as a single ?: expression. Regression test: wrapper with no retryable tokens in its own message but a 429 KoogHttpClientException as cause must still retry. (cherry picked from commit 0d1af974c1b536a089ab0fcfff982a59cc4c6514) --- .../clients/retry/RetryingLLMClient.kt | 56 ++++++++++--------- .../clients/retry/RetryingLLMClientTest.kt | 34 +++++++++++ 2 files changed, 65 insertions(+), 25 deletions(-) diff --git a/prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClient.kt b/prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClient.kt index 58d9992664..880c34506c 100644 --- a/prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClient.kt +++ b/prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClient.kt @@ -184,44 +184,44 @@ public class RetryingLLMClient @JvmOverloads constructor( private fun shouldRetry(error: Throwable): Boolean { if (error is IncompleteStreamException) return true - val message = error.message ?: return false + // Also consult the wrapped cause: some clients re-throw as a domain wrapper whose + // own message omits the retry-matching tokens (e.g. "LLM call failed") even though + // the underlying KoogHttpClientException carries the status code and keywords. + val messages = listOfNotNull(error.message, error.unwrapHttpException()?.message) + if (messages.isEmpty()) return false - // Check if error matches any retry pattern return config.retryablePatterns.any { pattern -> - pattern.matches(message) + messages.any { pattern.matches(it) } } } - private fun calculateDelay(attempt: Int, error: Throwable? = null): Duration { - // Check for retry-after hint. Prefer the header-aware overload when the caught - // throwable is a KoogHttpClientException (directly or wrapped one level deep — some - // clients re-throw as LLMClientException with the original exception as cause) so - // structured response metadata is available to the extractor; otherwise fall back - // to the message-based overload. - val extractor = config.retryAfterExtractor - if (extractor != null && error != null) { - val koogException = error as? KoogHttpClientException - ?: error.cause as? KoogHttpClientException - val retryAfter = if (koogException != null) { - extractor.extract(koogException) - } else { - error.message?.let { extractor.extract(it) } - } - if (retryAfter != null) return retryAfter + private fun calculateDelay(attempt: Int, error: Throwable? = null): Duration = + retryAfterHint(error) ?: exponentialBackoffWithJitter(attempt) + + // Prefers the header-aware extractor overload when the throwable carries a + // KoogHttpClientException (directly or wrapped one level deep) so structured response + // metadata is available; falls back to the message-based overload otherwise. + private fun retryAfterHint(error: Throwable?): Duration? { + val extractor = config.retryAfterExtractor ?: return null + if (error == null) return null + val koogException = error.unwrapHttpException() + return if (koogException != null) { + extractor.extract(koogException) + } else { + error.message?.let { extractor.extract(it) } } + } - // Exponential backoff with jitter + private fun exponentialBackoffWithJitter(attempt: Int): Duration { var exponentialMs = config.initialDelay.inWholeMilliseconds.toDouble() repeat(attempt) { exponentialMs *= config.backoffMultiplier } val boundedMs = minOf(exponentialMs, config.maxDelay.inWholeMilliseconds.toDouble()) - - // Add jitter (only increases delay, never decreases) + // Jitter is drawn from [0, bound) so it only increases the delay, never decreases it - + // guards against clients that shorten their backoff under load and stampede retries. val jitterMs = Random.nextDouble(0.0, boundedMs * config.jitterFactor) - val finalMs = (boundedMs + jitterMs).toLong() - - return finalMs.milliseconds + return (boundedMs + jitterMs).toLong().milliseconds } override fun close() { @@ -235,6 +235,12 @@ public class RetryingLLMClient @JvmOverloads constructor( override fun getBasicJsonSchemaGenerator(): BasicJsonSchemaGenerator { return delegate.getBasicJsonSchemaGenerator() } + + // Returns the throwable itself if it is a KoogHttpClientException, or its immediate cause + // when it wraps one. Keeps `shouldRetry` and `calculateDelay` symmetric so a wrapper whose + // own message is non-matching still has its underlying HTTP metadata consulted. + private fun Throwable.unwrapHttpException(): KoogHttpClientException? = + this as? KoogHttpClientException ?: this.cause as? KoogHttpClientException } /** diff --git a/prompt/prompt-executor/prompt-executor-clients/src/commonTest/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClientTest.kt b/prompt/prompt-executor/prompt-executor-clients/src/commonTest/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClientTest.kt index 3d740a8ad1..6d593cd759 100644 --- a/prompt/prompt-executor/prompt-executor-clients/src/commonTest/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClientTest.kt +++ b/prompt/prompt-executor/prompt-executor-clients/src/commonTest/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClientTest.kt @@ -155,6 +155,40 @@ class RetryingLLMClientTest { assertEquals(4_000L, elapsed) } + @Test + fun testRetriesWhenOuterWrapperIsNonRetryableButCauseIs429() = runTest { + // The outer wrapper's own message has no retry-matching tokens, but the wrapped + // KoogHttpClientException is a 429. shouldRetry must walk .cause the same way + // calculateDelay does; otherwise we'd bail before ever consulting the header hint. + val httpError = KoogHttpClientException( + clientName = "TestClient", + statusCode = 429, + errorBody = "rate limited", + message = "rate limited", + headers = mapOf("retry-after" to listOf("2")) + ) + val wrapped = RuntimeException("Unexpected failure in ollama transport layer", httpError) + + val mockClient = MockLLMClient( + executeResponse = testResponse, + failuresBeforeSuccess = 1, + failure = wrapped + ) + + val retryingClient = RetryingLLMClient( + mockClient, + RetryConfig(maxAttempts = 2, initialDelay = 10.milliseconds) + ) + + val startTime = testScheduler.currentTime + val result = retryingClient.execute(testPrompt, testModel, emptyList()) + val elapsed = testScheduler.currentTime - startTime + + assertEquals(testResponse, result) + assertEquals(2, mockClient.executeCalls) + assertEquals(2_000L, elapsed) + } + @Test fun testMessageBasedExtractorStillUsedForNonKoogExceptions() = runTest { // When the throwable isn't a KoogHttpClientException, calculateDelay must continue to From 88d5902af0d15974112e07525026684b51e803f4 Mon Sep 17 00:00:00 2001 From: qflen Date: Thu, 23 Apr 2026 02:22:57 +0200 Subject: [PATCH 03/10] test(http-client,prompt): cover SSE header capture and zero retry-after hint testCaptureHeadersOnNonSuccess only exercised processResponse, leaving the SSE error-before-stream paths in all three clients untested. Extend SSEEndpointConfig so the mock can return 429 + headers before streaming; add testCaptureHeadersOnSseError, overridden per client. Also: retry-after: 10 + x-ratelimit-reset-requests: 0s must resolve to 10s - zero must not clobber a usable hint. (cherry picked from commit 324f775d8bd407e40492d618e53236b38f0130dc) --- .../client/test/BaseKoogHttpClientTest.kt | 37 +++++++++++++++++++ .../ai/koog/http/client/test/MockWebServer.kt | 23 +++++++++++- .../StandardHeaderRetryAfterExtractorTest.kt | 14 +++++++ 3 files changed, 72 insertions(+), 2 deletions(-) diff --git a/http-client/http-client-test/src/main/kotlin/ai/koog/http/client/test/BaseKoogHttpClientTest.kt b/http-client/http-client-test/src/main/kotlin/ai/koog/http/client/test/BaseKoogHttpClientTest.kt index 951bdb105f..954f45fbe2 100644 --- a/http-client/http-client-test/src/main/kotlin/ai/koog/http/client/test/BaseKoogHttpClientTest.kt +++ b/http-client/http-client-test/src/main/kotlin/ai/koog/http/client/test/BaseKoogHttpClientTest.kt @@ -247,6 +247,43 @@ abstract class BaseKoogHttpClientTest { assertEquals(events, collected) } + open fun testCaptureHeadersOnSseError(): Unit = runTest { + mockServer.start( + sseEndpoints = listOf( + MockWebServer.SSEEndpointConfig( + path = "/stream", + events = emptyList(), + statusCode = HttpStatusCode.TooManyRequests, + responseHeaders = mapOf( + "Retry-After" to "5", + "X-RateLimit-Reset-Tokens" to "6m0s" + ) + ) + ) + ) + + val client = createClient() + + val flow = client.sse( + path = mockServer.url("/stream"), + request = "{}", + requestBodyType = String::class, + dataFilter = { it != "[DONE]" }, + decodeStreamingResponse = { it }, + processStreamingChunk = { it } + ) + + try { + flow.toList() + fail("Expected a KoogHttpClientException for 429 SSE response") + } catch (e: KoogHttpClientException) { + assertEquals(429, e.statusCode) + // Header keys are normalized to lowercase regardless of how the server sent them. + assertEquals(listOf("5"), e.headers["retry-after"]) + assertEquals(listOf("6m0s"), e.headers["x-ratelimit-reset-tokens"]) + } + } + @Suppress("FunctionName") open fun `test filter SSE events`(): Unit = runTest { val events = listOf("event1", "[DONE]", "event2", "[DONE]", "event3") diff --git a/http-client/http-client-test/src/main/kotlin/ai/koog/http/client/test/MockWebServer.kt b/http-client/http-client-test/src/main/kotlin/ai/koog/http/client/test/MockWebServer.kt index 491890327b..2eed88f19e 100644 --- a/http-client/http-client-test/src/main/kotlin/ai/koog/http/client/test/MockWebServer.kt +++ b/http-client/http-client-test/src/main/kotlin/ai/koog/http/client/test/MockWebServer.kt @@ -3,6 +3,7 @@ package ai.koog.http.client.test import io.ktor.http.ContentType import io.ktor.http.HttpHeaders import io.ktor.http.HttpStatusCode +import io.ktor.http.isSuccess import io.ktor.server.application.install import io.ktor.server.cio.CIO import io.ktor.server.cio.CIOApplicationEngine @@ -59,11 +60,17 @@ class MockWebServer { ) /** - * Configuration for a mock SSE endpoint + * Configuration for a mock SSE endpoint. + * + * When [statusCode] is a success code the endpoint streams [events]; when it is a non-success + * code the endpoint short-circuits with that status, [responseHeaders], and no SSE body. + * Useful for covering error-before-stream paths such as 429 rate limits. */ data class SSEEndpointConfig( val path: String, - val events: List + val events: List, + val statusCode: HttpStatusCode = HttpStatusCode.OK, + val responseHeaders: Map = emptyMap() ) /** @@ -218,6 +225,18 @@ class MockWebServer { post(config.path) { call.receiveText() + if (!config.statusCode.isSuccess()) { + config.responseHeaders.forEach { (name, value) -> + call.response.header(name, value) + } + call.respondText( + text = "", + contentType = ContentType.Text.Plain, + status = config.statusCode + ) + return@post + } + val handle: suspend ServerSSESession.() -> Unit = { config.events.forEach { event -> send(event) diff --git a/prompt/prompt-executor/prompt-executor-clients/src/commonTest/kotlin/ai/koog/prompt/executor/clients/retry/StandardHeaderRetryAfterExtractorTest.kt b/prompt/prompt-executor/prompt-executor-clients/src/commonTest/kotlin/ai/koog/prompt/executor/clients/retry/StandardHeaderRetryAfterExtractorTest.kt index fc12a9f405..4d5d5b3466 100644 --- a/prompt/prompt-executor/prompt-executor-clients/src/commonTest/kotlin/ai/koog/prompt/executor/clients/retry/StandardHeaderRetryAfterExtractorTest.kt +++ b/prompt/prompt-executor/prompt-executor-clients/src/commonTest/kotlin/ai/koog/prompt/executor/clients/retry/StandardHeaderRetryAfterExtractorTest.kt @@ -105,6 +105,20 @@ class StandardHeaderRetryAfterExtractorTest { assertNull(extractor.extract(error(mapOf("x-ratelimit-reset-tokens" to listOf(""))))) } + @Test + fun testZeroHintSkippedWhenAnotherHintWins() { + // A zero-duration hint must be filtered as "no hint" rather than chosen as the + // minimum, so the next-smallest positive value wins. Otherwise a server returning + // `x-ratelimit-reset-requests: 0s` would clobber a usable retry-after. + val e = error( + mapOf( + "retry-after" to listOf("10"), + "x-ratelimit-reset-requests" to listOf("0s") + ) + ) + assertEquals(10.seconds, extractor.extract(e)) + } + @Test fun testReturnsSmallestNonzeroWhenMultipleHeadersPresent() { val e = error( From dceda27cd2c3ca021ab17aab6fe1e8e7b3f6d2f9 Mon Sep 17 00:00:00 2001 From: qflen Date: Thu, 23 Apr 2026 02:23:07 +0200 Subject: [PATCH 04/10] docs(prompts): polish retry-after section and add Java tab - Rewrite the custom-extractor section: drop the noisy SAM-lambda + .let chaining example in favor of a clean object implementing RetryAfterExtractor. - Add a sentence clarifying that retry-after: 0 (and expired HTTP-dates / zero-length OpenAI durations) are treated as "no hint" rather than "retry immediately" - the caller falls back to exponential backoff. - Restructure the custom-extractor example into Kotlin/Java tabs matching the rest of the page; the Java tab plugs an anonymous RetryAfterExtractor into RetryConfig.builder(). (cherry picked from commit c741197c57fe842ec0d6e48ff14245ce9539ba18) --- docs/docs/prompts/handling-failures.md | 120 +++++++++++++++++-------- 1 file changed, 82 insertions(+), 38 deletions(-) diff --git a/docs/docs/prompts/handling-failures.md b/docs/docs/prompts/handling-failures.md index f16e444a7f..08405d23fd 100644 --- a/docs/docs/prompts/handling-failures.md +++ b/docs/docs/prompts/handling-failures.md @@ -230,37 +230,9 @@ When the underlying HTTP client raises a `KoogHttpClientException`, the exceptio response headers on its `headers: Map>` property. Keys are normalized to lowercase, so extractors can look them up without re-casing. -The default `RetryConfig.retryAfterExtractor` is a composite that consults headers first and -falls back to parsing the error message: - -```kotlin -public val DEFAULT_RETRY_AFTER_EXTRACTOR: RetryAfterExtractor = CompositeRetryAfterExtractor( - StandardHeaderRetryAfterExtractor.DEFAULT, - DefaultRetryAfterExtractor -) -``` - -So out of the box, a 429 carrying `retry-after: 5` is honored without any extra configuration. -If you want to supply a custom header-aware extractor, implement `RetryAfterExtractor` and -override `extract(error: KoogHttpClientException)`: - -```kotlin -val config = RetryConfig( - retryAfterExtractor = RetryAfterExtractor { message -> - // message-based fallback - null - }.let { fallback -> - CompositeRetryAfterExtractor( - object : RetryAfterExtractor { - override fun extract(message: String): Duration? = null - override fun extract(error: KoogHttpClientException): Duration? = - error.headers["x-my-custom-retry-after"]?.firstOrNull()?.toLongOrNull()?.seconds - }, - fallback - ) - } -) -``` +Out of the box, `RetryConfig.retryAfterExtractor` is a `CompositeRetryAfterExtractor` that +consults `StandardHeaderRetryAfterExtractor` first and falls back to `DefaultRetryAfterExtractor`. +So a 429 carrying `retry-after: 5` is honored without any extra configuration. `StandardHeaderRetryAfterExtractor` understands: @@ -270,6 +242,78 @@ val config = RetryConfig( When multiple hints are present it returns the smallest strictly-positive delay so the client doesn't sleep longer than any of the individual rate-limit buckets require. +A literal `0`, an expired HTTP-date, or a zero-length OpenAI-style duration is treated as +"no hint" rather than "retry immediately": the caller falls back to exponential backoff or to +the next extractor in the composite. + +To supply a custom header-aware extractor, implement `RetryAfterExtractor` and override +`extract(error: KoogHttpClientException)`. Existing single-arg SAM lambdas that only inspect +the error message continue to work unchanged. + +=== "Kotlin" + + + ```kotlin + // Header-aware extractor first, message-based fallback second. + val config = RetryConfig( + retryAfterExtractor = CompositeRetryAfterExtractor( + MyHeaderRetryAfterExtractor, + DefaultRetryAfterExtractor, + ) + ) + ``` + + +=== "Java" + + + + ```java + RetryAfterExtractor myExtractor = new RetryAfterExtractor() { + @Override + public Duration extract(String message) { + return null; + } + + @Override + public Duration extract(KoogHttpClientException error) { + List values = error.getHeaders().get("x-my-retry-after"); + if (values == null || values.isEmpty()) return null; + try { + long seconds = Long.parseLong(values.get(0)); + return DurationKt.toDuration(seconds, DurationUnit.SECONDS); + } catch (NumberFormatException ex) { + return null; + } + } + }; + + RetryConfig config = RetryConfig.builder() + .retryAfterExtractor(new CompositeRetryAfterExtractor( + myExtractor, + DefaultRetryAfterExtractor.INSTANCE + )) + .build(); + ``` + ### Streaming with retry @@ -301,7 +345,7 @@ val config = RetryConfig( val client = RetryingLLMClient(baseClient, config) val stream = client.executeStreaming(prompt, OpenAIModels.Chat.GPT4o) ``` - + !!!note Streaming retries only apply to connection failures that occur before the first token is received. @@ -353,7 +397,7 @@ To learn more about prompt executors, see [Prompt executors](prompt-executors.md ), ) ``` - + === "Java" @@ -390,7 +434,7 @@ To learn more about prompt executors, see [Prompt executors](prompt-executors.md MultiLLMPromptExecutor multiExecutor = new MultiLLMPromptExecutor(clients); ``` - + ## Timeout configuration @@ -428,7 +472,7 @@ You can customize these values for your specific needs. For example: ) ) ``` - + === "Java" @@ -456,7 +500,7 @@ You can customize these values for your specific needs. For example: ); OpenAILLMClient client = openAIClient(apiKey, settings); ``` - + !!! tip For long-running or streaming calls, set higher values for `requestTimeoutMillis` and `socketTimeoutMillis`. @@ -525,7 +569,7 @@ Here is an example of error handling in Kotlin and Java: } } ``` - + === "Java" @@ -580,4 +624,4 @@ Here is an example of error handling in Kotlin and Java: } } ``` - + From c12b8fe6901ed621f378c5c36b4031697a576cb7 Mon Sep 17 00:00:00 2001 From: qflen Date: Thu, 23 Apr 2026 02:23:15 +0200 Subject: [PATCH 05/10] refactor(prompt): improve retry-after extractor readability - CompositeRetryAfterExtractor: override toString() to render the consultation chain so retry logs make extractor order obvious. - parseOpenAIDuration: replace the unit-millis Double arithmetic with amount.toDuration(unit) using kotlin.time.DurationUnit. (cherry picked from commit b65f376fc18dbc26a61c32529bd3f66439f0cfd6) --- .../executor/clients/retry/RetryConfig.kt | 22 ++++++++++++------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/retry/RetryConfig.kt b/prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/retry/RetryConfig.kt index 8b740fdf29..573628b671 100644 --- a/prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/retry/RetryConfig.kt +++ b/prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/retry/RetryConfig.kt @@ -7,7 +7,9 @@ import kotlin.time.Clock import kotlin.time.Duration import kotlin.time.Duration.Companion.milliseconds import kotlin.time.Duration.Companion.seconds +import kotlin.time.DurationUnit import kotlin.time.Instant +import kotlin.time.toDuration /** * Configuration for retry behavior in LLM client operations. @@ -310,19 +312,19 @@ public class StandardHeaderRetryAfterExtractor( if (tokens.isEmpty()) return null if (tokens.sumOf { it.value.length } != trimmed.length) return null - var totalMillis = 0.0 + var total = Duration.ZERO for (token in tokens) { val amount = token.groupValues[1].toDoubleOrNull() ?: return null - val unitMillis = when (token.groupValues[2].lowercase()) { - "ms" -> 1.0 - "s" -> 1000.0 - "m" -> 60_000.0 - "h" -> 3_600_000.0 + val unit = when (token.groupValues[2].lowercase()) { + "ms" -> DurationUnit.MILLISECONDS + "s" -> DurationUnit.SECONDS + "m" -> DurationUnit.MINUTES + "h" -> DurationUnit.HOURS else -> return null } - totalMillis += amount * unitMillis + total += amount.toDuration(unit) } - return totalMillis.toLong().milliseconds + return total } /** @@ -369,4 +371,8 @@ public class CompositeRetryAfterExtractor( override fun extract(error: KoogHttpClientException): Duration? = extractors.firstNotNullOfOrNull { it.extract(error) } + + /** Renders the chain so retry logs make the consultation order obvious. */ + override fun toString(): String = + extractors.joinToString(prefix = "CompositeRetryAfterExtractor(", postfix = ")") } From b26aa589ef35a3359a89e72c1a6c708f32877f67 Mon Sep 17 00:00:00 2001 From: qflen Date: Tue, 12 May 2026 00:21:10 +0200 Subject: [PATCH 06/10] docs(prompts): fix Java retry-after example builder call RetryConfig.builder() is a Kotlin-only extension on the Companion and isn't callable from Java; the builder class itself is the Java entry point. Switch the example to `new RetryConfigBuilder()` so it actually compiles when copy-pasted. (cherry picked from commit c2ed5675f41e8019ed76e98d8b543cfe1484baa7) --- docs/docs/prompts/handling-failures.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/docs/prompts/handling-failures.md b/docs/docs/prompts/handling-failures.md index 08405d23fd..a0a0a7f41c 100644 --- a/docs/docs/prompts/handling-failures.md +++ b/docs/docs/prompts/handling-failures.md @@ -306,7 +306,7 @@ the error message continue to work unchanged. } }; - RetryConfig config = RetryConfig.builder() + RetryConfig config = new RetryConfigBuilder() .retryAfterExtractor(new CompositeRetryAfterExtractor( myExtractor, DefaultRetryAfterExtractor.INSTANCE From c196f7b75e56490b83ce19d8a59201e20cc9d3b5 Mon Sep 17 00:00:00 2001 From: qflen Date: Tue, 12 May 2026 00:44:56 +0200 Subject: [PATCH 07/10] fix(prompt): guard jitter draw against zero upper bound in RetryingLLMClient Random.nextDouble requires from < until, so `Random.nextDouble(0.0, 0.0)` throws IllegalArgumentException. The retry path hit this whenever `jitterFactor == 0.0` (documented as a valid "no jitter" setting and accepted by RetryConfig.require) or `initialDelay == Duration.ZERO`, silently converting a retryable error into an unrelated IAE. Skip the draw when the upper bound collapses to zero and add a regression test exercising jitterFactor = 0.0 through the public execute() retry path. (cherry picked from commit bbf7dd4c93702190d57dbc6c6f093bcc7049ceda) --- .../clients/retry/RetryingLLMClient.kt | 5 +++- .../clients/retry/RetryingLLMClientTest.kt | 23 +++++++++++++++++++ 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClient.kt b/prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClient.kt index 880c34506c..6246aa4ea6 100644 --- a/prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClient.kt +++ b/prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClient.kt @@ -220,7 +220,10 @@ public class RetryingLLMClient @JvmOverloads constructor( val boundedMs = minOf(exponentialMs, config.maxDelay.inWholeMilliseconds.toDouble()) // Jitter is drawn from [0, bound) so it only increases the delay, never decreases it - // guards against clients that shorten their backoff under load and stampede retries. - val jitterMs = Random.nextDouble(0.0, boundedMs * config.jitterFactor) + // Random.nextDouble requires from < until, so skip the draw when the upper bound + // collapses to zero (jitterFactor == 0.0 or initialDelay == Duration.ZERO). + val jitterUpper = boundedMs * config.jitterFactor + val jitterMs = if (jitterUpper > 0.0) Random.nextDouble(0.0, jitterUpper) else 0.0 return (boundedMs + jitterMs).toLong().milliseconds } diff --git a/prompt/prompt-executor/prompt-executor-clients/src/commonTest/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClientTest.kt b/prompt/prompt-executor/prompt-executor-clients/src/commonTest/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClientTest.kt index 6d593cd759..5d6ac3b2a7 100644 --- a/prompt/prompt-executor/prompt-executor-clients/src/commonTest/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClientTest.kt +++ b/prompt/prompt-executor/prompt-executor-clients/src/commonTest/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClientTest.kt @@ -408,6 +408,29 @@ class RetryingLLMClientTest { assertEquals(1, mockClient.executeCalls) // No retry on cancellation } + @Test + fun testJitterFactorZeroRetriesWithoutThrowing() = runTest { + // jitterFactor = 0.0 is documented as a valid "no jitter" setting and accepted by + // RetryConfig's require(). Before the fix, calculateDelay called Random.nextDouble(0.0, 0.0) + // and threw IllegalArgumentException instead of retrying, so a transient 503 would surface + // as an unrelated IAE. + val mockClient = MockLLMClient( + executeResponse = testResponse, + failuresBeforeSuccess = 1, + failureMessage = "Error: 503" + ) + + val retryingClient = RetryingLLMClient( + mockClient, + RetryConfig(maxAttempts = 2, initialDelay = 10.milliseconds, jitterFactor = 0.0) + ) + + val result = retryingClient.execute(testPrompt, testModel, emptyList()) + + assertEquals(testResponse, result) + assertEquals(2, mockClient.executeCalls) + } + @Test fun testDisabledRetryConfig() = runTest { val mockClient = MockLLMClient( From e4c2e358a167c7b036218bec3e8d3f3ae6c0637c Mon Sep 17 00:00:00 2001 From: qflen Date: Sat, 6 Jun 2026 01:20:56 +0200 Subject: [PATCH 08/10] fix(prompt): harden header-aware retry-after extraction Address review findings on the header-exposure work: - Normalize header keys once in the KoogHttpClientException constructor instead of at every throw site; clients now pass native header maps. - Capture headers at the lines() streaming throw sites in the Ktor, OkHttp, and Java clients and at the Spring WebClient error chokepoint, so streaming 429s (e.g. Ollama) keep their rate-limit hints; cover the path with testCaptureHeadersOnLinesError in all four client suites. - Treat retry-after as authoritative over x-ratelimit-reset-* instead of taking the global minimum; scan every value of repeated headers, parse comma-folded values, and reject non-positive or non-finite hints. - Make header names configurable on StandardHeaderRetryAfterExtractor, switch its clock to KoogClock, and parse reset durations with Duration.parseOrNull instead of a hand-rolled token grammar. - Cap extracted hints at RetryConfig.maxDelay, walk the full cause chain when unwrapping KoogHttpClientException, and fall back to the outer wrapper's message when header extraction yields nothing. - Guard the OkHttp SSE onFailure body read so a stripped response body cannot skip close(exception) and hang the collector. - Refresh ABI dumps, docs, and tests accordingly. (cherry picked from commit 8669aadba3529d50d13b5c94d2228a213d7f01a1) --- docs/docs/prompts/handling-failures.md | 56 ++++---- .../api/android/http-client-core.api | 12 +- .../api/http-client-core.klib.api | 5 +- .../api/jvm/http-client-core.api | 12 +- .../kotlin/ai/koog/http/client/Exceptions.kt | 21 ++- .../kotlin/ai/koog/http/client/Headers.kt | 6 +- .../kotlin/KoogHttpClientExceptionTest.kt | 47 ++++++ .../http/client/java/JavaKoogHttpClient.kt | 6 +- .../client/java/JavaKoogHttpClientTestBase.kt | 4 + .../http/client/ktor/KtorKoogHttpClient.kt | 6 +- .../client/ktor/KtorKoogHttpClientTestBase.kt | 4 + .../client/okhttp/OkHttpKoogHttpClient.kt | 13 +- .../okhttp/OkHttpKoogHttpClientTestBase.kt | 4 + .../spring/SpringWebClientKoogHttpClient.kt | 3 + .../SpringWebClientKoogHttpClientTestBase.kt | 12 ++ .../http-client-test/api/http-client-test.api | 50 ++++--- .../client/test/BaseKoogHttpClientTest.kt | 30 ++++ .../ai/koog/http/client/test/MockWebServer.kt | 4 + .../api/android/prompt-executor-clients.api | 32 +++++ .../api/jvm/prompt-executor-clients.api | 32 +++++ .../api/prompt-executor-clients.klib.api | 28 ++++ .../executor/clients/retry/RetryConfig.kt | 134 +++++++++++------- .../clients/retry/RetryingLLMClient.kt | 32 +++-- .../clients/retry/RetryingLLMClientTest.kt | 134 ++++++++++++++++++ .../StandardHeaderRetryAfterExtractorTest.kt | 105 +++++++++++--- 25 files changed, 638 insertions(+), 154 deletions(-) create mode 100644 http-client/http-client-core/src/commonTest/kotlin/KoogHttpClientExceptionTest.kt diff --git a/docs/docs/prompts/handling-failures.md b/docs/docs/prompts/handling-failures.md index a0a0a7f41c..6f58da9370 100644 --- a/docs/docs/prompts/handling-failures.md +++ b/docs/docs/prompts/handling-failures.md @@ -227,8 +227,9 @@ to) the error body. OpenAI in particular uses `retry-after`, `x-ratelimit-reset- `x-ratelimit-reset-tokens` to signal when a client should back off. When the underlying HTTP client raises a `KoogHttpClientException`, the exception exposes the -response headers on its `headers: Map>` property. Keys are normalized to -lowercase, so extractors can look them up without re-casing. +response headers on its `headers: Map>` property. The exception constructor +normalizes keys to lowercase, so extractors can look them up without re-casing no matter which +HTTP client implementation produced the error. Out of the box, `RetryConfig.retryAfterExtractor` is a `CompositeRetryAfterExtractor` that consults `StandardHeaderRetryAfterExtractor` first and falls back to `DefaultRetryAfterExtractor`. @@ -237,41 +238,44 @@ So a 429 carrying `retry-after: 5` is honored without any extra configuration. `StandardHeaderRetryAfterExtractor` understands: - `retry-after` as either delta-seconds or an IMF-fixdate (RFC 9110 §10.2.3); -- `x-ratelimit-reset-requests` / `x-ratelimit-reset-tokens` as OpenAI-style durations like - `1s`, `6m0s`, `100ms`. - -When multiple hints are present it returns the smallest strictly-positive delay so the client -doesn't sleep longer than any of the individual rate-limit buckets require. -A literal `0`, an expired HTTP-date, or a zero-length OpenAI-style duration is treated as -"no hint" rather than "retry immediately": the caller falls back to exponential backoff or to -the next extractor in the composite. - -To supply a custom header-aware extractor, implement `RetryAfterExtractor` and override -`extract(error: KoogHttpClientException)`. Existing single-arg SAM lambdas that only inspect -the error message continue to work unchanged. +- `x-ratelimit-reset-requests` / `x-ratelimit-reset-tokens` as Go-style durations like + `1s`, `6m0s`, `100ms` (the format OpenAI uses). + +`retry-after` is authoritative: when it carries a usable delay it wins outright. Only when it is +absent or unusable are the reset headers consulted, taking the smallest strictly-positive value +so the client retries as soon as the first rate-limit bucket refills. Every value of a repeated +header is considered, not just the first. A literal `0`, a negative or expired value, and +anything that fails to parse are all treated as "no hint" rather than "retry immediately": the +caller falls back to exponential backoff or to the next extractor in the composite. Whatever the +hint's source, `RetryingLLMClient` caps the resulting delay at `RetryConfig.maxDelay`, so a +misbehaving server cannot stall the retry loop beyond the configured bound. + +Retry eligibility follows the same principle: `RetryingLLMClient` matches its retryable patterns +against both the thrown error's message and the message of any `KoogHttpClientException` found in +its cause chain, so a provider wrapper whose own message omits the status code still retries when +the underlying HTTP error is transient. + +Provider-specific header names do not require a custom extractor: pass them to +`StandardHeaderRetryAfterExtractor` directly, as shown below. For full control, implement +`RetryAfterExtractor` and override `extract(error: KoogHttpClientException)`. Existing single-arg +SAM lambdas that only inspect the error message continue to work unchanged. === "Kotlin" ```kotlin - // Header-aware extractor first, message-based fallback second. + // Provider-specific header names first, message-based fallback second. val config = RetryConfig( retryAfterExtractor = CompositeRetryAfterExtractor( - MyHeaderRetryAfterExtractor, + StandardHeaderRetryAfterExtractor( + retryAfterHeaders = listOf("x-my-retry-after"), + resetDurationHeaders = listOf("x-my-ratelimit-reset"), + ), DefaultRetryAfterExtractor, ) ) diff --git a/http-client/http-client-core/api/android/http-client-core.api b/http-client/http-client-core/api/android/http-client-core.api index 74fd66d686..4df845b986 100644 --- a/http-client/http-client-core/api/android/http-client-core.api +++ b/http-client/http-client-core/api/android/http-client-core.api @@ -1,3 +1,7 @@ +public final class ai/koog/http/client/HeadersKt { + public static final fun lowercaseHeaderKeys (Ljava/util/Map;)Ljava/util/Map; +} + public final class ai/koog/http/client/HttpClientFactoryResolver { public static final field INSTANCE Lai/koog/http/client/HttpClientFactoryResolver; public final fun resolve ()Lai/koog/http/client/KoogHttpClient$Factory; @@ -44,10 +48,16 @@ public final class ai/koog/http/client/KoogHttpClient$Factory$DefaultImpls { public final class ai/koog/http/client/KoogHttpClientException : java/lang/Exception { public fun ()V + public fun (Ljava/lang/String;)V + public fun (Ljava/lang/String;Ljava/lang/Integer;)V + public fun (Ljava/lang/String;Ljava/lang/Integer;Ljava/lang/String;)V + public fun (Ljava/lang/String;Ljava/lang/Integer;Ljava/lang/String;Ljava/lang/String;)V public fun (Ljava/lang/String;Ljava/lang/Integer;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Throwable;)V - public synthetic fun (Ljava/lang/String;Ljava/lang/Integer;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Throwable;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun (Ljava/lang/String;Ljava/lang/Integer;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Throwable;Ljava/util/Map;)V + public synthetic fun (Ljava/lang/String;Ljava/lang/Integer;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Throwable;Ljava/util/Map;ILkotlin/jvm/internal/DefaultConstructorMarker;)V public final fun getClientName ()Ljava/lang/String; public final fun getErrorBody ()Ljava/lang/String; + public final fun getHeaders ()Ljava/util/Map; public final fun getStatusCode ()Ljava/lang/Integer; } diff --git a/http-client/http-client-core/api/http-client-core.klib.api b/http-client/http-client-core/api/http-client-core.klib.api index 9bc08b46de..48f6c9386a 100644 --- a/http-client/http-client-core/api/http-client-core.klib.api +++ b/http-client/http-client-core/api/http-client-core.klib.api @@ -32,16 +32,19 @@ abstract interface ai.koog.http.client/KoogHttpClient : kotlin/AutoCloseable { / } final class ai.koog.http.client/KoogHttpClientException : kotlin/Exception { // ai.koog.http.client/KoogHttpClientException|null[0] - constructor (kotlin/String? = ..., kotlin/Int? = ..., kotlin/String? = ..., kotlin/String? = ..., kotlin/Throwable? = ...) // ai.koog.http.client/KoogHttpClientException.|(kotlin.String?;kotlin.Int?;kotlin.String?;kotlin.String?;kotlin.Throwable?){}[0] + constructor (kotlin/String? = ..., kotlin/Int? = ..., kotlin/String? = ..., kotlin/String? = ..., kotlin/Throwable? = ..., kotlin.collections/Map> = ...) // ai.koog.http.client/KoogHttpClientException.|(kotlin.String?;kotlin.Int?;kotlin.String?;kotlin.String?;kotlin.Throwable?;kotlin.collections.Map>){}[0] final val clientName // ai.koog.http.client/KoogHttpClientException.clientName|{}clientName[0] final fun (): kotlin/String? // ai.koog.http.client/KoogHttpClientException.clientName.|(){}[0] final val errorBody // ai.koog.http.client/KoogHttpClientException.errorBody|{}errorBody[0] final fun (): kotlin/String? // ai.koog.http.client/KoogHttpClientException.errorBody.|(){}[0] + final val headers // ai.koog.http.client/KoogHttpClientException.headers|{}headers[0] + final fun (): kotlin.collections/Map> // ai.koog.http.client/KoogHttpClientException.headers.|(){}[0] final val statusCode // ai.koog.http.client/KoogHttpClientException.statusCode|{}statusCode[0] final fun (): kotlin/Int? // ai.koog.http.client/KoogHttpClientException.statusCode.|(){}[0] } +final fun (kotlin.collections/Map>).ai.koog.http.client/lowercaseHeaderKeys(): kotlin.collections/Map> // ai.koog.http.client/lowercaseHeaderKeys|lowercaseHeaderKeys@kotlin.collections.Map>(){}[0] final fun ai.koog.http.client/mergeHeaders(kotlin/Array>...): kotlin.collections/Map // ai.koog.http.client/mergeHeaders|mergeHeaders(kotlin.Array>...){}[0] final inline fun <#A: reified kotlin/Any, #B: reified kotlin/Any, #C: kotlin/Any> (ai.koog.http.client/KoogHttpClient).ai.koog.http.client/sse(kotlin/String, #A, noinline kotlin/Function1 = ..., noinline kotlin/Function1, noinline kotlin/Function1<#B, #C?>, kotlin.collections/Map = ..., kotlin.collections/Map = ...): kotlinx.coroutines.flow/Flow<#C> // ai.koog.http.client/sse|sse@ai.koog.http.client.KoogHttpClient(kotlin.String;0:0;kotlin.Function1;kotlin.Function1;kotlin.Function1<0:1,0:2?>;kotlin.collections.Map;kotlin.collections.Map){0§;1§;2§}[0] final inline fun <#A: reified kotlin/Any> (ai.koog.http.client/KoogHttpClient).ai.koog.http.client/lines(kotlin/String, #A, kotlin.collections/Map = ..., kotlin.collections/Map = ...): kotlinx.coroutines.flow/Flow // ai.koog.http.client/lines|lines@ai.koog.http.client.KoogHttpClient(kotlin.String;0:0;kotlin.collections.Map;kotlin.collections.Map){0§}[0] diff --git a/http-client/http-client-core/api/jvm/http-client-core.api b/http-client/http-client-core/api/jvm/http-client-core.api index 74fd66d686..4df845b986 100644 --- a/http-client/http-client-core/api/jvm/http-client-core.api +++ b/http-client/http-client-core/api/jvm/http-client-core.api @@ -1,3 +1,7 @@ +public final class ai/koog/http/client/HeadersKt { + public static final fun lowercaseHeaderKeys (Ljava/util/Map;)Ljava/util/Map; +} + public final class ai/koog/http/client/HttpClientFactoryResolver { public static final field INSTANCE Lai/koog/http/client/HttpClientFactoryResolver; public final fun resolve ()Lai/koog/http/client/KoogHttpClient$Factory; @@ -44,10 +48,16 @@ public final class ai/koog/http/client/KoogHttpClient$Factory$DefaultImpls { public final class ai/koog/http/client/KoogHttpClientException : java/lang/Exception { public fun ()V + public fun (Ljava/lang/String;)V + public fun (Ljava/lang/String;Ljava/lang/Integer;)V + public fun (Ljava/lang/String;Ljava/lang/Integer;Ljava/lang/String;)V + public fun (Ljava/lang/String;Ljava/lang/Integer;Ljava/lang/String;Ljava/lang/String;)V public fun (Ljava/lang/String;Ljava/lang/Integer;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Throwable;)V - public synthetic fun (Ljava/lang/String;Ljava/lang/Integer;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Throwable;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun (Ljava/lang/String;Ljava/lang/Integer;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Throwable;Ljava/util/Map;)V + public synthetic fun (Ljava/lang/String;Ljava/lang/Integer;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Throwable;Ljava/util/Map;ILkotlin/jvm/internal/DefaultConstructorMarker;)V public final fun getClientName ()Ljava/lang/String; public final fun getErrorBody ()Ljava/lang/String; + public final fun getHeaders ()Ljava/util/Map; public final fun getStatusCode ()Ljava/lang/Integer; } diff --git a/http-client/http-client-core/src/commonMain/kotlin/ai/koog/http/client/Exceptions.kt b/http-client/http-client-core/src/commonMain/kotlin/ai/koog/http/client/Exceptions.kt index 1d39322ac5..16b17438e2 100644 --- a/http-client/http-client-core/src/commonMain/kotlin/ai/koog/http/client/Exceptions.kt +++ b/http-client/http-client-core/src/commonMain/kotlin/ai/koog/http/client/Exceptions.kt @@ -8,10 +8,9 @@ import kotlin.jvm.JvmOverloads * @property clientName Name of the HTTP client that produced the error, for log attribution. * @property statusCode HTTP status code returned by the server, or `null` if no response was received. * @property errorBody Raw body of the failed response, or `null` if it could not be read. - * @property headers HTTP response headers captured from the failed response. - * Keys are normalized to lowercase so consumers can look headers up without re-casing - * (see [lowercaseHeaderKeys]); values preserve the order and formatting from the server. - * Defaults to an empty map when no response was available (e.g., connection errors). + * @param headers HTTP response headers captured from the failed response, in whatever casing + * the producer has them. Defaults to an empty map when no response was available + * (e.g., connection errors). */ public class KoogHttpClientException @JvmOverloads constructor( public val clientName: String? = null, @@ -19,7 +18,7 @@ public class KoogHttpClientException @JvmOverloads constructor( public val errorBody: String? = null, message: String? = null, cause: Throwable? = null, - public val headers: Map> = emptyMap() + headers: Map> = emptyMap() ) : Exception( buildString { appendLine("Error from client: ${clientName ?: "unknown client"}") @@ -31,4 +30,14 @@ public class KoogHttpClientException @JvmOverloads constructor( } }, cause -) +) { + /** + * HTTP response headers captured from the failed response. + * + * Keys are normalized to lowercase here, in the constructor, so consumers can look headers + * up without re-casing no matter how the producing client cased them (see + * [lowercaseHeaderKeys]); values preserve the order and formatting from the server. + * Empty when no response was available (e.g., connection errors). + */ + public val headers: Map> = headers.lowercaseHeaderKeys() +} diff --git a/http-client/http-client-core/src/commonMain/kotlin/ai/koog/http/client/Headers.kt b/http-client/http-client-core/src/commonMain/kotlin/ai/koog/http/client/Headers.kt index 8f2a205d07..4c56a6aef0 100644 --- a/http-client/http-client-core/src/commonMain/kotlin/ai/koog/http/client/Headers.kt +++ b/http-client/http-client-core/src/commonMain/kotlin/ai/koog/http/client/Headers.kt @@ -7,9 +7,9 @@ package ai.koog.http.client * a response containing e.g. both `Set-Cookie` and `set-cookie` becomes a single lowercase * entry with both values. Empty maps are returned as [emptyMap] to avoid allocation. * - * This matches the contract documented on [KoogHttpClientException.headers] and lets - * individual HTTP client implementations adapt their native header types without each - * reinventing the normalization. + * [KoogHttpClientException] applies this normalization in its constructor, so HTTP client + * implementations can pass their native header maps through as-is; the helper stays public + * for consumers that need the same normalization elsewhere. */ public fun Map>.lowercaseHeaderKeys(): Map> { if (isEmpty()) return emptyMap() diff --git a/http-client/http-client-core/src/commonTest/kotlin/KoogHttpClientExceptionTest.kt b/http-client/http-client-core/src/commonTest/kotlin/KoogHttpClientExceptionTest.kt new file mode 100644 index 0000000000..d6f199ee63 --- /dev/null +++ b/http-client/http-client-core/src/commonTest/kotlin/KoogHttpClientExceptionTest.kt @@ -0,0 +1,47 @@ +import ai.koog.http.client.KoogHttpClientException +import ai.koog.http.client.lowercaseHeaderKeys +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertTrue + +class KoogHttpClientExceptionTest { + + @Test + fun testConstructorNormalizesHeaderKeysToLowercase() { + val exception = KoogHttpClientException( + clientName = "TestClient", + statusCode = 429, + headers = mapOf( + "Retry-After" to listOf("5"), + "X-RateLimit-Reset-Tokens" to listOf("6m0s") + ) + ) + + assertEquals(listOf("5"), exception.headers["retry-after"]) + assertEquals(listOf("6m0s"), exception.headers["x-ratelimit-reset-tokens"]) + assertEquals(setOf("retry-after", "x-ratelimit-reset-tokens"), exception.headers.keys) + } + + @Test + fun testConstructorMergesKeysDifferingOnlyInCase() { + val exception = KoogHttpClientException( + headers = mapOf( + "Set-Cookie" to listOf("a=1"), + "set-cookie" to listOf("b=2") + ) + ) + + assertEquals(listOf("a=1", "b=2"), exception.headers["set-cookie"]) + } + + @Test + fun testHeadersDefaultToEmpty() { + assertTrue(KoogHttpClientException(clientName = "TestClient").headers.isEmpty()) + } + + @Test + fun testLowercaseHeaderKeysOnEmptyMapAvoidsAllocation() { + val empty: Map> = emptyMap() + assertTrue(empty.lowercaseHeaderKeys().isEmpty()) + } +} 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 716c4fc6a3..aa57bbe34f 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 @@ -2,7 +2,6 @@ package ai.koog.http.client.java import ai.koog.http.client.KoogHttpClient import ai.koog.http.client.KoogHttpClientException -import ai.koog.http.client.lowercaseHeaderKeys import ai.koog.http.client.mergeHeaders import ai.koog.utils.io.SuitableForIO import io.github.oshai.kotlinlogging.KLogger @@ -72,7 +71,7 @@ public class JavaKoogHttpClient internal constructor( clientName = clientName, statusCode = response.statusCode(), errorBody = response.body(), - headers = response.headers().map().lowercaseHeaderKeys(), + headers = response.headers().map(), ) } @@ -198,7 +197,7 @@ public class JavaKoogHttpClient internal constructor( KoogHttpClientException( clientName = clientName, statusCode = response.statusCode(), - headers = response.headers().map().lowercaseHeaderKeys(), + headers = response.headers().map(), ) ) return@callbackFlow @@ -288,6 +287,7 @@ public class JavaKoogHttpClient internal constructor( KoogHttpClientException( clientName = clientName, statusCode = response.statusCode(), + headers = response.headers().map(), ) ) return@launch diff --git a/http-client/http-client-java/src/test/kotlin/ai/koog/http/client/java/JavaKoogHttpClientTestBase.kt b/http-client/http-client-java/src/test/kotlin/ai/koog/http/client/java/JavaKoogHttpClientTestBase.kt index a48ed63348..f673c7708e 100644 --- a/http-client/http-client-java/src/test/kotlin/ai/koog/http/client/java/JavaKoogHttpClientTestBase.kt +++ b/http-client/http-client-java/src/test/kotlin/ai/koog/http/client/java/JavaKoogHttpClientTestBase.kt @@ -78,6 +78,10 @@ abstract class JavaKoogHttpClientTestBase : BaseKoogHttpClientTest() { override fun `test lines surfaces non-2xx as KoogHttpClientException`() = super.`test lines surfaces non-2xx as KoogHttpClientException`() + @Test + override fun testCaptureHeadersOnLinesError() = + super.testCaptureHeadersOnLinesError() + @Test override fun `test lines propagates cancellation`() = super.`test lines propagates cancellation`() diff --git a/http-client/http-client-ktor/src/commonMain/kotlin/ai/koog/http/client/ktor/KtorKoogHttpClient.kt b/http-client/http-client-ktor/src/commonMain/kotlin/ai/koog/http/client/ktor/KtorKoogHttpClient.kt index a5937deb76..d60a42c8f5 100644 --- a/http-client/http-client-ktor/src/commonMain/kotlin/ai/koog/http/client/ktor/KtorKoogHttpClient.kt +++ b/http-client/http-client-ktor/src/commonMain/kotlin/ai/koog/http/client/ktor/KtorKoogHttpClient.kt @@ -2,7 +2,6 @@ package ai.koog.http.client.ktor import ai.koog.http.client.KoogHttpClient import ai.koog.http.client.KoogHttpClientException -import ai.koog.http.client.lowercaseHeaderKeys import ai.koog.http.client.mergeHeaders import ai.koog.utils.io.SuitableForIO import io.github.oshai.kotlinlogging.KLogger @@ -101,6 +100,10 @@ public class KtorKoogHttpClient internal constructor( ) } + // Keys keep Ktor's casing; KoogHttpClientException normalizes them to lowercase itself. + private fun Headers.toResponseHeaderMap(): Map> = + entries().associate { it.key to it.value } + private fun HttpRequestBuilder.applyRequestHeaders(headers: Map) { headers.forEach { (name, value) -> if (name.equals(HttpHeaders.ContentType, ignoreCase = true)) { @@ -256,6 +259,7 @@ public class KtorKoogHttpClient internal constructor( clientName = clientName, statusCode = response.status.value, errorBody = response.bodyAsText(), + headers = response.headers.toResponseHeaderMap(), ) } diff --git a/http-client/http-client-ktor/src/jvmTest/kotlin/ai/koog/http/client/ktor/KtorKoogHttpClientTestBase.kt b/http-client/http-client-ktor/src/jvmTest/kotlin/ai/koog/http/client/ktor/KtorKoogHttpClientTestBase.kt index 8b3c213b12..88d00c9e84 100644 --- a/http-client/http-client-ktor/src/jvmTest/kotlin/ai/koog/http/client/ktor/KtorKoogHttpClientTestBase.kt +++ b/http-client/http-client-ktor/src/jvmTest/kotlin/ai/koog/http/client/ktor/KtorKoogHttpClientTestBase.kt @@ -92,6 +92,10 @@ abstract class KtorKoogHttpClientTestBase : BaseKoogHttpClientTest() { override fun `test lines surfaces non-2xx as KoogHttpClientException`() = super.`test lines surfaces non-2xx as KoogHttpClientException`() + @Test + override fun testCaptureHeadersOnLinesError() = + super.testCaptureHeadersOnLinesError() + @Test override fun `test lines propagates cancellation`() = super.`test lines propagates cancellation`() diff --git a/http-client/http-client-okhttp/src/main/kotlin/ai/koog/http/client/okhttp/OkHttpKoogHttpClient.kt b/http-client/http-client-okhttp/src/main/kotlin/ai/koog/http/client/okhttp/OkHttpKoogHttpClient.kt index fac6fd2f5c..7f80c49bac 100644 --- a/http-client/http-client-okhttp/src/main/kotlin/ai/koog/http/client/okhttp/OkHttpKoogHttpClient.kt +++ b/http-client/http-client-okhttp/src/main/kotlin/ai/koog/http/client/okhttp/OkHttpKoogHttpClient.kt @@ -2,7 +2,6 @@ package ai.koog.http.client.okhttp import ai.koog.http.client.KoogHttpClient import ai.koog.http.client.KoogHttpClientException -import ai.koog.http.client.lowercaseHeaderKeys import ai.koog.http.client.mergeHeaders import ai.koog.utils.io.SuitableForIO import io.github.oshai.kotlinlogging.KLogger @@ -73,7 +72,7 @@ public class OkHttpKoogHttpClient internal constructor( clientName = clientName, statusCode = response.code, errorBody = response.body.string(), - headers = response.headers.toMultimap().lowercaseHeaderKeys(), + headers = response.headers.toMultimap(), ) } @@ -206,13 +205,17 @@ public class OkHttpKoogHttpClient internal constructor( } override fun onFailure(eventSource: EventSource, t: Throwable?, response: Response?) { + // On mid-stream failures okhttp-sse passes a response whose body is stripped and + // throws on read; swallow that so the exception is still built and the flow is + // always closed instead of hanging the collector. Headers are in-memory and safe. + val errorBody = response?.let { r -> runCatching { r.body.string() }.getOrNull() } val exception = KoogHttpClientException( clientName = clientName, statusCode = response?.code, - errorBody = response?.body?.string(), + errorBody = errorBody, message = t?.message, cause = t, - headers = response?.headers?.toMultimap()?.lowercaseHeaderKeys() ?: emptyMap() + headers = response?.headers?.toMultimap() ?: emptyMap() ) logger.error(exception) { exception.message } close(exception) @@ -256,12 +259,14 @@ public class OkHttpKoogHttpClient internal constructor( if (!response.isSuccessful) { val errorBody = response.body.string() + val responseHeaders = response.headers.toMultimap() response.close() close( KoogHttpClientException( clientName = clientName, statusCode = response.code, errorBody = errorBody, + headers = responseHeaders, ) ) return@launch diff --git a/http-client/http-client-okhttp/src/test/kotlin/ai/koog/http/client/okhttp/OkHttpKoogHttpClientTestBase.kt b/http-client/http-client-okhttp/src/test/kotlin/ai/koog/http/client/okhttp/OkHttpKoogHttpClientTestBase.kt index 0a835eca40..88555b1d32 100644 --- a/http-client/http-client-okhttp/src/test/kotlin/ai/koog/http/client/okhttp/OkHttpKoogHttpClientTestBase.kt +++ b/http-client/http-client-okhttp/src/test/kotlin/ai/koog/http/client/okhttp/OkHttpKoogHttpClientTestBase.kt @@ -78,6 +78,10 @@ abstract class OkHttpKoogHttpClientTestBase : BaseKoogHttpClientTest() { override fun `test lines surfaces non-2xx as KoogHttpClientException`() = super.`test lines surfaces non-2xx as KoogHttpClientException`() + @Test + override fun testCaptureHeadersOnLinesError() = + super.testCaptureHeadersOnLinesError() + @Test override fun `test lines propagates cancellation`() = super.`test lines propagates cancellation`() diff --git a/http-client/http-client-spring-webclient/src/main/kotlin/ai/koog/http/client/spring/SpringWebClientKoogHttpClient.kt b/http-client/http-client-spring-webclient/src/main/kotlin/ai/koog/http/client/spring/SpringWebClientKoogHttpClient.kt index c25c6b7c06..89e4210094 100644 --- a/http-client/http-client-spring-webclient/src/main/kotlin/ai/koog/http/client/spring/SpringWebClientKoogHttpClient.kt +++ b/http-client/http-client-spring-webclient/src/main/kotlin/ai/koog/http/client/spring/SpringWebClientKoogHttpClient.kt @@ -334,6 +334,9 @@ public class SpringWebClientKoogHttpClient( clientName = clientName, statusCode = statusCode().value(), errorBody = errorBody, + // HttpHeaders is a MultiValueMap; the exception copies and + // lowercases the keys itself. + headers = headers().asHttpHeaders(), ) private fun Flow.chunkedToLines(): Flow = flow { diff --git a/http-client/http-client-spring-webclient/src/test/kotlin/ai/koog/http/client/spring/SpringWebClientKoogHttpClientTestBase.kt b/http-client/http-client-spring-webclient/src/test/kotlin/ai/koog/http/client/spring/SpringWebClientKoogHttpClientTestBase.kt index 919e91648e..b550692a9b 100644 --- a/http-client/http-client-spring-webclient/src/test/kotlin/ai/koog/http/client/spring/SpringWebClientKoogHttpClientTestBase.kt +++ b/http-client/http-client-spring-webclient/src/test/kotlin/ai/koog/http/client/spring/SpringWebClientKoogHttpClientTestBase.kt @@ -32,10 +32,18 @@ abstract class SpringWebClientKoogHttpClientTestBase : BaseKoogHttpClientTest() override fun `test handle on non-success status`() = super.`test handle on non-success status`() + @Test + override fun testCaptureHeadersOnNonSuccess() = + super.testCaptureHeadersOnNonSuccess() + @Test override fun `test get SSE flow and collect events`() = super.`test get SSE flow and collect events`() + @Test + override fun testCaptureHeadersOnSseError() = + super.testCaptureHeadersOnSseError() + @Test override fun `test filter SSE events`() = super.`test filter SSE events`() @@ -68,6 +76,10 @@ abstract class SpringWebClientKoogHttpClientTestBase : BaseKoogHttpClientTest() override fun `test lines surfaces non-2xx as KoogHttpClientException`() = super.`test lines surfaces non-2xx as KoogHttpClientException`() + @Test + override fun testCaptureHeadersOnLinesError() = + super.testCaptureHeadersOnLinesError() + @Test override fun `test lines propagates cancellation`() = super.`test lines propagates cancellation`() diff --git a/http-client/http-client-test/api/http-client-test.api b/http-client/http-client-test/api/http-client-test.api index 1ca17e1eaa..741f31c190 100644 --- a/http-client/http-client-test/api/http-client-test.api +++ b/http-client/http-client-test/api/http-client-test.api @@ -18,6 +18,9 @@ public abstract class ai/koog/http/client/test/BaseKoogHttpClientTest { public fun test return success string response on get with parameters ()V public fun test return success string response on post ()V public fun test return success string response on post with parameters ()V + public fun testCaptureHeadersOnLinesError ()V + public fun testCaptureHeadersOnNonSuccess ()V + public fun testCaptureHeadersOnSseError ()V } public final class ai/koog/http/client/test/BaseKoogHttpClientTest$TestRequest { @@ -85,39 +88,42 @@ public final class ai/koog/http/client/test/MockWebServer { } public final class ai/koog/http/client/test/MockWebServer$GetEndpointConfig { - public fun (Ljava/lang/String;Ljava/lang/String;Lio/ktor/http/HttpStatusCode;Lio/ktor/http/ContentType;Ljava/util/Map;)V - public synthetic fun (Ljava/lang/String;Ljava/lang/String;Lio/ktor/http/HttpStatusCode;Lio/ktor/http/ContentType;Ljava/util/Map;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun (Ljava/lang/String;Ljava/lang/String;Lio/ktor/http/HttpStatusCode;Lio/ktor/http/ContentType;Ljava/util/Map;Ljava/util/Map;)V + public synthetic fun (Ljava/lang/String;Ljava/lang/String;Lio/ktor/http/HttpStatusCode;Lio/ktor/http/ContentType;Ljava/util/Map;Ljava/util/Map;ILkotlin/jvm/internal/DefaultConstructorMarker;)V public final fun component1 ()Ljava/lang/String; public final fun component2 ()Ljava/lang/String; public final fun component3 ()Lio/ktor/http/HttpStatusCode; public final fun component4 ()Lio/ktor/http/ContentType; public final fun component5 ()Ljava/util/Map; - public final fun copy (Ljava/lang/String;Ljava/lang/String;Lio/ktor/http/HttpStatusCode;Lio/ktor/http/ContentType;Ljava/util/Map;)Lai/koog/http/client/test/MockWebServer$GetEndpointConfig; - public static synthetic fun copy$default (Lai/koog/http/client/test/MockWebServer$GetEndpointConfig;Ljava/lang/String;Ljava/lang/String;Lio/ktor/http/HttpStatusCode;Lio/ktor/http/ContentType;Ljava/util/Map;ILjava/lang/Object;)Lai/koog/http/client/test/MockWebServer$GetEndpointConfig; + public final fun component6 ()Ljava/util/Map; + public final fun copy (Ljava/lang/String;Ljava/lang/String;Lio/ktor/http/HttpStatusCode;Lio/ktor/http/ContentType;Ljava/util/Map;Ljava/util/Map;)Lai/koog/http/client/test/MockWebServer$GetEndpointConfig; + public static synthetic fun copy$default (Lai/koog/http/client/test/MockWebServer$GetEndpointConfig;Ljava/lang/String;Ljava/lang/String;Lio/ktor/http/HttpStatusCode;Lio/ktor/http/ContentType;Ljava/util/Map;Ljava/util/Map;ILjava/lang/Object;)Lai/koog/http/client/test/MockWebServer$GetEndpointConfig; public fun equals (Ljava/lang/Object;)Z public final fun getContentType ()Lio/ktor/http/ContentType; public final fun getExpectedParameters ()Ljava/util/Map; public final fun getPath ()Ljava/lang/String; public final fun getResponseBody ()Ljava/lang/String; + public final fun getResponseHeaders ()Ljava/util/Map; public final fun getStatusCode ()Lio/ktor/http/HttpStatusCode; public fun hashCode ()I public fun toString ()Ljava/lang/String; } public final class ai/koog/http/client/test/MockWebServer$LinesEndpointConfig { - public fun (Ljava/lang/String;Ljava/util/List;Lio/ktor/http/HttpStatusCode;Lio/ktor/http/ContentType;Ljava/util/Map;Ljava/util/Map;JLkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function0;)V - public synthetic fun (Ljava/lang/String;Ljava/util/List;Lio/ktor/http/HttpStatusCode;Lio/ktor/http/ContentType;Ljava/util/Map;Ljava/util/Map;JLkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function0;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun (Ljava/lang/String;Ljava/util/List;Lio/ktor/http/HttpStatusCode;Lio/ktor/http/ContentType;Ljava/util/Map;Ljava/util/Map;Ljava/util/Map;JLkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function0;)V + public synthetic fun (Ljava/lang/String;Ljava/util/List;Lio/ktor/http/HttpStatusCode;Lio/ktor/http/ContentType;Ljava/util/Map;Ljava/util/Map;Ljava/util/Map;JLkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function0;ILkotlin/jvm/internal/DefaultConstructorMarker;)V public final fun component1 ()Ljava/lang/String; + public final fun component10 ()Lkotlin/jvm/functions/Function0; public final fun component2 ()Ljava/util/List; public final fun component3 ()Lio/ktor/http/HttpStatusCode; public final fun component4 ()Lio/ktor/http/ContentType; public final fun component5 ()Ljava/util/Map; public final fun component6 ()Ljava/util/Map; - public final fun component7 ()J - public final fun component8 ()Lkotlin/jvm/functions/Function1; - public final fun component9 ()Lkotlin/jvm/functions/Function0; - public final fun copy (Ljava/lang/String;Ljava/util/List;Lio/ktor/http/HttpStatusCode;Lio/ktor/http/ContentType;Ljava/util/Map;Ljava/util/Map;JLkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function0;)Lai/koog/http/client/test/MockWebServer$LinesEndpointConfig; - public static synthetic fun copy$default (Lai/koog/http/client/test/MockWebServer$LinesEndpointConfig;Ljava/lang/String;Ljava/util/List;Lio/ktor/http/HttpStatusCode;Lio/ktor/http/ContentType;Ljava/util/Map;Ljava/util/Map;JLkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function0;ILjava/lang/Object;)Lai/koog/http/client/test/MockWebServer$LinesEndpointConfig; + public final fun component7 ()Ljava/util/Map; + public final fun component8 ()J + public final fun component9 ()Lkotlin/jvm/functions/Function1; + public final fun copy (Ljava/lang/String;Ljava/util/List;Lio/ktor/http/HttpStatusCode;Lio/ktor/http/ContentType;Ljava/util/Map;Ljava/util/Map;Ljava/util/Map;JLkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function0;)Lai/koog/http/client/test/MockWebServer$LinesEndpointConfig; + public static synthetic fun copy$default (Lai/koog/http/client/test/MockWebServer$LinesEndpointConfig;Ljava/lang/String;Ljava/util/List;Lio/ktor/http/HttpStatusCode;Lio/ktor/http/ContentType;Ljava/util/Map;Ljava/util/Map;Ljava/util/Map;JLkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function0;ILjava/lang/Object;)Lai/koog/http/client/test/MockWebServer$LinesEndpointConfig; public fun equals (Ljava/lang/Object;)Z public final fun getContentType ()Lio/ktor/http/ContentType; public final fun getExpectedHeaders ()Ljava/util/Map; @@ -127,42 +133,50 @@ public final class ai/koog/http/client/test/MockWebServer$LinesEndpointConfig { public final fun getOnLineWritten ()Lkotlin/jvm/functions/Function1; public final fun getOnStreamClosed ()Lkotlin/jvm/functions/Function0; public final fun getPath ()Ljava/lang/String; + public final fun getResponseHeaders ()Ljava/util/Map; public final fun getStatusCode ()Lio/ktor/http/HttpStatusCode; public fun hashCode ()I public fun toString ()Ljava/lang/String; } public final class ai/koog/http/client/test/MockWebServer$PostEndpointConfig { - public fun (Ljava/lang/String;Ljava/lang/String;Lio/ktor/http/HttpStatusCode;Lio/ktor/http/ContentType;Ljava/util/Map;Ljava/util/Map;)V - public synthetic fun (Ljava/lang/String;Ljava/lang/String;Lio/ktor/http/HttpStatusCode;Lio/ktor/http/ContentType;Ljava/util/Map;Ljava/util/Map;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun (Ljava/lang/String;Ljava/lang/String;Lio/ktor/http/HttpStatusCode;Lio/ktor/http/ContentType;Ljava/util/Map;Ljava/util/Map;Ljava/util/Map;)V + public synthetic fun (Ljava/lang/String;Ljava/lang/String;Lio/ktor/http/HttpStatusCode;Lio/ktor/http/ContentType;Ljava/util/Map;Ljava/util/Map;Ljava/util/Map;ILkotlin/jvm/internal/DefaultConstructorMarker;)V public final fun component1 ()Ljava/lang/String; public final fun component2 ()Ljava/lang/String; public final fun component3 ()Lio/ktor/http/HttpStatusCode; public final fun component4 ()Lio/ktor/http/ContentType; public final fun component5 ()Ljava/util/Map; public final fun component6 ()Ljava/util/Map; - public final fun copy (Ljava/lang/String;Ljava/lang/String;Lio/ktor/http/HttpStatusCode;Lio/ktor/http/ContentType;Ljava/util/Map;Ljava/util/Map;)Lai/koog/http/client/test/MockWebServer$PostEndpointConfig; - public static synthetic fun copy$default (Lai/koog/http/client/test/MockWebServer$PostEndpointConfig;Ljava/lang/String;Ljava/lang/String;Lio/ktor/http/HttpStatusCode;Lio/ktor/http/ContentType;Ljava/util/Map;Ljava/util/Map;ILjava/lang/Object;)Lai/koog/http/client/test/MockWebServer$PostEndpointConfig; + public final fun component7 ()Ljava/util/Map; + public final fun copy (Ljava/lang/String;Ljava/lang/String;Lio/ktor/http/HttpStatusCode;Lio/ktor/http/ContentType;Ljava/util/Map;Ljava/util/Map;Ljava/util/Map;)Lai/koog/http/client/test/MockWebServer$PostEndpointConfig; + public static synthetic fun copy$default (Lai/koog/http/client/test/MockWebServer$PostEndpointConfig;Ljava/lang/String;Ljava/lang/String;Lio/ktor/http/HttpStatusCode;Lio/ktor/http/ContentType;Ljava/util/Map;Ljava/util/Map;Ljava/util/Map;ILjava/lang/Object;)Lai/koog/http/client/test/MockWebServer$PostEndpointConfig; public fun equals (Ljava/lang/Object;)Z public final fun getContentType ()Lio/ktor/http/ContentType; public final fun getExpectedHeaders ()Ljava/util/Map; public final fun getExpectedParameters ()Ljava/util/Map; public final fun getPath ()Ljava/lang/String; public final fun getResponseBody ()Ljava/lang/String; + public final fun getResponseHeaders ()Ljava/util/Map; public final fun getStatusCode ()Lio/ktor/http/HttpStatusCode; public fun hashCode ()I public fun toString ()Ljava/lang/String; } public final class ai/koog/http/client/test/MockWebServer$SSEEndpointConfig { - public fun (Ljava/lang/String;Ljava/util/List;)V + public fun (Ljava/lang/String;Ljava/util/List;Lio/ktor/http/HttpStatusCode;Ljava/util/Map;)V + public synthetic fun (Ljava/lang/String;Ljava/util/List;Lio/ktor/http/HttpStatusCode;Ljava/util/Map;ILkotlin/jvm/internal/DefaultConstructorMarker;)V public final fun component1 ()Ljava/lang/String; public final fun component2 ()Ljava/util/List; - public final fun copy (Ljava/lang/String;Ljava/util/List;)Lai/koog/http/client/test/MockWebServer$SSEEndpointConfig; - public static synthetic fun copy$default (Lai/koog/http/client/test/MockWebServer$SSEEndpointConfig;Ljava/lang/String;Ljava/util/List;ILjava/lang/Object;)Lai/koog/http/client/test/MockWebServer$SSEEndpointConfig; + public final fun component3 ()Lio/ktor/http/HttpStatusCode; + public final fun component4 ()Ljava/util/Map; + public final fun copy (Ljava/lang/String;Ljava/util/List;Lio/ktor/http/HttpStatusCode;Ljava/util/Map;)Lai/koog/http/client/test/MockWebServer$SSEEndpointConfig; + public static synthetic fun copy$default (Lai/koog/http/client/test/MockWebServer$SSEEndpointConfig;Ljava/lang/String;Ljava/util/List;Lio/ktor/http/HttpStatusCode;Ljava/util/Map;ILjava/lang/Object;)Lai/koog/http/client/test/MockWebServer$SSEEndpointConfig; public fun equals (Ljava/lang/Object;)Z public final fun getEvents ()Ljava/util/List; public final fun getPath ()Ljava/lang/String; + public final fun getResponseHeaders ()Ljava/util/Map; + public final fun getStatusCode ()Lio/ktor/http/HttpStatusCode; public fun hashCode ()I public fun toString ()Ljava/lang/String; } diff --git a/http-client/http-client-test/src/main/kotlin/ai/koog/http/client/test/BaseKoogHttpClientTest.kt b/http-client/http-client-test/src/main/kotlin/ai/koog/http/client/test/BaseKoogHttpClientTest.kt index 954f45fbe2..80641d4d44 100644 --- a/http-client/http-client-test/src/main/kotlin/ai/koog/http/client/test/BaseKoogHttpClientTest.kt +++ b/http-client/http-client-test/src/main/kotlin/ai/koog/http/client/test/BaseKoogHttpClientTest.kt @@ -459,6 +459,36 @@ abstract class BaseKoogHttpClientTest { assertEquals(400, failure.statusCode) } + open fun testCaptureHeadersOnLinesError(): Unit = runTest { + mockServer.start( + linesEndpoints = listOf( + MockWebServer.LinesEndpointConfig( + path = "/stream", + lines = emptyList(), + statusCode = HttpStatusCode.TooManyRequests, + contentType = ContentType.Text.Plain, + responseHeaders = mapOf( + "Retry-After" to "5", + "X-RateLimit-Reset-Tokens" to "6m0s" + ) + ) + ) + ) + + val client = createClient() + + val failure = assertThrows { + client.lines( + path = mockServer.url("/stream"), + requestBody = "{}" + ).toList() + } + assertEquals(429, failure.statusCode) + // Header keys are normalized to lowercase regardless of how the server sent them. + assertEquals(listOf("5"), failure.headers["retry-after"]) + assertEquals(listOf("6m0s"), failure.headers["x-ratelimit-reset-tokens"]) + } + @Suppress("FunctionName") open fun `test lines propagates cancellation`(): Unit = runTest { // Given: server that emits up to 1000 lines diff --git a/http-client/http-client-test/src/main/kotlin/ai/koog/http/client/test/MockWebServer.kt b/http-client/http-client-test/src/main/kotlin/ai/koog/http/client/test/MockWebServer.kt index 2eed88f19e..e6ad22ba45 100644 --- a/http-client/http-client-test/src/main/kotlin/ai/koog/http/client/test/MockWebServer.kt +++ b/http-client/http-client-test/src/main/kotlin/ai/koog/http/client/test/MockWebServer.kt @@ -83,6 +83,7 @@ class MockWebServer { val contentType: ContentType = ContentType.Application.Json, val expectedParameters: Map = emptyMap(), val expectedHeaders: Map = emptyMap(), + val responseHeaders: Map = emptyMap(), val lineDelayMillis: Long = 10, val onLineWritten: (Int) -> Unit = {}, val onStreamClosed: () -> Unit = {}, @@ -204,6 +205,9 @@ class MockWebServer { } } + config.responseHeaders.forEach { (name, value) -> + call.response.header(name, value) + } call.respondTextWriter(contentType = config.contentType, status = config.statusCode) { try { config.lines.forEachIndexed { index, line -> diff --git a/prompt/prompt-executor/prompt-executor-clients/api/android/prompt-executor-clients.api b/prompt/prompt-executor/prompt-executor-clients/api/android/prompt-executor-clients.api index dcce8e930e..979232bcc2 100644 --- a/prompt/prompt-executor/prompt-executor-clients/api/android/prompt-executor-clients.api +++ b/prompt/prompt-executor/prompt-executor-clients/api/android/prompt-executor-clients.api @@ -88,21 +88,36 @@ public final class ai/koog/prompt/executor/clients/LLModelDefinitionsKt { public static final fun modelsById (Lai/koog/prompt/executor/clients/LLModelDefinitions;)Ljava/util/Map; } +public final class ai/koog/prompt/executor/clients/retry/CompositeRetryAfterExtractor : ai/koog/prompt/executor/clients/retry/RetryAfterExtractor { + public fun (Ljava/util/List;)V + public fun ([Lai/koog/prompt/executor/clients/retry/RetryAfterExtractor;)V + public fun extract-LV8wdWc (Lai/koog/http/client/KoogHttpClientException;)Lkotlin/time/Duration; + public fun extract-LV8wdWc (Ljava/lang/String;)Lkotlin/time/Duration; + public fun toString ()Ljava/lang/String; +} + public final class ai/koog/prompt/executor/clients/retry/DefaultRetryAfterExtractor : ai/koog/prompt/executor/clients/retry/RetryAfterExtractor { public static final field INSTANCE Lai/koog/prompt/executor/clients/retry/DefaultRetryAfterExtractor; + public fun extract-LV8wdWc (Lai/koog/http/client/KoogHttpClientException;)Lkotlin/time/Duration; public fun extract-LV8wdWc (Ljava/lang/String;)Lkotlin/time/Duration; } public abstract interface class ai/koog/prompt/executor/clients/retry/RetryAfterExtractor { + public fun extract-LV8wdWc (Lai/koog/http/client/KoogHttpClientException;)Lkotlin/time/Duration; public abstract fun extract-LV8wdWc (Ljava/lang/String;)Lkotlin/time/Duration; } +public final class ai/koog/prompt/executor/clients/retry/RetryAfterExtractor$DefaultImpls { + public static fun extract-LV8wdWc (Lai/koog/prompt/executor/clients/retry/RetryAfterExtractor;Lai/koog/http/client/KoogHttpClientException;)Lkotlin/time/Duration; +} + public final class ai/koog/prompt/executor/clients/retry/RetryConfig { public static final field AGGRESSIVE Lai/koog/prompt/executor/clients/retry/RetryConfig; public static final field CONSERVATIVE Lai/koog/prompt/executor/clients/retry/RetryConfig; public static final field Companion Lai/koog/prompt/executor/clients/retry/RetryConfig$Companion; public static final field DEFAULT Lai/koog/prompt/executor/clients/retry/RetryConfig; public static final field DEFAULT_PATTERNS Ljava/util/List; + public static final field DEFAULT_RETRY_AFTER_EXTRACTOR Lai/koog/prompt/executor/clients/retry/RetryAfterExtractor; public static final field DISABLED Lai/koog/prompt/executor/clients/retry/RetryConfig; public static final field PRODUCTION Lai/koog/prompt/executor/clients/retry/RetryConfig; public fun ()V @@ -205,6 +220,23 @@ public final class ai/koog/prompt/executor/clients/retry/RetryingLLMClientKt { public static synthetic fun toRetryingClient$default (Lai/koog/prompt/executor/clients/LLMClient;Lai/koog/prompt/executor/clients/retry/RetryConfig;ILjava/lang/Object;)Lai/koog/prompt/executor/clients/retry/RetryingLLMClient; } +public final class ai/koog/prompt/executor/clients/retry/StandardHeaderRetryAfterExtractor : ai/koog/prompt/executor/clients/retry/RetryAfterExtractor { + public static final field Companion Lai/koog/prompt/executor/clients/retry/StandardHeaderRetryAfterExtractor$Companion; + public static final field DEFAULT Lai/koog/prompt/executor/clients/retry/StandardHeaderRetryAfterExtractor; + public static final field DEFAULT_RESET_DURATION_HEADERS Ljava/util/List; + public static final field DEFAULT_RETRY_AFTER_HEADERS Ljava/util/List; + public fun ()V + public fun (Lai/koog/utils/time/KoogClock;)V + public fun (Lai/koog/utils/time/KoogClock;Ljava/util/List;)V + public fun (Lai/koog/utils/time/KoogClock;Ljava/util/List;Ljava/util/List;)V + public synthetic fun (Lai/koog/utils/time/KoogClock;Ljava/util/List;Ljava/util/List;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun extract-LV8wdWc (Lai/koog/http/client/KoogHttpClientException;)Lkotlin/time/Duration; + public fun extract-LV8wdWc (Ljava/lang/String;)Lkotlin/time/Duration; +} + +public final class ai/koog/prompt/executor/clients/retry/StandardHeaderRetryAfterExtractor$Companion { +} + public abstract class ai/koog/prompt/executor/clients/serialization/AdditionalPropertiesFlatteningSerializer : kotlinx/serialization/json/JsonTransformingSerializer { public fun (Lkotlinx/serialization/KSerializer;Ljava/lang/String;)V public synthetic fun (Lkotlinx/serialization/KSerializer;Ljava/lang/String;ILkotlin/jvm/internal/DefaultConstructorMarker;)V diff --git a/prompt/prompt-executor/prompt-executor-clients/api/jvm/prompt-executor-clients.api b/prompt/prompt-executor/prompt-executor-clients/api/jvm/prompt-executor-clients.api index 387003abfa..7d40c84cc8 100644 --- a/prompt/prompt-executor/prompt-executor-clients/api/jvm/prompt-executor-clients.api +++ b/prompt/prompt-executor/prompt-executor-clients/api/jvm/prompt-executor-clients.api @@ -105,21 +105,36 @@ public final class ai/koog/prompt/executor/clients/LLModelDefinitionsKt { public static final fun modelsById (Lai/koog/prompt/executor/clients/LLModelDefinitions;)Ljava/util/Map; } +public final class ai/koog/prompt/executor/clients/retry/CompositeRetryAfterExtractor : ai/koog/prompt/executor/clients/retry/RetryAfterExtractor { + public fun (Ljava/util/List;)V + public fun ([Lai/koog/prompt/executor/clients/retry/RetryAfterExtractor;)V + public fun extract-LV8wdWc (Lai/koog/http/client/KoogHttpClientException;)Lkotlin/time/Duration; + public fun extract-LV8wdWc (Ljava/lang/String;)Lkotlin/time/Duration; + public fun toString ()Ljava/lang/String; +} + public final class ai/koog/prompt/executor/clients/retry/DefaultRetryAfterExtractor : ai/koog/prompt/executor/clients/retry/RetryAfterExtractor { public static final field INSTANCE Lai/koog/prompt/executor/clients/retry/DefaultRetryAfterExtractor; + public fun extract-LV8wdWc (Lai/koog/http/client/KoogHttpClientException;)Lkotlin/time/Duration; public fun extract-LV8wdWc (Ljava/lang/String;)Lkotlin/time/Duration; } public abstract interface class ai/koog/prompt/executor/clients/retry/RetryAfterExtractor { + public fun extract-LV8wdWc (Lai/koog/http/client/KoogHttpClientException;)Lkotlin/time/Duration; public abstract fun extract-LV8wdWc (Ljava/lang/String;)Lkotlin/time/Duration; } +public final class ai/koog/prompt/executor/clients/retry/RetryAfterExtractor$DefaultImpls { + public static fun extract-LV8wdWc (Lai/koog/prompt/executor/clients/retry/RetryAfterExtractor;Lai/koog/http/client/KoogHttpClientException;)Lkotlin/time/Duration; +} + public final class ai/koog/prompt/executor/clients/retry/RetryConfig { public static final field AGGRESSIVE Lai/koog/prompt/executor/clients/retry/RetryConfig; public static final field CONSERVATIVE Lai/koog/prompt/executor/clients/retry/RetryConfig; public static final field Companion Lai/koog/prompt/executor/clients/retry/RetryConfig$Companion; public static final field DEFAULT Lai/koog/prompt/executor/clients/retry/RetryConfig; public static final field DEFAULT_PATTERNS Ljava/util/List; + public static final field DEFAULT_RETRY_AFTER_EXTRACTOR Lai/koog/prompt/executor/clients/retry/RetryAfterExtractor; public static final field DISABLED Lai/koog/prompt/executor/clients/retry/RetryConfig; public static final field PRODUCTION Lai/koog/prompt/executor/clients/retry/RetryConfig; public synthetic fun (IJJDDLjava/util/List;Lai/koog/prompt/executor/clients/retry/RetryAfterExtractor;ILkotlin/jvm/internal/DefaultConstructorMarker;)V @@ -233,6 +248,23 @@ public final class ai/koog/prompt/executor/clients/retry/RetryingLLMClientKt { public static synthetic fun toRetryingClient$default (Lai/koog/prompt/executor/clients/LLMClient;Lai/koog/prompt/executor/clients/retry/RetryConfig;ILjava/lang/Object;)Lai/koog/prompt/executor/clients/retry/RetryingLLMClient; } +public final class ai/koog/prompt/executor/clients/retry/StandardHeaderRetryAfterExtractor : ai/koog/prompt/executor/clients/retry/RetryAfterExtractor { + public static final field Companion Lai/koog/prompt/executor/clients/retry/StandardHeaderRetryAfterExtractor$Companion; + public static final field DEFAULT Lai/koog/prompt/executor/clients/retry/StandardHeaderRetryAfterExtractor; + public static final field DEFAULT_RESET_DURATION_HEADERS Ljava/util/List; + public static final field DEFAULT_RETRY_AFTER_HEADERS Ljava/util/List; + public fun ()V + public fun (Lai/koog/utils/time/KoogClock;)V + public fun (Lai/koog/utils/time/KoogClock;Ljava/util/List;)V + public fun (Lai/koog/utils/time/KoogClock;Ljava/util/List;Ljava/util/List;)V + public synthetic fun (Lai/koog/utils/time/KoogClock;Ljava/util/List;Ljava/util/List;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun extract-LV8wdWc (Lai/koog/http/client/KoogHttpClientException;)Lkotlin/time/Duration; + public fun extract-LV8wdWc (Ljava/lang/String;)Lkotlin/time/Duration; +} + +public final class ai/koog/prompt/executor/clients/retry/StandardHeaderRetryAfterExtractor$Companion { +} + public abstract class ai/koog/prompt/executor/clients/serialization/AdditionalPropertiesFlatteningSerializer : kotlinx/serialization/json/JsonTransformingSerializer { public fun (Lkotlinx/serialization/KSerializer;Ljava/lang/String;)V public synthetic fun (Lkotlinx/serialization/KSerializer;Ljava/lang/String;ILkotlin/jvm/internal/DefaultConstructorMarker;)V diff --git a/prompt/prompt-executor/prompt-executor-clients/api/prompt-executor-clients.klib.api b/prompt/prompt-executor/prompt-executor-clients/api/prompt-executor-clients.klib.api index 6eae666dab..82c702027e 100644 --- a/prompt/prompt-executor/prompt-executor-clients/api/prompt-executor-clients.klib.api +++ b/prompt/prompt-executor/prompt-executor-clients/api/prompt-executor-clients.klib.api @@ -12,6 +12,7 @@ open annotation class ai.koog.prompt.executor.clients/InternalLLMClientApi : kot abstract fun interface ai.koog.prompt.executor.clients.retry/RetryAfterExtractor { // ai.koog.prompt.executor.clients.retry/RetryAfterExtractor|null[0] abstract fun extract(kotlin/String): kotlin.time/Duration? // ai.koog.prompt.executor.clients.retry/RetryAfterExtractor.extract|extract(kotlin.String){}[0] + open fun extract(ai.koog.http.client/KoogHttpClientException): kotlin.time/Duration? // ai.koog.prompt.executor.clients.retry/RetryAfterExtractor.extract|extract(ai.koog.http.client.KoogHttpClientException){}[0] } abstract interface ai.koog.prompt.executor.clients/LLMClientAPI : kotlin/AutoCloseable { // ai.koog.prompt.executor.clients/LLMClientAPI|null[0] @@ -56,6 +57,15 @@ abstract class ai.koog.prompt.executor.clients/LLMEmbeddingProvider : ai.koog.pr constructor () // ai.koog.prompt.executor.clients/LLMEmbeddingProvider.|(){}[0] } +final class ai.koog.prompt.executor.clients.retry/CompositeRetryAfterExtractor : ai.koog.prompt.executor.clients.retry/RetryAfterExtractor { // ai.koog.prompt.executor.clients.retry/CompositeRetryAfterExtractor|null[0] + constructor (kotlin.collections/List) // ai.koog.prompt.executor.clients.retry/CompositeRetryAfterExtractor.|(kotlin.collections.List){}[0] + constructor (kotlin/Array...) // ai.koog.prompt.executor.clients.retry/CompositeRetryAfterExtractor.|(kotlin.Array...){}[0] + + final fun extract(ai.koog.http.client/KoogHttpClientException): kotlin.time/Duration? // ai.koog.prompt.executor.clients.retry/CompositeRetryAfterExtractor.extract|extract(ai.koog.http.client.KoogHttpClientException){}[0] + final fun extract(kotlin/String): kotlin.time/Duration? // ai.koog.prompt.executor.clients.retry/CompositeRetryAfterExtractor.extract|extract(kotlin.String){}[0] + final fun toString(): kotlin/String // ai.koog.prompt.executor.clients.retry/CompositeRetryAfterExtractor.toString|toString(){}[0] +} + final class ai.koog.prompt.executor.clients.retry/RetryConfig { // ai.koog.prompt.executor.clients.retry/RetryConfig|null[0] constructor (kotlin/Int = ..., kotlin.time/Duration = ..., kotlin.time/Duration = ..., kotlin/Double = ..., kotlin/Double = ..., kotlin.collections/List = ..., ai.koog.prompt.executor.clients.retry/RetryAfterExtractor? = ...) // ai.koog.prompt.executor.clients.retry/RetryConfig.|(kotlin.Int;kotlin.time.Duration;kotlin.time.Duration;kotlin.Double;kotlin.Double;kotlin.collections.List;ai.koog.prompt.executor.clients.retry.RetryAfterExtractor?){}[0] @@ -95,6 +105,8 @@ final class ai.koog.prompt.executor.clients.retry/RetryConfig { // ai.koog.promp final fun (): ai.koog.prompt.executor.clients.retry/RetryConfig // ai.koog.prompt.executor.clients.retry/RetryConfig.Companion.DEFAULT.|(){}[0] final val DEFAULT_PATTERNS // ai.koog.prompt.executor.clients.retry/RetryConfig.Companion.DEFAULT_PATTERNS|{}DEFAULT_PATTERNS[0] final fun (): kotlin.collections/List // ai.koog.prompt.executor.clients.retry/RetryConfig.Companion.DEFAULT_PATTERNS.|(){}[0] + final val DEFAULT_RETRY_AFTER_EXTRACTOR // ai.koog.prompt.executor.clients.retry/RetryConfig.Companion.DEFAULT_RETRY_AFTER_EXTRACTOR|{}DEFAULT_RETRY_AFTER_EXTRACTOR[0] + final fun (): ai.koog.prompt.executor.clients.retry/RetryAfterExtractor // ai.koog.prompt.executor.clients.retry/RetryConfig.Companion.DEFAULT_RETRY_AFTER_EXTRACTOR.|(){}[0] final val DISABLED // ai.koog.prompt.executor.clients.retry/RetryConfig.Companion.DISABLED|{}DISABLED[0] final fun (): ai.koog.prompt.executor.clients.retry/RetryConfig // ai.koog.prompt.executor.clients.retry/RetryConfig.Companion.DISABLED.|(){}[0] final val PRODUCTION // ai.koog.prompt.executor.clients.retry/RetryConfig.Companion.PRODUCTION|{}PRODUCTION[0] @@ -118,6 +130,22 @@ final class ai.koog.prompt.executor.clients.retry/RetryingLLMClient : ai.koog.pr final suspend fun moderate(ai.koog.prompt/Prompt, ai.koog.prompt.llm/LLModel): ai.koog.prompt.dsl/ModerationResult // ai.koog.prompt.executor.clients.retry/RetryingLLMClient.moderate|moderate(ai.koog.prompt.Prompt;ai.koog.prompt.llm.LLModel){}[0] } +final class ai.koog.prompt.executor.clients.retry/StandardHeaderRetryAfterExtractor : ai.koog.prompt.executor.clients.retry/RetryAfterExtractor { // ai.koog.prompt.executor.clients.retry/StandardHeaderRetryAfterExtractor|null[0] + constructor (ai.koog.utils.time/KoogClock = ..., kotlin.collections/List = ..., kotlin.collections/List = ...) // ai.koog.prompt.executor.clients.retry/StandardHeaderRetryAfterExtractor.|(ai.koog.utils.time.KoogClock;kotlin.collections.List;kotlin.collections.List){}[0] + + final fun extract(ai.koog.http.client/KoogHttpClientException): kotlin.time/Duration? // ai.koog.prompt.executor.clients.retry/StandardHeaderRetryAfterExtractor.extract|extract(ai.koog.http.client.KoogHttpClientException){}[0] + final fun extract(kotlin/String): kotlin.time/Duration? // ai.koog.prompt.executor.clients.retry/StandardHeaderRetryAfterExtractor.extract|extract(kotlin.String){}[0] + + final object Companion { // ai.koog.prompt.executor.clients.retry/StandardHeaderRetryAfterExtractor.Companion|null[0] + final val DEFAULT // ai.koog.prompt.executor.clients.retry/StandardHeaderRetryAfterExtractor.Companion.DEFAULT|{}DEFAULT[0] + final fun (): ai.koog.prompt.executor.clients.retry/StandardHeaderRetryAfterExtractor // ai.koog.prompt.executor.clients.retry/StandardHeaderRetryAfterExtractor.Companion.DEFAULT.|(){}[0] + final val DEFAULT_RESET_DURATION_HEADERS // ai.koog.prompt.executor.clients.retry/StandardHeaderRetryAfterExtractor.Companion.DEFAULT_RESET_DURATION_HEADERS|{}DEFAULT_RESET_DURATION_HEADERS[0] + final fun (): kotlin.collections/List // ai.koog.prompt.executor.clients.retry/StandardHeaderRetryAfterExtractor.Companion.DEFAULT_RESET_DURATION_HEADERS.|(){}[0] + final val DEFAULT_RETRY_AFTER_HEADERS // ai.koog.prompt.executor.clients.retry/StandardHeaderRetryAfterExtractor.Companion.DEFAULT_RETRY_AFTER_HEADERS|{}DEFAULT_RETRY_AFTER_HEADERS[0] + final fun (): kotlin.collections/List // ai.koog.prompt.executor.clients.retry/StandardHeaderRetryAfterExtractor.Companion.DEFAULT_RETRY_AFTER_HEADERS.|(){}[0] + } +} + final class ai.koog.prompt.executor.clients/ConnectionTimeoutConfig { // ai.koog.prompt.executor.clients/ConnectionTimeoutConfig|null[0] constructor (kotlin/Long = ..., kotlin/Long = ..., kotlin/Long = ...) // ai.koog.prompt.executor.clients/ConnectionTimeoutConfig.|(kotlin.Long;kotlin.Long;kotlin.Long){}[0] diff --git a/prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/retry/RetryConfig.kt b/prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/retry/RetryConfig.kt index 573628b671..4757038bd6 100644 --- a/prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/retry/RetryConfig.kt +++ b/prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/retry/RetryConfig.kt @@ -1,15 +1,13 @@ package ai.koog.prompt.executor.clients.retry import ai.koog.http.client.KoogHttpClientException +import ai.koog.utils.time.KoogClock import kotlin.jvm.JvmField import kotlin.jvm.JvmOverloads -import kotlin.time.Clock import kotlin.time.Duration import kotlin.time.Duration.Companion.milliseconds import kotlin.time.Duration.Companion.seconds -import kotlin.time.DurationUnit import kotlin.time.Instant -import kotlin.time.toDuration /** * Configuration for retry behavior in LLM client operations. @@ -243,52 +241,81 @@ public object DefaultRetryAfterExtractor : RetryAfterExtractor { * Extractor that reads structured rate-limit hints from HTTP response headers attached to a * [KoogHttpClientException]. * - * Recognized headers (case-insensitive; keys on [KoogHttpClientException.headers] are - * already normalized to lowercase): + * Recognized headers (keys on [KoogHttpClientException.headers] are already normalized to + * lowercase; configured names are lowercased on construction): * - * - `retry-after` — integer delta-seconds or an IMF-fixdate (RFC 9110 §10.2.3). Non-integer, - * non-date values are ignored. - * - `x-ratelimit-reset-requests`, `x-ratelimit-reset-tokens` — OpenAI-style durations such as - * `1s`, `6m0s`, or `100ms` (RFC 9110 does not standardize these; the format is OpenAI's). + * - [retryAfterHeaders] (default `retry-after`) - integer delta-seconds or an IMF-fixdate + * (RFC 9110 §10.2.3). Non-integer, non-date values are ignored. + * - [resetDurationHeaders] (default `x-ratelimit-reset-requests`, `x-ratelimit-reset-tokens`) - + * Go-style durations such as `1s`, `6m0s`, or `100ms`, the format OpenAI uses (RFC 9110 does + * not standardize these). * - * When multiple hints are present, the smallest strictly-positive delay wins; otherwise the - * extractor returns `null` so the caller can fall back to exponential backoff or a secondary - * extractor (see [CompositeRetryAfterExtractor]). + * `retry-after` is authoritative (RFC 9110), so any usable value from [retryAfterHeaders] wins + * outright; only when none is present are [resetDurationHeaders] consulted, taking the smallest + * strictly-positive value so the client retries as soon as the first bucket refills. Every value + * of a repeated header is considered, not just the first. Non-positive, non-finite, and + * unparseable values count as "no hint", so the caller can fall back to exponential backoff or + * a secondary extractor (see [CompositeRetryAfterExtractor]). * * This extractor only consults headers; [extract] with a bare message always returns `null`. * Compose it with [DefaultRetryAfterExtractor] to retain message-based fallback behavior. * * @property clock Clock used to evaluate HTTP-date values relative to "now". Override in tests. + * @param retryAfterHeaders Names of headers carrying an authoritative delta-seconds or HTTP-date + * delay, consulted first. Defaults to [DEFAULT_RETRY_AFTER_HEADERS]. + * @param resetDurationHeaders Names of headers carrying informational Go-style reset durations, + * consulted only when no retry-after hint is usable. Defaults to + * [DEFAULT_RESET_DURATION_HEADERS]. */ -public class StandardHeaderRetryAfterExtractor( - private val clock: Clock = Clock.System +public class StandardHeaderRetryAfterExtractor @JvmOverloads constructor( + private val clock: KoogClock = KoogClock.System, + retryAfterHeaders: List = DEFAULT_RETRY_AFTER_HEADERS, + resetDurationHeaders: List = DEFAULT_RESET_DURATION_HEADERS, ) : RetryAfterExtractor { + private val retryAfterHeaders: List = retryAfterHeaders.map { it.lowercase() } + private val resetDurationHeaders: List = resetDurationHeaders.map { it.lowercase() } + override fun extract(message: String): Duration? = null override fun extract(error: KoogHttpClientException): Duration? { val headers = error.headers if (headers.isEmpty()) return null - val candidates = buildList { - headers.firstValue("retry-after")?.let { parseRetryAfter(it)?.let(::add) } - headers.firstValue("x-ratelimit-reset-requests")?.let { parseOpenAIDuration(it)?.let(::add) } - headers.firstValue("x-ratelimit-reset-tokens")?.let { parseOpenAIDuration(it)?.let(::add) } - } + // retry-after is authoritative per RFC 9110, so it is honored outright; the reset + // durations are informational and only consulted when no retry-after value is usable. + retryAfterHeaders + .flatMap { name -> headers.usableHints(name, ::parseRetryAfter) } + .minOrNull() + ?.let { return it } - return candidates.filter { it > Duration.ZERO }.minOrNull() + return resetDurationHeaders + .flatMap { name -> headers.usableHints(name, ::parseResetDuration) } + .minOrNull() } - private fun Map>.firstValue(name: String): String? = - this[name]?.firstOrNull()?.trim()?.takeIf { it.isNotEmpty() } + // Considers every value of a (possibly repeated or case-merged) header, keeping only usable + // hints: parseable, strictly positive, and finite. Zero, negative, expired-date, and + // saturated-to-infinite values are all "no hint" rather than "retry immediately/never". + private fun Map>.usableHints( + name: String, + parse: (String) -> Duration? + ): List = this[name].orEmpty().mapNotNull { value -> + parse(value.trim())?.takeIf { it > Duration.ZERO && it.isFinite() } + } private fun parseRetryAfter(value: String): Duration? { - val trimmed = value.trim() - trimmed.toLongOrNull()?.let { return if (it >= 0) it.seconds else null } - val target = parseHttpDate(trimmed) ?: return null - val now = clock.now() - val delta = target - now - return if (delta > Duration.ZERO) delta else Duration.ZERO + if (value.isEmpty()) return null + value.toLongOrNull()?.let { return it.seconds } + parseHttpDate(value)?.let { return it - clock.now() } + // A proxy may fold repeated delta-seconds into one comma-separated value. An IMF-fixdate + // also contains a comma but was matched whole above, so splitting cannot corrupt it. + if ("," in value) { + return value.split(',') + .mapNotNull { part -> part.trim().toLongOrNull()?.seconds?.takeIf { it > Duration.ZERO } } + .minOrNull() + } + return null } /** @@ -305,33 +332,34 @@ public class StandardHeaderRetryAfterExtractor( return runCatching { Instant.parse("$year-$month-${day}T${time}Z") }.getOrNull() } - private fun parseOpenAIDuration(value: String): Duration? { - val trimmed = value.trim() - if (trimmed.isEmpty()) return null - val tokens = DURATION_TOKEN.findAll(trimmed).toList() - if (tokens.isEmpty()) return null - if (tokens.sumOf { it.value.length } != trimmed.length) return null - - var total = Duration.ZERO - for (token in tokens) { - val amount = token.groupValues[1].toDoubleOrNull() ?: return null - val unit = when (token.groupValues[2].lowercase()) { - "ms" -> DurationUnit.MILLISECONDS - "s" -> DurationUnit.SECONDS - "m" -> DurationUnit.MINUTES - "h" -> DurationUnit.HOURS - else -> return null - } - total += amount.toDuration(unit) - } - return total - } + /** + * Parses a Go-style duration (`1s`, `6m0s`, `100ms`, `7.66s`) the way OpenAI formats its + * reset headers. [Duration.parseOrNull] accepts exactly this grammar (plus harmless + * supersets such as `1d`); non-finite results like `Infinity` are discarded by the + * caller's finiteness guard. + */ + private fun parseResetDuration(value: String): Duration? = Duration.parseOrNull(value) /** - * Companion exposing a shared default instance backed by [Clock.System]. + * Companion exposing the default header names and a shared default instance backed by + * [KoogClock.System]. */ public companion object { - /** Shared instance using the system clock. */ + // Note: DEFAULT must stay declared below the header-name lists; companion properties + // initialize top to bottom and its constructor defaults read both lists. + + /** Default names of authoritative delay headers: the standard `retry-after` (RFC 9110). */ + @JvmField + public val DEFAULT_RETRY_AFTER_HEADERS: List = listOf("retry-after") + + /** Default names of informational reset-duration headers (OpenAI's rate-limit buckets). */ + @JvmField + public val DEFAULT_RESET_DURATION_HEADERS: List = listOf( + "x-ratelimit-reset-requests", + "x-ratelimit-reset-tokens", + ) + + /** Shared instance using the system clock and the default header names. */ @JvmField public val DEFAULT: StandardHeaderRetryAfterExtractor = StandardHeaderRetryAfterExtractor() @@ -348,14 +376,12 @@ public class StandardHeaderRetryAfterExtractor( "May" to "05", "Jun" to "06", "Jul" to "07", "Aug" to "08", "Sep" to "09", "Oct" to "10", "Nov" to "11", "Dec" to "12" ) - - private val DURATION_TOKEN = Regex("(\\d+(?:\\.\\d+)?)(ms|h|m|s)", RegexOption.IGNORE_CASE) } } /** * Runs a list of [RetryAfterExtractor]s in order and returns the first non-null result. Useful - * for stacking a header-aware extractor over a message-based fallback — the ordering models + * for stacking a header-aware extractor over a message-based fallback: the ordering models * the idea that structured hints are preferred when present. * * @param extractors Extractors to consult, in order. diff --git a/prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClient.kt b/prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClient.kt index 6246aa4ea6..fd20e7f800 100644 --- a/prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClient.kt +++ b/prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClient.kt @@ -56,6 +56,9 @@ public class RetryingLLMClient @JvmOverloads constructor( private companion object { private val logger = KotlinLogging.logger { } + + // Guards the cause-chain walk against pathological self-referential cycles. + private const val MAX_CAUSE_CHAIN_DEPTH = 10 } override suspend fun execute( @@ -188,28 +191,28 @@ public class RetryingLLMClient @JvmOverloads constructor( // own message omits the retry-matching tokens (e.g. "LLM call failed") even though // the underlying KoogHttpClientException carries the status code and keywords. val messages = listOfNotNull(error.message, error.unwrapHttpException()?.message) - if (messages.isEmpty()) return false return config.retryablePatterns.any { pattern -> messages.any { pattern.matches(it) } } } + // Server hints are capped at maxDelay so a (possibly misbehaving) server cannot stall the + // retry loop beyond the configured bound; the exponential path enforces the same cap itself. private fun calculateDelay(attempt: Int, error: Throwable? = null): Duration = - retryAfterHint(error) ?: exponentialBackoffWithJitter(attempt) + retryAfterHint(error)?.coerceAtMost(config.maxDelay) ?: exponentialBackoffWithJitter(attempt) // Prefers the header-aware extractor overload when the throwable carries a - // KoogHttpClientException (directly or wrapped one level deep) so structured response - // metadata is available; falls back to the message-based overload otherwise. + // KoogHttpClientException anywhere in its cause chain, then falls back to the message-based + // overload on the outer error so hints present only in a wrapper's own message (the + // pre-header behavior) are still honored. private fun retryAfterHint(error: Throwable?): Duration? { val extractor = config.retryAfterExtractor ?: return null if (error == null) return null - val koogException = error.unwrapHttpException() - return if (koogException != null) { - extractor.extract(koogException) - } else { - error.message?.let { extractor.extract(it) } + error.unwrapHttpException()?.let { koogException -> + extractor.extract(koogException)?.let { return it } } + return error.message?.let { extractor.extract(it) } } private fun exponentialBackoffWithJitter(attempt: Int): Duration { @@ -239,11 +242,14 @@ public class RetryingLLMClient @JvmOverloads constructor( return delegate.getBasicJsonSchemaGenerator() } - // Returns the throwable itself if it is a KoogHttpClientException, or its immediate cause - // when it wraps one. Keeps `shouldRetry` and `calculateDelay` symmetric so a wrapper whose - // own message is non-matching still has its underlying HTTP metadata consulted. + // Returns the first KoogHttpClientException in the cause chain (starting with the throwable + // itself), so HTTP metadata is found no matter how many wrapper layers a client adds. Used + // by both `shouldRetry` and `calculateDelay` to keep their views of the error symmetric. private fun Throwable.unwrapHttpException(): KoogHttpClientException? = - this as? KoogHttpClientException ?: this.cause as? KoogHttpClientException + generateSequence(this) { it.cause } + .take(MAX_CAUSE_CHAIN_DEPTH) + .filterIsInstance() + .firstOrNull() } /** diff --git a/prompt/prompt-executor/prompt-executor-clients/src/commonTest/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClientTest.kt b/prompt/prompt-executor/prompt-executor-clients/src/commonTest/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClientTest.kt index 5d6ac3b2a7..de7699aaff 100644 --- a/prompt/prompt-executor/prompt-executor-clients/src/commonTest/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClientTest.kt +++ b/prompt/prompt-executor/prompt-executor-clients/src/commonTest/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClientTest.kt @@ -214,6 +214,140 @@ class RetryingLLMClientTest { assertEquals(3_000L, elapsed) } + @Test + fun testRetryAfterHintIsCappedAtMaxDelay() = runTest { + // A server-controlled hint must never exceed the user's configured ceiling, otherwise + // a misbehaving server could stall the retry loop arbitrarily long (or forever). + val error = KoogHttpClientException( + clientName = "TestClient", + statusCode = 429, + errorBody = "rate limited", + message = "rate limited", + headers = mapOf("retry-after" to listOf("3600")) + ) + + val mockClient = MockLLMClient( + executeResponse = testResponse, + failuresBeforeSuccess = 1, + failure = error + ) + + val retryingClient = RetryingLLMClient( + mockClient, + RetryConfig(maxAttempts = 2, initialDelay = 10.milliseconds, maxDelay = 2.seconds) + ) + + val startTime = testScheduler.currentTime + val result = retryingClient.execute(testPrompt, testModel, emptyList()) + val elapsed = testScheduler.currentTime - startTime + + assertEquals(testResponse, result) + assertEquals(2, mockClient.executeCalls) + assertEquals(2_000L, elapsed) + } + + @Test + fun testHugeResetHeaderFallsBackToExponentialBackoff() = runTest { + // A reset value that saturates to Duration.INFINITE must be discarded as a hint so the + // client backs off normally instead of sleeping forever. + val error = KoogHttpClientException( + clientName = "TestClient", + statusCode = 429, + errorBody = "rate limited", + message = "rate limited", + headers = mapOf("x-ratelimit-reset-tokens" to listOf("99999999999999999999999999s")) + ) + + val mockClient = MockLLMClient( + executeResponse = testResponse, + failuresBeforeSuccess = 1, + failure = error + ) + + val retryingClient = RetryingLLMClient( + mockClient, + RetryConfig(maxAttempts = 2, initialDelay = 50.milliseconds, jitterFactor = 0.0) + ) + + val startTime = testScheduler.currentTime + val result = retryingClient.execute(testPrompt, testModel, emptyList()) + val elapsed = testScheduler.currentTime - startTime + + assertEquals(testResponse, result) + assertEquals(2, mockClient.executeCalls) + assertEquals(50L, elapsed) + } + + @Test + fun testOuterWrapperMessageHintUsedWhenInnerHasNone() = runTest { + // The wrapper's own message carries the only parseable hint; the wrapped + // KoogHttpClientException has no headers and no hint in its message. The retry layer + // must fall back to the outer message (the pre-header behavior) instead of silently + // using exponential backoff. + val inner = KoogHttpClientException( + clientName = "TestClient", + statusCode = 429, + errorBody = null, + message = "too many requests", + headers = emptyMap() + ) + val wrapped = RuntimeException("Error 429 rate limited, retry-after: 6", inner) + + val mockClient = MockLLMClient( + executeResponse = testResponse, + failuresBeforeSuccess = 1, + failure = wrapped + ) + + val retryingClient = RetryingLLMClient( + mockClient, + RetryConfig(maxAttempts = 2, initialDelay = 10.milliseconds) + ) + + val startTime = testScheduler.currentTime + val result = retryingClient.execute(testPrompt, testModel, emptyList()) + val elapsed = testScheduler.currentTime - startTime + + assertEquals(testResponse, result) + assertEquals(2, mockClient.executeCalls) + assertEquals(6_000L, elapsed) + } + + @Test + fun testUnwrapsKoogHttpClientExceptionNestedSeveralLevelsDeep() = runTest { + // Header metadata must be found regardless of how many wrapper layers a client adds. + val httpError = KoogHttpClientException( + clientName = "TestClient", + statusCode = 429, + errorBody = "rate limited", + message = "rate limited", + headers = mapOf("retry-after" to listOf("3")) + ) + val wrapped = RuntimeException( + "Outer wrapper without retry tokens", + IllegalStateException("Intermediate wrapper", httpError) + ) + + val mockClient = MockLLMClient( + executeResponse = testResponse, + failuresBeforeSuccess = 1, + failure = wrapped + ) + + val retryingClient = RetryingLLMClient( + mockClient, + RetryConfig(maxAttempts = 2, initialDelay = 10.milliseconds) + ) + + val startTime = testScheduler.currentTime + val result = retryingClient.execute(testPrompt, testModel, emptyList()) + val elapsed = testScheduler.currentTime - startTime + + assertEquals(testResponse, result) + assertEquals(2, mockClient.executeCalls) + assertEquals(3_000L, elapsed) + } + @Test fun testSucceedOnFirstAttempt() = runTest { val mockClient = MockLLMClient( diff --git a/prompt/prompt-executor/prompt-executor-clients/src/commonTest/kotlin/ai/koog/prompt/executor/clients/retry/StandardHeaderRetryAfterExtractorTest.kt b/prompt/prompt-executor/prompt-executor-clients/src/commonTest/kotlin/ai/koog/prompt/executor/clients/retry/StandardHeaderRetryAfterExtractorTest.kt index 4d5d5b3466..71b6bbfa8d 100644 --- a/prompt/prompt-executor/prompt-executor-clients/src/commonTest/kotlin/ai/koog/prompt/executor/clients/retry/StandardHeaderRetryAfterExtractorTest.kt +++ b/prompt/prompt-executor/prompt-executor-clients/src/commonTest/kotlin/ai/koog/prompt/executor/clients/retry/StandardHeaderRetryAfterExtractorTest.kt @@ -1,10 +1,10 @@ package ai.koog.prompt.executor.clients.retry import ai.koog.http.client.KoogHttpClientException +import ai.koog.utils.time.KoogClock import kotlin.test.Test import kotlin.test.assertEquals import kotlin.test.assertNull -import kotlin.time.Clock import kotlin.time.Duration.Companion.milliseconds import kotlin.time.Duration.Companion.seconds import kotlin.time.Instant @@ -42,7 +42,7 @@ class StandardHeaderRetryAfterExtractorTest { @Test fun testIgnoresRetryAfterZeroAsNonCandidate() { - // Zero is not a "nonzero" candidate per the extractor contract; no other hints → null. + // Zero means "no usable hint", not "retry immediately"; no other hints -> null. assertNull(extractor.extract(error(mapOf("retry-after" to listOf("0"))))) } @@ -57,26 +57,47 @@ class StandardHeaderRetryAfterExtractorTest { assertNull(extractor.extract(error(mapOf("retry-after" to listOf(""))))) } + @Test + fun testIgnoresHugeRetryAfterThatSaturatesToInfinite() { + // Long.MAX_VALUE seconds saturates to Duration.INFINITE; an infinite "hint" must be + // discarded instead of stalling the retry loop forever. + assertNull(extractor.extract(error(mapOf("retry-after" to listOf("9223372036854775807"))))) + } + @Test fun testExtractsHttpDateRetryAfter() { // Reference "now" = 2026-01-01T00:00:00Z; server says retry at 2026-01-01T00:00:42Z. val now = Instant.parse("2026-01-01T00:00:00Z") - val fixedExtractor = StandardHeaderRetryAfterExtractor(clock = fixedClockAt(now)) + val fixedExtractor = StandardHeaderRetryAfterExtractor(clock = KoogClock { now }) val e = error(mapOf("retry-after" to listOf("Thu, 01 Jan 2026 00:00:42 GMT"))) assertEquals(42.seconds, fixedExtractor.extract(e)) } @Test - fun testHttpDateRetryAfterInPastClampsToZeroAndIsIgnored() { + fun testHttpDateRetryAfterInPastIsIgnored() { val now = Instant.parse("2026-01-01T00:01:00Z") - val fixedExtractor = StandardHeaderRetryAfterExtractor(clock = fixedClockAt(now)) + val fixedExtractor = StandardHeaderRetryAfterExtractor(clock = KoogClock { now }) val e = error(mapOf("retry-after" to listOf("Thu, 01 Jan 2026 00:00:00 GMT"))) - // Past date → Duration.ZERO → filtered out as non-positive → null. + // Past date -> negative delta -> filtered out as non-positive -> null. assertNull(fixedExtractor.extract(e)) } + @Test + fun testConsidersAllValuesOfRepeatedHeader() { + // A merged/repeated header may carry an unusable first value; later usable values + // must still be consulted instead of being masked by the first. + assertEquals(5.seconds, extractor.extract(error(mapOf("retry-after" to listOf("", "5"))))) + assertEquals(30.seconds, extractor.extract(error(mapOf("retry-after" to listOf("0", "30"))))) + } + + @Test + fun testParsesCommaFoldedRetryAfter() { + // Some proxies fold repeated headers into one comma-separated value. + assertEquals(5.seconds, extractor.extract(error(mapOf("retry-after" to listOf("5, 10"))))) + } + @Test fun testExtractsOpenAIStyleResetRequests() { assertEquals(1.seconds, extractor.extract(error(mapOf("x-ratelimit-reset-requests" to listOf("1s"))))) @@ -99,45 +120,87 @@ class StandardHeaderRetryAfterExtractorTest { @Test fun testRejectsMalformedOpenAIDuration() { - // Extra characters that don't match any token should reject the whole value. + // Values that are not a well-formed duration must reject as a whole. assertNull(extractor.extract(error(mapOf("x-ratelimit-reset-tokens" to listOf("6minutes"))))) assertNull(extractor.extract(error(mapOf("x-ratelimit-reset-tokens" to listOf("1s foo"))))) assertNull(extractor.extract(error(mapOf("x-ratelimit-reset-tokens" to listOf(""))))) } @Test - fun testZeroHintSkippedWhenAnotherHintWins() { - // A zero-duration hint must be filtered as "no hint" rather than chosen as the - // minimum, so the next-smallest positive value wins. Otherwise a server returning - // `x-ratelimit-reset-requests: 0s` would clobber a usable retry-after. + fun testRejectsNonFiniteResetDuration() { + // Duration.parseOrNull accepts "Infinity"; the extractor must not surface it as a hint. + assertNull(extractor.extract(error(mapOf("x-ratelimit-reset-tokens" to listOf("Infinity"))))) + assertNull( + extractor.extract(error(mapOf("x-ratelimit-reset-tokens" to listOf("99999999999999999999999999s")))) + ) + } + + @Test + fun testRetryAfterPreferredOverResetDurations() { + // retry-after is authoritative (RFC 9110): even a much smaller informational bucket + // reset must not undercut it and trigger an immediate re-throttled retry. val e = error( mapOf( "retry-after" to listOf("10"), - "x-ratelimit-reset-requests" to listOf("0s") + "x-ratelimit-reset-requests" to listOf("500ms"), + "x-ratelimit-reset-tokens" to listOf("6m0s") ) ) assertEquals(10.seconds, extractor.extract(e)) } @Test - fun testReturnsSmallestNonzeroWhenMultipleHeadersPresent() { + fun testZeroResetDurationDoesNotClobberRetryAfter() { val e = error( mapOf( "retry-after" to listOf("10"), + "x-ratelimit-reset-requests" to listOf("0s") + ) + ) + assertEquals(10.seconds, extractor.extract(e)) + } + + @Test + fun testSmallestResetDurationWinsWhenRetryAfterUnusable() { + // With retry-after absent or unusable, the smallest positive bucket reset is the + // earliest moment a retry can succeed. + val withoutRetryAfter = error( + mapOf( "x-ratelimit-reset-requests" to listOf("500ms"), "x-ratelimit-reset-tokens" to listOf("6m0s") ) ) - assertEquals(500.milliseconds, extractor.extract(e)) + assertEquals(500.milliseconds, extractor.extract(withoutRetryAfter)) + + val withZeroRetryAfter = error( + mapOf( + "retry-after" to listOf("0"), + "x-ratelimit-reset-tokens" to listOf("2s") + ) + ) + assertEquals(2.seconds, extractor.extract(withZeroRetryAfter)) } @Test - fun testIgnoresHeadersWithUppercaseKeysFromCaller() { - // Extractor looks headers up by lowercase key. KoogHttpClientException.headers is - // contractually lowercase; if a caller builds an exception with uppercase keys by - // mistake, the extractor simply finds nothing. + fun testNormalizesUppercaseKeysFromCaller() { + // KoogHttpClientException lowercases header keys in its constructor, so even an + // exception built directly with native-cased keys works with the extractor. val e = error(mapOf("Retry-After" to listOf("5"))) - assertNull(extractor.extract(e)) + assertEquals(5.seconds, extractor.extract(e)) + } + + @Test + fun testCustomHeaderNames() { + // Header names are configurable per provider and lowercased on construction. + val custom = StandardHeaderRetryAfterExtractor( + retryAfterHeaders = listOf("X-Custom-Retry"), + resetDurationHeaders = listOf("x-custom-reset") + ) + + assertEquals(7.seconds, custom.extract(error(mapOf("x-custom-retry" to listOf("7"))))) + assertEquals(3.seconds, custom.extract(error(mapOf("x-custom-reset" to listOf("3s"))))) + // The defaults are replaced, not extended. + assertNull(custom.extract(error(mapOf("retry-after" to listOf("5"))))) } @Test @@ -172,8 +235,4 @@ class StandardHeaderRetryAfterExtractorTest { ) assertEquals(99.seconds, composite.extract(withoutHeader)) } - - private fun fixedClockAt(instant: Instant): Clock = object : Clock { - override fun now(): Instant = instant - } } From 6e2c9d8a7a052bc8a8136e49509523a0f7b736cb Mon Sep 17 00:00:00 2001 From: kyboi Date: Tue, 29 Sep 2026 18:26:30 +0100 Subject: [PATCH 09/10] fix(http-client): adapt header retry changes to current develop --- .../kotlin/ai/koog/http/client/ktor/KtorKoogHttpClient.kt | 2 -- .../kotlin/ai/koog/http/client/test/BaseKoogHttpClientTest.kt | 4 ++-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/http-client/http-client-ktor/src/commonMain/kotlin/ai/koog/http/client/ktor/KtorKoogHttpClient.kt b/http-client/http-client-ktor/src/commonMain/kotlin/ai/koog/http/client/ktor/KtorKoogHttpClient.kt index d60a42c8f5..4e87f3341b 100644 --- a/http-client/http-client-ktor/src/commonMain/kotlin/ai/koog/http/client/ktor/KtorKoogHttpClient.kt +++ b/http-client/http-client-ktor/src/commonMain/kotlin/ai/koog/http/client/ktor/KtorKoogHttpClient.kt @@ -116,8 +116,6 @@ public class KtorKoogHttpClient internal constructor( } } } - private fun Headers.toResponseHeaderMap(): Map> = - entries().associate { it.key to it.value }.lowercaseHeaderKeys() override suspend fun get( path: String, diff --git a/http-client/http-client-test/src/main/kotlin/ai/koog/http/client/test/BaseKoogHttpClientTest.kt b/http-client/http-client-test/src/main/kotlin/ai/koog/http/client/test/BaseKoogHttpClientTest.kt index 80641d4d44..3411a05fa2 100644 --- a/http-client/http-client-test/src/main/kotlin/ai/koog/http/client/test/BaseKoogHttpClientTest.kt +++ b/http-client/http-client-test/src/main/kotlin/ai/koog/http/client/test/BaseKoogHttpClientTest.kt @@ -206,7 +206,7 @@ abstract class BaseKoogHttpClientTest { try { client.post( path = mockServer.url("/rate-limited"), - request = "PAYLOAD", + requestBody = "PAYLOAD", ) fail("Expected a KoogHttpClientException for 429") } catch (e: KoogHttpClientException) { @@ -266,7 +266,7 @@ abstract class BaseKoogHttpClientTest { val flow = client.sse( path = mockServer.url("/stream"), - request = "{}", + requestBody = "{}", requestBodyType = String::class, dataFilter = { it != "[DONE]" }, decodeStreamingResponse = { it }, From ae8529b323b0627f012c2b0e40858941bf97d76c Mon Sep 17 00:00:00 2001 From: kyboi Date: Tue, 29 Sep 2026 18:26:33 +0100 Subject: [PATCH 10/10] fix(prompt): cap retry delay after jitter --- .../clients/retry/RetryingLLMClient.kt | 2 +- .../clients/retry/RetryingLLMClientTest.kt | 45 +++++++++++++++++++ 2 files changed, 46 insertions(+), 1 deletion(-) diff --git a/prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClient.kt b/prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClient.kt index fd20e7f800..0a61f542b1 100644 --- a/prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClient.kt +++ b/prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClient.kt @@ -227,7 +227,7 @@ public class RetryingLLMClient @JvmOverloads constructor( // collapses to zero (jitterFactor == 0.0 or initialDelay == Duration.ZERO). val jitterUpper = boundedMs * config.jitterFactor val jitterMs = if (jitterUpper > 0.0) Random.nextDouble(0.0, jitterUpper) else 0.0 - return (boundedMs + jitterMs).toLong().milliseconds + return (boundedMs + jitterMs).toLong().milliseconds.coerceAtMost(config.maxDelay) } override fun close() { diff --git a/prompt/prompt-executor/prompt-executor-clients/src/commonTest/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClientTest.kt b/prompt/prompt-executor/prompt-executor-clients/src/commonTest/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClientTest.kt index de7699aaff..0e0f68dcfa 100644 --- a/prompt/prompt-executor/prompt-executor-clients/src/commonTest/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClientTest.kt +++ b/prompt/prompt-executor/prompt-executor-clients/src/commonTest/kotlin/ai/koog/prompt/executor/clients/retry/RetryingLLMClientTest.kt @@ -31,6 +31,7 @@ import kotlin.test.Test import kotlin.test.assertEquals import kotlin.test.assertFailsWith import kotlin.test.assertSame +import kotlin.test.assertTrue import kotlin.time.Duration.Companion.milliseconds import kotlin.time.Duration.Companion.seconds @@ -542,6 +543,48 @@ class RetryingLLMClientTest { assertEquals(1, mockClient.executeCalls) // No retry on cancellation } + @Test + fun testJitterNeverExceedsMaxDelay() = runTest { + repeat(20) { + val mockClient = MockLLMClient( + executeResponse = testResponse, + failuresBeforeSuccess = 1, + failureMessage = "Error: 503" + ) + val retryingClient = RetryingLLMClient( + mockClient, + RetryConfig(maxAttempts = 2, initialDelay = 1.seconds, maxDelay = 1.seconds, jitterFactor = 1.0) + ) + + val startTime = testScheduler.currentTime + val result = retryingClient.execute(testPrompt, testModel, emptyList()) + + assertEquals(testResponse, result) + assertEquals(2, mockClient.executeCalls) + assertEquals(1_000L, testScheduler.currentTime - startTime) + } + } + + @Test + fun testJitterBelowMaxDelayRetainsItsRange() = runTest { + val mockClient = MockLLMClient( + executeResponse = testResponse, + failuresBeforeSuccess = 1, + failureMessage = "Error: 503" + ) + val retryingClient = RetryingLLMClient( + mockClient, + RetryConfig(maxAttempts = 2, initialDelay = 100.milliseconds, maxDelay = 1.seconds, jitterFactor = 1.0) + ) + + val startTime = testScheduler.currentTime + val result = retryingClient.execute(testPrompt, testModel, emptyList()) + + assertEquals(testResponse, result) + assertEquals(2, mockClient.executeCalls) + assertTrue(testScheduler.currentTime - startTime in 100L..199L) + } + @Test fun testJitterFactorZeroRetriesWithoutThrowing() = runTest { // jitterFactor = 0.0 is documented as a valid "no jitter" setting and accepted by @@ -559,10 +602,12 @@ class RetryingLLMClientTest { RetryConfig(maxAttempts = 2, initialDelay = 10.milliseconds, jitterFactor = 0.0) ) + val startTime = testScheduler.currentTime val result = retryingClient.execute(testPrompt, testModel, emptyList()) assertEquals(testResponse, result) assertEquals(2, mockClient.executeCalls) + assertEquals(10L, testScheduler.currentTime - startTime) } @Test