Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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 {}
Expand Down Expand Up @@ -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 ->
Expand Down Expand Up @@ -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)

Expand Down
Original file line number Diff line number Diff line change
@@ -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()
}
}
}