fix(prompt): propagate wrapped cancellation through retries

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<Throwable>()
+        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
 

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<StreamFrame>()
+                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<Int, Throwable>(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<Int, Throwable>(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<ToolDescriptor>
         ): Message.Assistant {
             executeCalls++
+            executeFailure?.let { throw it }
 
             if (throwCancellation) {
                 throw CancellationException("Cancelled")