diff --git a/docs/docs/prompts/handling-failures.md b/docs/docs/prompts/handling-failures.md index 1ccdaaf13d..6f58da9370 100644 --- a/docs/docs/prompts/handling-failures.md +++ b/docs/docs/prompts/handling-failures.md @@ -220,6 +220,105 @@ 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. 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`. +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 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 + // Provider-specific header names first, message-based fallback second. + val config = RetryConfig( + retryAfterExtractor = CompositeRetryAfterExtractor( + StandardHeaderRetryAfterExtractor( + retryAfterHeaders = listOf("x-my-retry-after"), + resetDurationHeaders = listOf("x-my-ratelimit-reset"), + ), + 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 = new RetryConfigBuilder() + .retryAfterExtractor(new CompositeRetryAfterExtractor( + myExtractor, + DefaultRetryAfterExtractor.INSTANCE + )) + .build(); + ``` + + ### Streaming with retry Streaming operations can optionally be retried. This feature is disabled by default. @@ -250,7 +349,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. @@ -302,7 +401,7 @@ To learn more about prompt executors, see [Prompt executors](prompt-executors.md ), ) ``` - + === "Java" @@ -339,7 +438,7 @@ To learn more about prompt executors, see [Prompt executors](prompt-executors.md MultiLLMPromptExecutor multiExecutor = new MultiLLMPromptExecutor(clients); ``` - + ## Timeout configuration @@ -377,7 +476,7 @@ You can customize these values for your specific needs. For example: ) ) ``` - + === "Java" @@ -405,7 +504,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`. @@ -474,7 +573,7 @@ Here is an example of error handling in Kotlin and Java: } } ``` - + === "Java" @@ -529,4 +628,4 @@ Here is an example of error handling in Kotlin and Java: } } ``` - + 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 751ce63de0..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 @@ -1,14 +1,24 @@ 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. + * @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( +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, + headers: Map> = emptyMap() ) : Exception( buildString { appendLine("Error from client: ${clientName ?: "unknown client"}") @@ -20,4 +30,14 @@ public class KoogHttpClientException( } }, 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 new file mode 100644 index 0000000000..4c56a6aef0 --- /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. + * + * [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() + 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-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 ff5c5fa694..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 @@ -71,6 +71,7 @@ public class JavaKoogHttpClient internal constructor( clientName = clientName, statusCode = response.statusCode(), errorBody = response.body(), + headers = response.headers().map(), ) } @@ -196,6 +197,7 @@ public class JavaKoogHttpClient internal constructor( KoogHttpClientException( clientName = clientName, statusCode = response.statusCode(), + headers = response.headers().map(), ) ) return@callbackFlow @@ -285,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 2156bcee73..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 @@ -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`() @@ -70,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 074e01e89b..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 @@ -27,6 +27,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,9 +96,14 @@ public class KtorKoogHttpClient internal constructor( clientName = clientName, statusCode = response.status.value, errorBody = response.bodyAsText(), + headers = response.headers.toResponseHeaderMap(), ) } + // 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)) { @@ -210,7 +216,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 @@ -250,6 +257,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 96fc8e83fe..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 @@ -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`() @@ -84,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 742f2a89de..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 @@ -72,6 +72,7 @@ public class OkHttpKoogHttpClient internal constructor( clientName = clientName, statusCode = response.code, errorBody = response.body.string(), + headers = response.headers.toMultimap(), ) } @@ -204,12 +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 + cause = t, + headers = response?.headers?.toMultimap() ?: emptyMap() ) logger.error(exception) { exception.message } close(exception) @@ -253,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 abd3b562f0..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 @@ -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`() @@ -70,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 2377ea4795..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 @@ -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"), + requestBody = "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") @@ -215,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"), + requestBody = "{}", + 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") @@ -390,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 7393d738b9..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 @@ -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 @@ -41,7 +42,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,15 +55,22 @@ 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() ) /** - * 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() ) /** @@ -74,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 = {}, @@ -116,6 +126,9 @@ class MockWebServer { } } + config.responseHeaders.forEach { (name, value) -> + call.response.header(name, value) + } call.respondText( text = config.responseBody, contentType = config.contentType, @@ -153,6 +166,9 @@ class MockWebServer { } } + config.responseHeaders.forEach { (name, value) -> + call.response.header(name, value) + } call.respondText( text = config.responseBody, contentType = config.contentType, @@ -189,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 -> @@ -210,6 +229,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/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/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..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,10 +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.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,169 @@ public object DefaultRetryAfterExtractor : RetryAfterExtractor { return null } } + +/** + * Extractor that reads structured rate-limit hints from HTTP response headers attached to a + * [KoogHttpClientException]. + * + * Recognized headers (keys on [KoogHttpClientException.headers] are already normalized to + * lowercase; configured names are lowercased on construction): + * + * - [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). + * + * `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 @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 + + // 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 resetDurationHeaders + .flatMap { name -> headers.usableHints(name, ::parseResetDuration) } + .minOrNull() + } + + // 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? { + 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 + } + + /** + * 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() + } + + /** + * 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 the default header names and a shared default instance backed by + * [KoogClock.System]. + */ + public companion object { + // 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() + + 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" + ) + } +} + +/** + * 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) } + + /** Renders the chain so retry logs make the consultation order obvious. */ + override fun toString(): String = + extractors.joinToString(prefix = "CompositeRetryAfterExtractor(", postfix = ")") +} 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..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 @@ -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 @@ -55,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( @@ -183,34 +187,47 @@ 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) - // 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 in error message - error?.message?.let { message -> - config.retryAfterExtractor?.extract(message)?.let { retryAfter -> - return retryAfter - } + // 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)?.coerceAtMost(config.maxDelay) ?: exponentialBackoffWithJitter(attempt) + + // Prefers the header-aware extractor overload when the throwable carries a + // 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 + error.unwrapHttpException()?.let { koogException -> + extractor.extract(koogException)?.let { return it } } + return 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) - val jitterMs = Random.nextDouble(0.0, boundedMs * config.jitterFactor) - val finalMs = (boundedMs + jitterMs).toLong() - - return finalMs.milliseconds + // 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. + // 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.coerceAtMost(config.maxDelay) } override fun close() { @@ -224,6 +241,15 @@ public class RetryingLLMClient @JvmOverloads constructor( override fun getBasicJsonSchemaGenerator(): BasicJsonSchemaGenerator { return delegate.getBasicJsonSchemaGenerator() } + + // 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? = + 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 daffef7aa9..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 @@ -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 @@ -30,7 +31,9 @@ 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 class RetryingLLMClientTest { @@ -53,6 +56,299 @@ 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 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 + // 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 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( @@ -247,6 +543,73 @@ 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 + // 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 startTime = testScheduler.currentTime + val result = retryingClient.execute(testPrompt, testModel, emptyList()) + + assertEquals(testResponse, result) + assertEquals(2, mockClient.executeCalls) + assertEquals(10L, testScheduler.currentTime - startTime) + } + @Test fun testDisabledRetryConfig() = runTest { val mockClient = MockLLMClient( @@ -544,9 +907,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 +945,7 @@ class RetryingLLMClientTest { if (executeFailures < failuresBeforeSuccess) { executeFailures++ - throw RuntimeException(failureMessage) + throw failureToThrow() } return executeResponse!! @@ -595,7 +960,7 @@ class RetryingLLMClientTest { if (streamFailures < streamFailuresBeforeSuccess) { streamFailures++ - throw RuntimeException(failureMessage) + throw failureToThrow() } streamResponse.collect { emit(it) } @@ -610,7 +975,7 @@ class RetryingLLMClientTest { if (multipleChoicesFailures < failuresBeforeSuccess) { multipleChoicesFailures++ - throw RuntimeException(failureMessage) + throw failureToThrow() } return multipleChoicesResponse!! @@ -621,7 +986,7 @@ class RetryingLLMClientTest { if (moderateFailures < failuresBeforeSuccess) { moderateFailures++ - throw RuntimeException(failureMessage) + throw failureToThrow() } return moderateResponse @@ -635,7 +1000,7 @@ class RetryingLLMClientTest { if (embedFailures < failuresBeforeSuccess) { embedFailures++ - throw RuntimeException(failureMessage) + throw failureToThrow() } return embedResponse @@ -649,7 +1014,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..71b6bbfa8d --- /dev/null +++ b/prompt/prompt-executor/prompt-executor-clients/src/commonTest/kotlin/ai/koog/prompt/executor/clients/retry/StandardHeaderRetryAfterExtractorTest.kt @@ -0,0 +1,238 @@ +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.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 means "no usable hint", not "retry immediately"; 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 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 = 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 testHttpDateRetryAfterInPastIsIgnored() { + val now = Instant.parse("2026-01-01T00:01:00Z") + val fixedExtractor = StandardHeaderRetryAfterExtractor(clock = KoogClock { now }) + + val e = error(mapOf("retry-after" to listOf("Thu, 01 Jan 2026 00:00:00 GMT"))) + // 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"))))) + 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() { + // 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 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("500ms"), + "x-ratelimit-reset-tokens" to listOf("6m0s") + ) + ) + assertEquals(10.seconds, extractor.extract(e)) + } + + @Test + 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(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 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"))) + 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 + 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)) + } +}