From 6379d417b94a89061dff9294e98ac7cd8c47e13b Mon Sep 17 00:00:00 2001
From: kyboi <kevin@kreoh.com>
Date: Tue, 29 Sep 2026 16:06:17 +0100
Subject: [PATCH] fix(skills): propagate discovery cancellation

---
 .../koog/skills/discovery/SkillsDiscovery.kt  |  9 +-
 .../SkillsDiscoveryCancellationTest.kt        | 88 +++++++++++++++++++
 2 files changed, 95 insertions(+), 2 deletions(-)
 create mode 100644 skills/src/jvmTest/kotlin/ai/koog/skills/discovery/SkillsDiscoveryCancellationTest.kt

diff --git a/skills/src/commonMain/kotlin/ai/koog/skills/discovery/SkillsDiscovery.kt b/skills/src/commonMain/kotlin/ai/koog/skills/discovery/SkillsDiscovery.kt
index 7192885313..63cbf1079d 100644
--- a/skills/src/commonMain/kotlin/ai/koog/skills/discovery/SkillsDiscovery.kt
+++ b/skills/src/commonMain/kotlin/ai/koog/skills/discovery/SkillsDiscovery.kt
@@ -5,6 +5,7 @@ import ai.koog.rag.base.files.FileMetadata
 import ai.koog.rag.base.files.FileSystemProvider
 import ai.koog.skills.model.Skill
 import io.github.oshai.kotlinlogging.KotlinLogging
+import kotlin.coroutines.cancellation.CancellationException
 import kotlin.jvm.JvmName
 
 private val logger = KotlinLogging.logger {}
@@ -80,7 +81,9 @@ public suspend fun <Path> discoverSkills(
         val (directory, depth) = queue.removeFirst()
         visitedDirectories++
 
-        val children = runCatching { fs.list(directory) }.getOrDefault(emptyList())
+        val children = runCatching { fs.list(directory) }
+            .onFailure { if (it is CancellationException) throw it }
+            .getOrDefault(emptyList())
         val skillFile = children.firstOrNull { fs.name(it) == skillFileName }
         if (skillFile != null) {
             fs.parseSkill(skillFile, skillNamePattern, warningLogger)?.let { skill ->
@@ -139,7 +142,9 @@ private suspend fun <Path> FileSystemProvider.ReadOnly<Path>.parseSkill(
     skillNamePattern: Regex,
     warningLogger: (String) -> Unit,
 ): Skill? {
-    val content = runCatching { readBytes(skillFile).decodeToString() }.getOrNull() ?: return null
+    val content = runCatching { readBytes(skillFile).decodeToString() }
+        .onFailure { if (it is CancellationException) throw it }
+        .getOrNull() ?: return null
     val frontmatter = parseFrontmatter(content)
     val location = toAbsolutePathString(skillFile)
 
diff --git a/skills/src/jvmTest/kotlin/ai/koog/skills/discovery/SkillsDiscoveryCancellationTest.kt b/skills/src/jvmTest/kotlin/ai/koog/skills/discovery/SkillsDiscoveryCancellationTest.kt
new file mode 100644
index 0000000000..c6b18ad830
--- /dev/null
+++ b/skills/src/jvmTest/kotlin/ai/koog/skills/discovery/SkillsDiscoveryCancellationTest.kt
@@ -0,0 +1,88 @@
+package ai.koog.skills.discovery
+
+import ai.koog.rag.base.files.FileMetadata
+import ai.koog.rag.base.files.FileSystemProvider
+import ai.koog.rag.base.files.JVMFileSystemProvider
+import kotlinx.coroutines.CompletableDeferred
+import kotlinx.coroutines.awaitCancellation
+import kotlinx.coroutines.cancelAndJoin
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.test.runTest
+import java.io.IOException
+import java.nio.file.Path
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertFalse
+
+class SkillsDiscoveryCancellationTest {
+    @Test
+    fun testCancellationDuringListingStopsDiscovery() = testCancellation(read = false)
+
+    @Test
+    fun testCancellationDuringReadingStopsDiscovery() = testCancellation(read = true)
+
+    @Test
+    fun testListingFailureSkipsDirectory() = testOrdinaryFailure(read = false)
+
+    @Test
+    fun testReadingFailureSkipsSkill() = testOrdinaryFailure(read = true)
+
+    private fun testCancellation(read: Boolean) = runTest {
+        val entered = CompletableDeferred<Unit>()
+        val fs = DiscoveryFileSystem(read) {
+            entered.complete(Unit)
+            awaitCancellation()
+        }
+        var continued = false
+        val discovery = launch {
+            discoverSkills(fs, listOf(fs.broken.toString(), fs.healthy.toString()))
+            continued = true
+        }
+        entered.await()
+        val callsAtCancellation = fs.calls.toList()
+        discovery.cancelAndJoin()
+
+        assertEquals(callsAtCancellation, fs.calls, "Discovery must perform no later I/O after cancellation")
+        assertFalse(continued, "Caller code must not continue after discoverSkills")
+    }
+
+    private fun testOrdinaryFailure(read: Boolean) = runTest {
+        val fs = DiscoveryFileSystem(read) { throw IOException("Cannot access skill") }
+
+        val skills = discoverSkills(fs, listOf(fs.broken.toString(), fs.healthy.toString()))
+
+        assertEquals(listOf("healthy"), skills.map { it.name })
+        assertEquals(fs.toAbsolutePathString(fs.healthy.resolve("SKILL.md")), skills.single().location)
+    }
+
+    private class DiscoveryFileSystem(
+        private val failRead: Boolean,
+        private val failure: suspend () -> Nothing,
+    ) : FileSystemProvider.ReadOnly<Path> by JVMFileSystemProvider.ReadOnly {
+        val broken = Path.of("broken").toAbsolutePath()
+        val healthy = Path.of("healthy").toAbsolutePath()
+        val calls = mutableListOf<String>()
+
+        override suspend fun metadata(path: Path): FileMetadata {
+            calls += "metadata:$path"
+            val type = if (path.fileName.toString() == "SKILL.md") {
+                FileMetadata.FileType.File
+            } else {
+                FileMetadata.FileType.Directory
+            }
+            return FileMetadata(type, hidden = false)
+        }
+
+        override suspend fun list(directory: Path): List<Path> {
+            calls += "list:$directory"
+            if (!failRead && directory == broken) failure()
+            return listOf(directory.resolve("SKILL.md"))
+        }
+
+        override suspend fun readBytes(path: Path): ByteArray {
+            calls += "read:$path"
+            if (failRead && path.parent == broken) failure()
+            return "---\nname: healthy\ndescription: A healthy skill.\n---".encodeToByteArray()
+        }
+    }
+}