From f5cf56ab71509f1b57da4e2ce1ecbee611bea34a Mon Sep 17 00:00:00 2001 From: kyboi Date: Tue, 29 Sep 2026 18:00:20 +0100 Subject: [PATCH] fix(prompt): propagate wrapped cancellation through retries --- .../clients/retry/RetryingLLMClient.kt | 15 +++ .../clients/retry/RetryingLLMClientTest.kt | 121 ++++++++++++++++++ 2 files changed, 136 insertions(+) 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..d330a5a35d 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 @@ -55,6 +55,7 @@ public class RetryingLLMClient @JvmOverloads constructor( private companion object { private val logger = KotlinLogging.logger { } + private const val MAX_CAUSE_DEPTH = 32 } override suspend fun execute( @@ -84,6 +85,7 @@ public class RetryingLLMClient @JvmOverloads constructor( } catch (e: CancellationException) { throw e // Never retry cancellations } catch (e: Throwable) { + e.cancellationCause()?.let { throw it } // If we already received tokens, don't retry - pass error through if (firstFrameReceived) { throw e @@ -162,6 +164,7 @@ public class RetryingLLMClient @JvmOverloads constructor( } catch (e: CancellationException) { throw e // Never retry cancellations } catch (e: Throwable) { + e.cancellationCause()?.let { throw it } lastException = e if (!shouldRetry(e) || attempt >= config.maxAttempts - 1) { @@ -180,6 +183,18 @@ public class RetryingLLMClient @JvmOverloads constructor( throw lastException!! } + private fun Throwable.cancellationCause(): CancellationException? { + val visited = mutableListOf() + var current: Throwable? = this + // Bound traversal and compare by identity to handle malformed cause chains safely. + while (current != null && visited.size < MAX_CAUSE_DEPTH && visited.none { it === current }) { + if (current is CancellationException) return current + visited += current + current = current.cause + } + return null + } + private fun shouldRetry(error: Throwable): Boolean { if (error is IncompleteStreamException) return true 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..62dab404cb 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 @@ -25,11 +25,13 @@ import kotlinx.coroutines.flow.collect import kotlinx.coroutines.flow.flow import kotlinx.coroutines.flow.flowOf import kotlinx.coroutines.flow.toList +import kotlinx.coroutines.isActive import kotlinx.coroutines.test.runTest 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 class RetryingLLMClientTest { @@ -533,6 +535,123 @@ class RetryingLLMClientTest { assertEquals(mockClient.standardJsonSchemaGeneratorDefault, result) } + @Test + fun testWrappedCancellationIsNotRetried() = runTest { + val cancellation = CancellationException("Cancelled operation") + val delegate = MockLLMClient( + streamResponse = flow { + throw IllegalStateException("429 Too Many Requests", cancellation) + } + ) + val client = RetryingLLMClient( + delegate, + RetryConfig(maxAttempts = 3, initialDelay = 10.milliseconds) + ) + + val failure = runCatching { + client.executeStreaming(testPrompt, testModel).collect() + }.exceptionOrNull() + + assertEquals(1, delegate.streamCalls) + assertSame(cancellation, failure) + } + + @Test + fun testCancellationCausesInExecute() = runTest { + val cancellation = CancellationException("Delegate cancelled") + val failures = listOf( + cancellation, + IllegalStateException("429 Too Many Requests", cancellation), + IllegalStateException("429 Too Many Requests", RuntimeException("Nested", cancellation)), + IllegalArgumentException("Invalid request", cancellation) + ) + for (failure in failures) { + val mockClient = MockLLMClient(executeFailure = failure) + val client = RetryingLLMClient(mockClient, RetryConfig(maxAttempts = 3)) + + val actual = runCatching { client.execute(testPrompt, testModel) }.exceptionOrNull() + + assertEquals(1, mockClient.executeCalls) + assertSame(cancellation, actual) + assertTrue(coroutineContext.isActive) + } + } + + @Test + fun testStreamingCancellationCauses() = runTest { + val cancellation = CancellationException("Delegate cancelled") + val failures = listOf( + cancellation, + IllegalStateException("429 Too Many Requests", RuntimeException("Nested", cancellation)), + IllegalArgumentException("Invalid request", cancellation) + ) + for (failure in failures) { + for (emitFirstFrame in listOf(false, true)) { + val frames = mutableListOf() + val mockClient = MockLLMClient( + streamResponse = flow { + if (emitFirstFrame) emit(StreamFrame.TextDelta("first")) + throw failure + } + ) + val client = RetryingLLMClient(mockClient, RetryConfig(maxAttempts = 3)) + + val actual = runCatching { + client.executeStreaming(testPrompt, testModel).collect { frames += it } + }.exceptionOrNull() + + assertEquals(1, mockClient.streamCalls) + assertSame(cancellation, actual) + assertEquals(if (emitFirstFrame) listOf(StreamFrame.TextDelta("first")) else emptyList(), frames) + assertTrue(coroutineContext.isActive) + } + } + } + + @Test + fun testCancellationAtCauseTraversalBoundary() = runTest { + val cancellation = CancellationException("Delegate cancelled") + val failure = (1..31).fold(cancellation) { cause, _ -> + IllegalStateException("429 Too Many Requests", cause) + } + val mockClient = MockLLMClient(executeFailure = failure) + val client = RetryingLLMClient(mockClient) + + val actual = runCatching { client.execute(testPrompt, testModel) }.exceptionOrNull() + + assertSame(cancellation, actual) + assertEquals(1, mockClient.executeCalls) + } + + @Test + fun testCauseTraversalStopsAtBound() = runTest { + val cancellation = CancellationException("Delegate cancelled") + val failure = (1..32).fold(cancellation) { cause, _ -> + IllegalStateException("429 Too Many Requests", cause) + } + val mockClient = MockLLMClient(executeFailure = failure) + val client = RetryingLLMClient(mockClient, RetryConfig(maxAttempts = 2)) + + val actual = runCatching { client.execute(testPrompt, testModel) }.exceptionOrNull() + + assertSame(failure, actual) + assertEquals(2, mockClient.executeCalls) + } + + @Test + fun testCyclicCausesPreserveRetry() = runTest { + val failure = object : RuntimeException("429 Too Many Requests") { + override val cause: Throwable get() = this + } + val mockClient = MockLLMClient(executeFailure = failure) + val client = RetryingLLMClient(mockClient, RetryConfig(maxAttempts = 2)) + + val actual = runCatching { client.execute(testPrompt, testModel) }.exceptionOrNull() + + assertSame(failure, actual) + assertEquals(2, mockClient.executeCalls) + } + // Mock LLMClient for testing private class MockLLMClient( private val executeResponse: Message.Assistant? = null, @@ -545,6 +664,7 @@ class RetryingLLMClientTest { private var streamFailuresBeforeSuccess: Int = 0, private val failureMessage: String = "Mock failure", private val throwCancellation: Boolean = false, + private val executeFailure: Throwable? = null, private val llmProvider: LLMProvider = LLMProvider.OpenAI, ) : LLMClient() { @@ -573,6 +693,7 @@ class RetryingLLMClientTest { tools: List ): Message.Assistant { executeCalls++ + executeFailure?.let { throw it } if (throwCancellation) { throw CancellationException("Cancelled")