epic-12: after epic audit and init commit

This commit is contained in:
2026-05-16 11:42:00 +04:00
commit c77277af0b
461 changed files with 28958 additions and 0 deletions
+23
View File
@@ -0,0 +1,23 @@
plugins {
id 'java-library'
id 'org.jetbrains.kotlin.jvm'
id 'org.jetbrains.kotlin.plugin.serialization'
}
dependencies {
implementation project(":core:tools")
implementation project(":core:events")
implementation project(":core:inference")
implementation project(":core:approvals")
implementation project(":core:sessions")
implementation project(":infrastructure:inference")
implementation project(":infrastructure:inference:commons")
implementation project(":infrastructure:inference:llama_cpp")
implementation project(":infrastructure:persistence")
implementation project(":infrastructure:tools")
implementation project(":infrastructure:tools:filesystem")
implementation "io.ktor:ktor-client-core:$ktor_version"
implementation "io.ktor:ktor-client-cio:$ktor_version"
implementation "io.ktor:ktor-client-content-negotiation:$ktor_version"
implementation "io.ktor:ktor-serialization-kotlinx-json:$ktor_version"
}
+12
View File
@@ -0,0 +1,12 @@
plugins {
id 'java-library'
id 'org.jetbrains.kotlin.jvm'
id 'org.jetbrains.kotlin.plugin.serialization'
}
dependencies {
implementation project(':core:inference')
implementation project(':core:events')
testImplementation "org.junit.jupiter:junit-jupiter"
testImplementation "org.jetbrains.kotlinx:kotlinx-coroutines-test"
}
@@ -0,0 +1,19 @@
plugins {
id 'java-library'
id 'org.jetbrains.kotlin.jvm'
id 'org.jetbrains.kotlin.plugin.serialization'
}
ext {
ktor_version = '3.0.3'
}
dependencies {
implementation project(':core:inference')
implementation project(':core:events')
implementation project(':core:context')
implementation "io.ktor:ktor-client-core:$ktor_version"
implementation "io.ktor:ktor-client-cio:$ktor_version"
implementation "io.ktor:ktor-client-content-negotiation:$ktor_version"
implementation "io.ktor:ktor-serialization-kotlinx-json:$ktor_version"
}
@@ -0,0 +1,26 @@
package com.correx.infrastructure.inference.commons
import com.correx.core.events.types.ProviderId
import com.correx.core.inference.InferenceProvider
import com.correx.core.inference.Tokenizer
/**
* A wrapper around [InferenceProvider] that tracks model lifecycle.
* Ensures the provider is tied to a specific [ModelManager] instance and [ModelDescriptor].
*/
interface ManagedInferenceProvider : InferenceProvider {
/**
* Unload the managed model.
*/
suspend fun unload()
/**
* Check if the underlying model is still loaded.
*/
fun isLoaded(): Boolean
/**
* Load a different model, replacing this one.
*/
suspend fun load(newDescriptor: ModelDescriptor): ManagedInferenceProvider
}
@@ -0,0 +1,15 @@
package com.correx.infrastructure.inference.commons
import com.correx.core.inference.CapabilityScore
/**
* Descriptor for a model instance with its configuration.
*/
data class ModelDescriptor(
val modelId: String,
val modelPath: String,
val residencyMode: ResidencyMode,
val idleTimeoutMs: Long = 60_000L,
val contextSize: Int = 8192,
val capabilities: Set<CapabilityScore> = emptySet(),
)
@@ -0,0 +1,11 @@
package com.correx.infrastructure.inference.commons
/**
* Thrown when model loading fails during ModelManager.load().
* No event is emitted on failure per event sourcing principles.
*/
class ModelLoadException(
override val message: String,
val modelId: String,
override val cause: Throwable? = null,
) : RuntimeException(message, cause)
@@ -0,0 +1,36 @@
package com.correx.infrastructure.inference.commons
import com.correx.core.inference.InferenceProvider
import com.correx.core.inference.ProviderHealth
/**
* Manages the lifecycle of models.
*/
interface ModelManager {
/**
* Loads a model and returns an InferenceProvider instance.
*
* @param descriptor Model configuration
* @return An InferenceProvider for the loaded model
* @throws ModelLoadException if a different model is already loaded
*/
suspend fun load(descriptor: ModelDescriptor): InferenceProvider
/**
* Unloads the currently loaded model.
*
* @param modelId The model ID to unload
* @throws ModelLoadException if no model is currently loaded
*/
suspend fun unload(modelId: String)
/**
* Returns the descriptor of the currently loaded model, or null if none.
*/
fun currentModel(): ModelDescriptor?
/**
* Health check for the current model.
*/
fun healthCheck(): ProviderHealth
}
@@ -0,0 +1,10 @@
package com.correx.infrastructure.inference.commons
/**
* Residency modes for model lifecycle management.
*/
enum class ResidencyMode {
PERSISTENT, // never unload
DYNAMIC, // unload after idle timeout
EPHEMERAL, // unload immediately after infer completes
}
@@ -0,0 +1,22 @@
plugins {
id 'java-library'
id 'org.jetbrains.kotlin.jvm'
id 'org.jetbrains.kotlin.plugin.serialization'
}
ext {
ktor_version = '3.0.3'
// Ensure parent version is available
ext.ktor_version = '3.0.3'
}
dependencies {
implementation project(':core:inference')
implementation project(':core:events')
implementation project(':core:context')
implementation project(':infrastructure:inference:commons')
implementation "io.ktor:ktor-client-core:$ktor_version"
implementation "io.ktor:ktor-client-cio:$ktor_version"
implementation "io.ktor:ktor-client-content-negotiation:$ktor_version"
implementation "io.ktor:ktor-serialization-kotlinx-json:$ktor_version"
}
@@ -0,0 +1,208 @@
package com.correx.infrastructure.inference.llama.cpp
import com.correx.core.events.EventDispatcher
import com.correx.core.events.events.ModelLoadedEvent
import com.correx.core.events.events.ModelUnloadedEvent
import com.correx.core.events.stores.EventStore
import com.correx.core.events.types.ProviderId
import com.correx.core.events.types.SessionId
import com.correx.core.inference.InferenceProvider
import com.correx.core.inference.ProviderHealth
import com.correx.core.inference.Tokenizer
import com.correx.infrastructure.inference.commons.ManagedInferenceProvider
import com.correx.infrastructure.inference.commons.ModelDescriptor
import com.correx.infrastructure.inference.commons.ModelLoadException
import com.correx.infrastructure.inference.commons.ModelManager
import io.ktor.client.HttpClient
import io.ktor.client.call.body
import io.ktor.client.request.get
import kotlinx.coroutines.delay
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock
import kotlinx.datetime.Clock
import java.io.File
/**
* Implementation of ModelManager that manages llama.cpp server process lifecycle.
*
* Enforces single-model-at-a-time policy via Mutex.
* Model swaps are performed via process restart (kill existing, spawn new with --model flag).
*
* Thread-safe: all public methods use Mutex for synchronization.
*/
@Suppress("LongParameterList", "TooGenericExceptionCaught", "UnusedPrivateProperty")
class DefaultModelManager(
private val llamaServerBin: String = "llama-server",
private val host: String = "127.0.0.1",
private val port: Int = 10000,
private val healthTimeoutMs: Long = 30_000L,
private val eventStore: EventStore,
private val httpClient: HttpClient,
private val eventDispatcher: EventDispatcher,
) : ModelManager {
private var currentProcess: LlamaProcess? = null
private var currentDescriptor: ModelDescriptor? = null
private val mutex = Mutex()
private val logDirPath = System.getProperty("user.home") + "/.local/share/correx/logs"
override suspend fun load(descriptor: ModelDescriptor): InferenceProvider = mutex.withLock {
// If same model already loaded, return existing provider
if (currentDescriptor?.modelId == descriptor.modelId) {
return@withLock createProvider(descriptor)
}
// Kill existing process if running
stopCurrentModelInternal()
// Spawn new process
val process = try {
val logFile = File(logDirPath, "llama-server.log")
val command = listOf(
llamaServerBin,
"--model", descriptor.modelPath,
"--ctx-size", descriptor.contextSize.toString(),
"--host", host,
"--port", port.toString(),
)
val lp = LlamaProcess(command, logFile)
lp.start()
lp
} catch (e: Exception) {
throw ModelLoadException("Failed to start llama-server: ${e.message}", descriptor.modelId, e)
}
// Health check polling
if (!waitForHealthy(process)) {
process.stop()
throw ModelLoadException(
"Health check timeout after ${healthTimeoutMs}ms",
descriptor.modelId,
)
}
// Set state
currentProcess = process
currentDescriptor = descriptor
// Emit event
val providerId = ProviderId("llama-cpp:${descriptor.modelId}")
eventDispatcher.emit(
payload = ModelLoadedEvent(
modelId = descriptor.modelId,
providerId = providerId,
sessionId = SessionId("correx:model-manager"),
),
sessionId = SessionId("correx:model-manager")
)
// Return provider
createProvider(descriptor)
}
override suspend fun unload(modelId: String) = mutex.withLock {
val process = currentProcess ?: throw ModelLoadException("No model currently loaded", modelId)
val descriptor = currentDescriptor ?: throw ModelLoadException("No model currently loaded", modelId)
if (descriptor.modelId != modelId) {
throw ModelLoadException("Model ID mismatch: expected ${descriptor.modelId}, got $modelId", modelId)
}
process.stop()
currentProcess = null
currentDescriptor = null
// Emit event
val providerId = ProviderId("llama-cpp")
eventDispatcher.emit(
payload = ModelUnloadedEvent(
modelId = modelId,
providerId = providerId,
sessionId = SessionId("correx:model-manager"),
cancellationReason = null,
),
sessionId = SessionId("correx:model-manager")
)
}
override fun currentModel(): ModelDescriptor? = currentDescriptor
override fun healthCheck(): ProviderHealth {
val process = currentProcess ?: return ProviderHealth.Unavailable("No model loaded")
return if (process.isAlive) {
ProviderHealth.Healthy
} else {
ProviderHealth.Unavailable("Process not running")
}
}
private suspend fun stopCurrentModelInternal() {
currentProcess?.stop()
currentProcess = null
currentDescriptor = null
}
@Suppress("UnusedParameter")
private suspend fun waitForHealthy(process: LlamaProcess): Boolean {
val endTime = Clock.System.now().toEpochMilliseconds() + healthTimeoutMs
while (Clock.System.now().toEpochMilliseconds() < endTime) {
try {
val response = httpClient.get("http://$host:$port/health").body<String>()
if (response.contains("\"status\":\"healthy\"") || response.contains("ok")) {
return true
}
} catch (_: Exception) {
// Continue polling
}
delay(POLL_INTERVAL)
}
return false
}
private fun createProvider(descriptor: ModelDescriptor): ManagedInferenceProvider {
val providerId = ProviderId("llama-cpp:${descriptor.modelId}")
val baseUrl = "http://$host:$port"
val delegate = LlamaCppInferenceProvider(
descriptor = descriptor,
baseUrl = baseUrl,
httpClient = httpClient,
)
return DefaultManagedInferenceProvider(delegate, this, descriptor, providerId)
}
private companion object {
private const val POLL_INTERVAL = 1000L
}
}
/**
* Implementation of [ManagedInferenceProvider] that tracks model lifecycle.
* Ensures the provider is tied to a specific [ModelManager] instance and [ModelDescriptor].
*/
class DefaultManagedInferenceProvider(
private val delegate: InferenceProvider,
private val manager: ModelManager,
private val descriptor: ModelDescriptor,
override val id: ProviderId,
) : ManagedInferenceProvider, InferenceProvider by delegate {
override val name: String = "Managed(${delegate.name})"
override val tokenizer: Tokenizer = delegate.tokenizer
override suspend fun unload() {
manager.unload(descriptor.modelId)
}
override fun isLoaded(): Boolean {
return manager.currentModel()?.modelId == descriptor.modelId
}
override suspend fun load(newDescriptor: ModelDescriptor): ManagedInferenceProvider {
val provider = manager.load(newDescriptor) as? ManagedInferenceProvider
?: throw ModelLoadException("Failed to load new model", newDescriptor.modelId)
return provider
}
}
@@ -0,0 +1,45 @@
package com.correx.infrastructure.inference.llama.cpp
import kotlinx.serialization.SerialName
import kotlinx.serialization.Serializable
@Serializable
data class ChatCompletionRequest(
val model: String,
val messages: List<ChatMessage>,
val temperature: Double,
@SerialName("top_p") val topP: Double,
@SerialName("max_tokens") val maxTokens: Int,
@SerialName("stop") val stopSequences: List<String> = emptyList(),
val seed: Long? = null,
val stream: Boolean = false,
)
@Serializable
data class ChatMessage(val role: String, val content: String)
@Serializable
data class ChatCompletionResponse(
val id: String,
val choices: List<Choice>,
val usage: Usage,
)
@Serializable
data class Choice(
val message: ChatMessage,
@SerialName("finish_reason") val finishReason: String,
)
@Serializable
data class Usage(
@SerialName("prompt_tokens") val promptTokens: Int,
@SerialName("completion_tokens") val completionTokens: Int,
@SerialName("total_tokens") val totalTokens: Int,
)
@Serializable
data class TokenizeRequest(val tokens: List<String>)
@Serializable
data class TokenizeResponse(val tokens: List<Int>)
@@ -0,0 +1,120 @@
package com.correx.infrastructure.inference.llama.cpp
import com.correx.core.context.model.ContextLayer
import com.correx.core.context.model.ContextPack
import com.correx.core.events.types.ProviderId
import com.correx.core.inference.CapabilityScore
import com.correx.core.inference.FinishReason
import com.correx.core.inference.InferenceProvider
import com.correx.core.inference.InferenceRequest
import com.correx.core.inference.InferenceResponse
import com.correx.core.inference.ProviderHealth
import com.correx.core.inference.TokenUsage
import com.correx.core.inference.Tokenizer
import com.correx.infrastructure.inference.commons.ModelDescriptor
import io.ktor.client.HttpClient
import io.ktor.client.call.body
import io.ktor.client.engine.cio.CIO
import io.ktor.client.plugins.contentnegotiation.ContentNegotiation
import io.ktor.client.request.accept
import io.ktor.client.request.get
import io.ktor.client.request.post
import io.ktor.client.request.setBody
import io.ktor.http.ContentType
import io.ktor.http.contentType
import io.ktor.serialization.kotlinx.json.json
import kotlinx.serialization.json.Json
private fun defaultHttpClient(): HttpClient = HttpClient(CIO) {
install(ContentNegotiation) {
json(
Json {
ignoreUnknownKeys = true
isLenient = true
encodeDefaults = false
},
)
}
}
@Suppress("TooGenericExceptionCaught", "MagicNumber")
class LlamaCppInferenceProvider(
private val descriptor: ModelDescriptor,
private val baseUrl: String = "http://127.0.0.1:10000",
private val httpClient: HttpClient = defaultHttpClient(),
) : InferenceProvider {
private val modelId: String = descriptor.modelId
override val id: ProviderId = ProviderId("llama-cpp:$modelId")
override val name: String = "LlamaCpp ($modelId)"
override val tokenizer: Tokenizer = LlamaCppTokenizer(baseUrl, httpClient)
override suspend fun infer(request: InferenceRequest): InferenceResponse {
val startTime = System.currentTimeMillis()
val messages = buildMessages(request.contextPack)
val body = ChatCompletionRequest(
model = modelId,
messages = messages,
temperature = request.generationConfig.temperature,
topP = request.generationConfig.topP,
maxTokens = request.generationConfig.maxTokens,
stopSequences = request.generationConfig.stopSequences,
seed = request.generationConfig.seed,
stream = false,
)
val response = httpClient.post("$baseUrl/v1/chat/completions") {
contentType(ContentType.Application.Json)
accept(ContentType.Application.Json)
setBody(body)
}.body<ChatCompletionResponse>()
val finishReason = response.choices.first().finishReason.lowercase().let {
if (it == "length") FinishReason.Length else FinishReason.Stop
}
return InferenceResponse(
requestId = request.requestId,
text = response.choices.first().message.content,
finishReason = finishReason,
tokensUsed = TokenUsage(
promptTokens = response.usage.promptTokens,
completionTokens = response.usage.completionTokens,
),
latencyMs = System.currentTimeMillis() - startTime,
)
}
override suspend fun healthCheck(): ProviderHealth = try {
val response = httpClient.get("$baseUrl/health")
if (response.status.value in 200..299) {
ProviderHealth.Healthy
} else {
ProviderHealth.Unavailable("Health check returned non-2xx status: ${response.status}")
}
} catch (e: Exception) {
ProviderHealth.Unavailable("Health check failed: ${e.message}")
}
override fun capabilities(): Set<CapabilityScore> = descriptor.capabilities
private fun buildMessages(contextPack: ContextPack): List<ChatMessage> {
val messages = contextPack.layers.flatMap { (layer, entries) ->
when (layer) {
ContextLayer.L0 -> entries.map { ChatMessage("system", it.content) }
ContextLayer.L1 -> entries.map { ChatMessage("user", it.content) }
ContextLayer.L2 -> entries.map {
ChatMessage(
role = if (it.sourceType == "assistant") "assistant" else "user",
content = it.content
)
}
else -> emptyList()
}
}
return messages.ifEmpty { listOf(ChatMessage("user", "")) }
}
}
@@ -0,0 +1,26 @@
package com.correx.infrastructure.inference.llama.cpp
import com.correx.core.inference.Token
import com.correx.core.inference.Tokenizer
import io.ktor.client.HttpClient
import io.ktor.client.call.body
import io.ktor.client.request.accept
import io.ktor.client.request.post
import io.ktor.client.request.setBody
import io.ktor.http.ContentType
import io.ktor.http.contentType
class LlamaCppTokenizer(
private val baseUrl: String = "http://127.0.0.1:10000",
private val httpClient: HttpClient,
) : Tokenizer {
override suspend fun tokenize(text: String): List<Token> =
httpClient.post("$baseUrl/tokenize") {
contentType(ContentType.Application.Json)
accept(ContentType.Application.Json)
setBody(TokenizeRequest(tokens = listOf(text)))
}.body<TokenizeResponse>().tokens.map { Token(it) }
override suspend fun countTokens(text: String): Int = tokenize(text).size
}
@@ -0,0 +1,39 @@
package com.correx.infrastructure.inference.llama.cpp
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext
import java.io.File
/**
* Manages the lifecycle of a llama.cpp server process.
*/
class LlamaProcess(
private val command: List<String>,
private val logFile: File
) {
private var process: Process? = null
val isAlive: Boolean
get() = process?.isAlive ?: false
suspend fun start() = withContext(Dispatchers.IO) {
if (isAlive) return@withContext
logFile.parentFile?.mkdirs()
process = ProcessBuilder(command)
.redirectOutput(logFile)
.redirectErrorStream(true)
.start()
}
suspend fun stop() = withContext(Dispatchers.IO) {
process?.destroyForcibly()
try {
process?.waitFor()
} catch (_: Exception) {
// Process already dead or timeout
}
process = null
}
}
@@ -0,0 +1,38 @@
package com.correx.infrastructure.inference
import com.correx.core.events.types.ProviderId
import com.correx.core.inference.InferenceProvider
import com.correx.core.inference.ModelCapability
import com.correx.core.inference.ProviderHealth
import com.correx.core.inference.ProviderRegistry
import java.util.concurrent.CopyOnWriteArrayList
/**
* In-memory [ProviderRegistry] backed by a [CopyOnWriteArrayList] for safe concurrent reads.
*
* Providers supplied at construction are registered immediately.
* Additional providers may be registered at any time via [register].
*/
class DefaultProviderRegistry(
initialProviders: List<InferenceProvider> = emptyList(),
) : ProviderRegistry {
private val providers: CopyOnWriteArrayList<InferenceProvider> =
CopyOnWriteArrayList(initialProviders)
override fun register(provider: InferenceProvider) {
providers.add(provider)
}
override fun resolve(capability: ModelCapability): List<InferenceProvider> =
providers
.filter { p -> p.capabilities().any { it.capability == capability } }
.sortedByDescending { p ->
p.capabilities().firstOrNull { it.capability == capability }?.score ?: 0.0
}
override fun listAll(): List<InferenceProvider> = providers.toList()
override suspend fun healthCheckAll(): Map<ProviderId, ProviderHealth> =
providers.associate { it.id to it.healthCheck() }
}
@@ -0,0 +1,32 @@
package com.correx.infrastructure.inference
import com.correx.core.inference.InferenceProvider
import com.correx.core.inference.ModelCapability
import com.correx.core.inference.NoEligibleProviderException
import com.correx.core.inference.RoutingStrategy
import com.correx.core.events.types.StageId
/**
* Selects the first [InferenceProvider] from [candidates] whose declared capabilities
* cover all [requiredCapabilities].
*
* "First" means the list order as supplied by [ProviderRegistry.resolve] — callers that
* want score-ordered selection should pass a score-sorted list.
*
* @throws NoEligibleProviderException if no candidate satisfies all required capabilities.
*/
class FirstAvailableRoutingStrategy : RoutingStrategy {
override fun select(
candidates: List<InferenceProvider>,
requiredCapabilities: Set<ModelCapability>,
): InferenceProvider {
val declared = candidates.firstOrNull { provider ->
val providerCapabilities = provider.capabilities().map { it.capability }.toSet()
providerCapabilities.containsAll(requiredCapabilities)
}
return declared ?: throw NoEligibleProviderException(
StageId("routing"),
requiredCapabilities,
)
}
}
@@ -0,0 +1,3 @@
package com.correx.infrastructure.inference
object Module
@@ -0,0 +1,80 @@
package com.correx.infrastructure.inference
import com.correx.core.events.types.ProviderId
import com.correx.core.inference.CapabilityScore
import com.correx.core.inference.InferenceProvider
import com.correx.core.inference.InferenceRequest
import com.correx.core.inference.InferenceResponse
import com.correx.core.inference.ModelCapability
import com.correx.core.inference.ProviderHealth
import com.correx.core.inference.Tokenizer
import com.correx.core.inference.Token
import com.correx.core.utils.TypeId
import kotlinx.coroutines.test.runTest
import org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.Assertions.assertTrue
import org.junit.jupiter.api.Test
class DefaultProviderRegistryTest {
private fun fakeProvider(
id: String,
vararg caps: Pair<ModelCapability, Double>,
): InferenceProvider = object : InferenceProvider {
override val id: ProviderId = TypeId(id)
override val name: String = id
override val tokenizer: Tokenizer = object : Tokenizer {
override suspend fun tokenize(text: String): List<Token> = emptyList()
override suspend fun countTokens(text: String): Int = 0
}
override suspend fun infer(request: InferenceRequest): InferenceResponse =
error("not implemented in test")
override suspend fun healthCheck(): ProviderHealth = ProviderHealth.Healthy
override fun capabilities(): Set<CapabilityScore> =
caps.map { (cap, score) -> CapabilityScore(cap, score) }.toSet()
}
@Test
fun `listAll returns all registered providers`() {
val p1 = fakeProvider("p1", ModelCapability.Coding to 0.9)
val p2 = fakeProvider("p2", ModelCapability.General to 0.8)
val registry = DefaultProviderRegistry(listOf(p1, p2))
assertEquals(2, registry.listAll().size)
}
@Test
fun `register adds provider`() {
val registry = DefaultProviderRegistry()
val p = fakeProvider("p1", ModelCapability.Coding to 0.9)
registry.register(p)
assertEquals(1, registry.listAll().size)
}
@Test
fun `resolve returns providers with matching capability ordered by score desc`() {
val low = fakeProvider("low", ModelCapability.Coding to 0.5)
val high = fakeProvider("high", ModelCapability.Coding to 0.9)
val other = fakeProvider("other", ModelCapability.General to 0.8)
val registry = DefaultProviderRegistry(listOf(low, high, other))
val resolved = registry.resolve(ModelCapability.Coding)
assertEquals(2, resolved.size)
assertEquals("high", resolved.first().id.value)
}
@Test
fun `resolve returns empty list when no provider matches`() {
val registry = DefaultProviderRegistry(listOf(fakeProvider("p", ModelCapability.General to 1.0)))
assertTrue(registry.resolve(ModelCapability.Reasoning).isEmpty())
}
@Test
fun `healthCheckAll returns health for all providers`() = runTest {
val p1 = fakeProvider("p1", ModelCapability.Coding to 1.0)
val p2 = fakeProvider("p2", ModelCapability.General to 1.0)
val registry = DefaultProviderRegistry(listOf(p1, p2))
val health = registry.healthCheckAll()
assertEquals(2, health.size)
assertTrue(health.values.all { it is ProviderHealth.Healthy })
}
}
@@ -0,0 +1,77 @@
package com.correx.infrastructure.inference
import com.correx.core.events.types.ProviderId
import com.correx.core.inference.CapabilityScore
import com.correx.core.inference.InferenceProvider
import com.correx.core.inference.InferenceRequest
import com.correx.core.inference.InferenceResponse
import com.correx.core.inference.ModelCapability
import com.correx.core.inference.NoEligibleProviderException
import com.correx.core.inference.ProviderHealth
import com.correx.core.inference.Token
import com.correx.core.inference.Tokenizer
import com.correx.core.utils.TypeId
import org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.Test
import org.junit.jupiter.api.assertThrows
class FirstAvailableRoutingStrategyTest {
private val strategy = FirstAvailableRoutingStrategy()
private fun fakeProvider(id: String, vararg caps: ModelCapability): InferenceProvider =
object : InferenceProvider {
override val id: ProviderId = TypeId(id)
override val name: String = id
override val tokenizer: Tokenizer = object : Tokenizer {
override suspend fun tokenize(text: String): List<Token> = emptyList()
override suspend fun countTokens(text: String): Int = 0
}
override suspend fun infer(request: InferenceRequest): InferenceResponse =
error("not implemented in test")
override suspend fun healthCheck(): ProviderHealth = ProviderHealth.Healthy
override fun capabilities(): Set<CapabilityScore> =
caps.map { CapabilityScore(it, 1.0) }.toSet()
}
@Test
fun `selects first provider satisfying all required capabilities`() {
val coding = fakeProvider("codingOnly", ModelCapability.Coding)
val general = fakeProvider("generalAndCoding", ModelCapability.General, ModelCapability.Coding)
val required = setOf(ModelCapability.Coding, ModelCapability.General)
val selected = strategy.select(listOf(coding, general), required)
assertEquals("generalAndCoding", selected.id.value)
}
@Test
fun `selects first matching provider when multiple satisfy requirements`() {
val p1 = fakeProvider("p1", ModelCapability.Coding)
val p2 = fakeProvider("p2", ModelCapability.Coding)
val selected = strategy.select(listOf(p1, p2), setOf(ModelCapability.Coding))
assertEquals("p1", selected.id.value)
}
@Test
fun `throws NoEligibleProviderException when no candidate satisfies required capabilities`() {
val general = fakeProvider("general", ModelCapability.General)
assertThrows<NoEligibleProviderException> {
strategy.select(listOf(general), setOf(ModelCapability.Coding))
}
}
@Test
fun `throws NoEligibleProviderException when candidates list is empty`() {
assertThrows<NoEligibleProviderException> {
strategy.select(emptyList(), setOf(ModelCapability.Coding))
}
}
@Test
fun `selects any provider when required capabilities is empty`() {
val p = fakeProvider("p", ModelCapability.General)
val selected = strategy.select(listOf(p), emptySet())
assertEquals("p", selected.id.value)
}
}
+13
View File
@@ -0,0 +1,13 @@
plugins {
id 'java-library'
id 'org.jetbrains.kotlin.jvm'
id 'org.jetbrains.kotlin.plugin.serialization'
}
dependencies {
implementation(project(":core:events"))
implementation(project(":core:sessions"))
implementation "org.xerial:sqlite-jdbc"
testImplementation(testFixtures(project(":testing:contracts")))
testImplementation "org.junit.jupiter:junit-jupiter"
}
@@ -0,0 +1,68 @@
package com.correx.infrastructure.persistence
import com.correx.core.events.events.NewEvent
import com.correx.core.events.events.StoredEvent
import com.correx.core.events.stores.EventStore
import com.correx.core.events.types.EventId
import com.correx.core.events.types.SessionId
import java.util.concurrent.*
import java.util.concurrent.atomic.*
class InMemoryEventStore : EventStore {
private val streams = ConcurrentHashMap<SessionId, MutableList<StoredEvent>>()
private val sequences = ConcurrentHashMap<SessionId, AtomicLong>()
private val seenEventIds = ConcurrentHashMap.newKeySet<EventId>()
override fun append(event: NewEvent): StoredEvent {
val stream = streams.computeIfAbsent(event.metadata.sessionId) { mutableListOf() }
synchronized(stream) {
if (!seenEventIds.add(event.metadata.eventId)) {
error("duplicate event_id: ${event.metadata.eventId}")
}
return doAppend(event, stream)
}
}
override fun appendAll(events: List<NewEvent>): List<StoredEvent> {
if (events.isEmpty()) return emptyList()
val sessionId = events.first().metadata.sessionId
val stream = streams.computeIfAbsent(sessionId) { mutableListOf() }
synchronized(stream) {
return events.map { doAppend(it, stream) }
}
}
override fun read(sessionId: SessionId): List<StoredEvent> {
return streams[sessionId]
?.sortedBy { it.sequence }
?: emptyList()
}
override fun readFrom(sessionId: SessionId, fromSequence: Long): List<StoredEvent> {
return streams[sessionId]
?.filter { it.sequence > fromSequence }
?.toList()
?: emptyList()
}
override fun lastSequence(sessionId: SessionId): Long? {
return sequences[sessionId]?.get()
}
private fun doAppend(event: NewEvent, stream: MutableList<StoredEvent>): StoredEvent {
val seq = sequences.computeIfAbsent(event.metadata.sessionId) { AtomicLong(0) }
.incrementAndGet()
val stored = StoredEvent(metadata = event.metadata, sequence = seq, payload = event.payload)
// safety: enforce ordering invariant
if (stream.isNotEmpty() && stream.last().sequence >= stored.sequence) {
error("sequence violation for session ${event.metadata.sessionId}")
}
stream.add(stored)
return stored
}
}
@@ -0,0 +1,215 @@
package com.correx.infrastructure.persistence
import com.correx.core.events.events.EventMetadata
import com.correx.core.events.events.NewEvent
import com.correx.core.events.events.StoredEvent
import com.correx.core.events.serialization.JsonEventSerializer
import com.correx.core.events.serialization.eventJson
import com.correx.core.events.stores.EventStore
import com.correx.core.events.types.CausationId
import com.correx.core.events.types.CorrelationId
import com.correx.core.events.types.EventId
import com.correx.core.events.types.SessionId
import com.correx.infrastructure.persistence.util.JDBCHelper.transaction
import kotlinx.datetime.Instant
import java.sql.Connection
import java.sql.ResultSet
class SqliteEventStore(
private val connection: Connection,
private val jsonSerializer: JsonEventSerializer = JsonEventSerializer(eventJson)
) : EventStore {
init {
connection.createStatement().use { stmt ->
stmt.execute(
"""
CREATE TABLE IF NOT EXISTS events (
event_id TEXT PRIMARY KEY,
session_id TEXT NOT NULL,
sequence INTEGER NOT NULL,
timestamp TEXT NOT NULL,
schema_version INTEGER NOT NULL,
causation_id TEXT,
correlation_id TEXT,
payload TEXT NOT NULL
);
""".trimIndent()
)
}
}
override fun append(event: NewEvent): StoredEvent {
return connection.transaction {
val existing = findByEventId(event.metadata.eventId)
check(existing == null) { "duplicate event_id: ${event.metadata.eventId}" }
val seq = nextSequence(event.metadata.sessionId)
val stored = StoredEvent(
metadata = event.metadata,
sequence = seq,
payload = event.payload
)
insert(stored)
stored
}
}
override fun appendAll(events: List<NewEvent>): List<StoredEvent> {
if (events.isEmpty()) return emptyList()
val sessionId = events.first().metadata.sessionId
return connection.transaction {
events.map { event ->
val existing = findByEventId(event.metadata.eventId)
check(existing == null) { "duplicate event_id: ${event.metadata.eventId}" }
val seq = nextSequence(sessionId)
val stored = StoredEvent(
metadata = event.metadata,
sequence = seq,
payload = event.payload
)
insert(stored)
stored
}
}
}
override fun read(sessionId: SessionId): List<StoredEvent> =
connection.prepareStatement(
"""
SELECT * FROM events
WHERE session_id = ?
ORDER BY sequence ASC
""".trimIndent()
).use { ps ->
ps.setString(1, sessionId.value)
ps.executeQuery().use { rs ->
buildList {
while (rs.next()) {
add(map(rs))
}
}
}
}
override fun readFrom(sessionId: SessionId, fromSequence: Long): List<StoredEvent> =
connection.prepareStatement(
"""
SELECT * FROM events
WHERE session_id = ?
AND sequence > ?
ORDER BY sequence ASC
""".trimIndent()
).use { ps ->
ps.setString(1, sessionId.value)
ps.setLong(2, fromSequence)
ps.executeQuery().use { rs ->
buildList {
while (rs.next()) {
add(map(rs))
}
}
}
}
override fun lastSequence(sessionId: SessionId): Long? =
connection.prepareStatement(
"""
SELECT MAX(sequence)
FROM events
WHERE session_id = ?
""".trimIndent()
).use { ps ->
ps.setString(1, sessionId.value)
ps.executeQuery().use { rs ->
if (rs.next()) rs.getLong(1).takeIf { !rs.wasNull() }
else null
}
}
// ---------- helpers ----------
private fun nextSequence(sessionId: SessionId): Long =
connection.prepareStatement(
"""
SELECT COALESCE(MAX(sequence), 0)
FROM events
WHERE session_id = ?
""".trimIndent()
).use { ps ->
ps.setString(1, sessionId.value)
ps.executeQuery().use { rs ->
rs.next()
rs.getLong(1) + 1
}
}
@Suppress("MagicNumber")
private fun insert(event: StoredEvent) {
connection.prepareStatement(
"""
INSERT INTO events (
event_id,
session_id,
sequence,
timestamp,
schema_version,
causation_id,
correlation_id,
payload
) VALUES (?, ?, ?, ?, ?, ?, ?, ?)
""".trimIndent()
).use { ps ->
ps.setString(1, event.metadata.eventId.value)
ps.setString(2, event.metadata.sessionId.value)
ps.setLong(3, event.sequence)
ps.setString(4, event.metadata.timestamp.toString())
ps.setInt(5, event.metadata.schemaVersion)
ps.setString(6, event.metadata.causationId?.value)
ps.setString(7, event.metadata.correlationId?.value)
ps.setString(8, jsonSerializer.serialize(event.payload))
ps.executeUpdate()
}
}
private fun findByEventId(eventId: EventId): StoredEvent? =
connection.prepareStatement(
"""
SELECT * FROM events
WHERE event_id = ?
""".trimIndent()
).use { ps ->
ps.setString(1, eventId.value)
ps.executeQuery().use { rs ->
if (rs.next()) map(rs) else null
}
}
private fun map(rs: ResultSet): StoredEvent =
StoredEvent(
metadata = EventMetadata(
eventId = EventId(rs.getString("event_id")),
sessionId = SessionId(rs.getString("session_id")),
timestamp = Instant.parse(rs.getString("timestamp")),
schemaVersion = rs.getInt("schema_version"),
causationId = rs.getString("causation_id")?.let { CausationId(it) },
correlationId = rs.getString("correlation_id")?.let { CorrelationId(it) }
),
sequence = rs.getLong("sequence"),
payload = jsonSerializer.deserialize(rs.getString("payload"))
)
}
@@ -0,0 +1,23 @@
package com.correx.infrastructure.persistence.util
import java.sql.Connection
import java.sql.SQLException
object JDBCHelper {
inline fun <T> Connection.transaction(block: () -> T): T {
val oldAutoCommit = autoCommit
try {
autoCommit = false
val result = block()
commit()
return result
} catch (e: SQLException) {
rollback()
throw e
} finally {
autoCommit = oldAutoCommit
}
}
}
@@ -0,0 +1,8 @@
package com.correx.infrastructure.persistence
import com.correx.testing.contracts.fixtures.events.store.EventStoreContractTest
import com.correx.core.events.stores.EventStore
class InMemoryEventStoreTest : EventStoreContractTest() {
override fun store(): EventStore = InMemoryEventStore()
}
@@ -0,0 +1,23 @@
package com.correx.infrastructure.persistence
import com.correx.core.events.stores.EventStore
import com.correx.core.sessions.projections.replay.DefaultEventReplayer
import com.correx.core.sessions.projections.replay.EventReplayer
import com.correx.core.sessions.SessionCounterProjection
import com.correx.core.sessions.SessionCounterState
import com.correx.testing.contracts.fixtures.projections.SessionReplayContractTest
class InMemoryReplayTest : SessionReplayContractTest() {
override fun store(): EventStore =
InMemoryEventStore()
override fun replayer(
store: EventStore
): EventReplayer<SessionCounterState> {
return DefaultEventReplayer(
store,
SessionCounterProjection("s1")
)
}
}
@@ -0,0 +1,12 @@
package com.correx.infrastructure.persistence
import com.correx.testing.contracts.fixtures.events.store.EventStoreContractTest
import com.correx.core.events.stores.EventStore
import java.sql.DriverManager
class SqliteEventStoreTest : EventStoreContractTest() {
override fun store(): EventStore {
val conn = DriverManager.getConnection("jdbc:sqlite::memory:")
return SqliteEventStore(conn)
}
}
@@ -0,0 +1,28 @@
package com.correx.infrastructure.persistence
import com.correx.core.events.stores.EventStore
import com.correx.core.sessions.projections.replay.DefaultEventReplayer
import com.correx.core.sessions.projections.replay.EventReplayer
import com.correx.core.sessions.SessionCounterProjection
import com.correx.core.sessions.SessionCounterState
import com.correx.testing.contracts.fixtures.projections.SessionReplayContractTest
import java.sql.DriverManager
class SqliteReplayTest : SessionReplayContractTest() {
override fun store(): EventStore {
val connection =
DriverManager.getConnection("jdbc:sqlite::memory:")
return SqliteEventStore(connection)
}
override fun replayer(
store: EventStore
): EventReplayer<SessionCounterState> {
return DefaultEventReplayer(
store,
SessionCounterProjection("s1")
)
}
}
+5
View File
@@ -0,0 +1,5 @@
plugins {
id 'java-library'
id 'org.jetbrains.kotlin.jvm'
id 'org.jetbrains.kotlin.plugin.serialization'
}
@@ -0,0 +1,3 @@
package com.correx.infrastructure.scheduler
object Module
+5
View File
@@ -0,0 +1,5 @@
plugins {
id 'java-library'
id 'org.jetbrains.kotlin.jvm'
id 'org.jetbrains.kotlin.plugin.serialization'
}
@@ -0,0 +1,3 @@
package com.correx.infrastructure.security
object Module
@@ -0,0 +1,68 @@
package com.correx.infrastructure
import com.correx.core.approvals.domain.ApprovalEngine
import com.correx.core.events.EventDispatcher
import com.correx.core.events.stores.EventStore
import com.correx.core.inference.DefaultInferenceRouter
import com.correx.core.inference.InferenceProvider
import com.correx.core.inference.InferenceRouter
import com.correx.core.inference.ProviderRegistry
import com.correx.core.inference.RoutingStrategy
import com.correx.infrastructure.inference.DefaultProviderRegistry
import com.correx.infrastructure.inference.FirstAvailableRoutingStrategy
import com.correx.core.tools.contract.ToolExecutor
import com.correx.core.tools.registry.ToolRegistry
import com.correx.infrastructure.inference.commons.ModelManager
import com.correx.infrastructure.inference.llama.cpp.DefaultModelManager
import com.correx.infrastructure.persistence.SqliteEventStore
import com.correx.infrastructure.tools.DefaultToolRegistry
import com.correx.infrastructure.tools.DispatchingToolExecutor
import com.correx.infrastructure.tools.SandboxedToolExecutor
import com.correx.infrastructure.tools.ToolConfig
import com.correx.infrastructure.tools.buildTools
import com.correx.infrastructure.tools.filesystem.FileEditTool
import com.correx.infrastructure.tools.filesystem.FileReadTool
import com.correx.infrastructure.tools.filesystem.FileWriteTool
import com.correx.infrastructure.tools.shell.ShellTool
import io.ktor.client.HttpClient
import java.nio.file.Path
import java.sql.DriverManager
object InfrastructureModule {
fun createEventStore(dbPath: String): EventStore =
SqliteEventStore(DriverManager.getConnection("jdbc:sqlite:$dbPath"))
fun createModelManager(
eventStore: EventStore,
httpClient: HttpClient,
eventDispatcher: EventDispatcher,
): ModelManager = DefaultModelManager(
eventStore = eventStore,
httpClient = httpClient,
eventDispatcher = eventDispatcher,
)
fun createToolRegistry(config: ToolConfig): ToolRegistry =
DefaultToolRegistry.build(config.buildTools())
fun createInferenceRouter(
providers: List<InferenceProvider> = emptyList(),
registry: ProviderRegistry = DefaultProviderRegistry(providers),
strategy: RoutingStrategy = FirstAvailableRoutingStrategy(),
): InferenceRouter = DefaultInferenceRouter(registry, strategy)
fun createToolExecutor(
registry: ToolRegistry,
approvalEngine: ApprovalEngine,
eventStore: EventStore,
eventDispatcher: EventDispatcher,
workDir: Path,
): ToolExecutor = SandboxedToolExecutor(
delegate = DispatchingToolExecutor(registry),
registry = registry,
approvalEngine = approvalEngine,
eventStore = eventStore,
eventDispatcher = eventDispatcher,
workDir = workDir,
)
}
+5
View File
@@ -0,0 +1,5 @@
plugins {
id 'java-library'
id 'org.jetbrains.kotlin.jvm'
id 'org.jetbrains.kotlin.plugin.serialization'
}
@@ -0,0 +1,3 @@
package com.correx.infrastructure.telemetry
object Module
+13
View File
@@ -0,0 +1,13 @@
plugins {
id 'java-library'
id 'org.jetbrains.kotlin.jvm'
id 'org.jetbrains.kotlin.plugin.serialization'
}
dependencies {
implementation project(":core:tools")
implementation project(":core:events")
implementation project(":core:approvals")
implementation project(":core:sessions")
implementation project(":infrastructure:tools:filesystem")
}
@@ -0,0 +1,11 @@
plugins {
id 'java-library'
id 'org.jetbrains.kotlin.jvm'
id 'org.jetbrains.kotlin.plugin.serialization'
}
dependencies {
implementation project(':core:tools')
implementation project(':core:events')
implementation project(':core:approvals')
}
@@ -0,0 +1,231 @@
package com.correx.infrastructure.tools.filesystem
import com.correx.core.approvals.Tier
import com.correx.core.events.events.ToolRequest
import com.correx.core.events.types.ToolInvocationId
import com.correx.core.tools.contract.FileAffectingTool
import com.correx.core.tools.contract.Tool
import com.correx.core.tools.contract.ToolCapability
import com.correx.core.tools.contract.ToolExecutor
import com.correx.core.tools.contract.ToolResult
import com.correx.core.tools.contract.ValidationResult
import kotlinx.coroutines.CancellationException
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext
import java.io.IOException
import java.nio.file.Files
import java.nio.file.InvalidPathException
import java.nio.file.Path
import java.nio.file.Paths
import java.nio.file.StandardOpenOption
@Suppress("TooManyFunctions")
class FileEditTool(
allowedPaths: Set<Path> = emptySet(),
) : Tool, FileAffectingTool, ToolExecutor {
private val normalizedAllowedPaths: Set<Path> = allowedPaths.map { it.normalize().toAbsolutePath() }.toSet()
override val name: String = "file_edit"
override val tier: Tier = Tier.T3
override val requiredCapabilities: Set<ToolCapability> = setOf(ToolCapability.FILE_WRITE)
override fun affectedPaths(request: ToolRequest): Set<Path> {
val pathString = request.parameters["path"] as? String ?: return emptySet()
return setOf(Paths.get(pathString).normalize().toAbsolutePath())
}
override fun validateRequest(request: ToolRequest): ValidationResult {
val operation = request.parameters["operation"] as? String
val pathString = request.parameters["path"] as? String
return when {
operation == null ->
ValidationResult.Invalid("Missing 'operation' parameter. Expected 'append', 'replace', or 'patch'.")
pathString == null ->
ValidationResult.Invalid("Missing 'path' parameter. Expected String.")
else -> runCatching {
val path = Paths.get(pathString).normalize().toAbsolutePath()
checkPathAllowed(path, pathString, operation, request)
}.getOrElse { e ->
mapExceptionToValidationResult(e)
}
}
}
private fun checkPathAllowed(
path: Path,
pathString: String,
operation: String,
request: ToolRequest,
): ValidationResult {
return when {
normalizedAllowedPaths.isEmpty() -> ValidationResult.Invalid("No paths are allowed.")
!isPathAllowed(path) ->
ValidationResult.Invalid("Path '$pathString' is not in the allowed list.")
!Files.exists(path) ->
ValidationResult.Invalid("File not found: $pathString")
operation == "append" && !request.parameters.containsKey("content") ->
ValidationResult.Invalid("Missing 'content' parameter for 'append' operation.")
operation == "replace" && (
!request.parameters.containsKey("target") ||
!request.parameters.containsKey("replacement")
) ->
ValidationResult.Invalid("Missing 'target' or 'replacement' parameter for 'replace' operation.")
operation == "patch" && !request.parameters.containsKey("patch") ->
ValidationResult.Invalid("Missing 'patch' parameter for 'patch' operation.")
operation !in listOf("append", "replace", "patch") ->
ValidationResult.Invalid("Unknown operation: $operation")
else -> ValidationResult.Valid
}
}
private fun isPathAllowed(path: Path): Boolean =
normalizedAllowedPaths.any { allowedPath ->
path.startsWith(allowedPath)
}
private fun mapExceptionToValidationResult(e: Throwable): ValidationResult =
when (e) {
is InvalidPathException -> ValidationResult.Invalid("Invalid path format: ${e.message}")
is IOException -> ValidationResult.Invalid("IO error: ${e.message}")
is SecurityException -> ValidationResult.Invalid("Security error: ${e.message}")
else -> ValidationResult.Invalid(e.message ?: "Unknown error occurred")
}
override suspend fun execute(request: ToolRequest): ToolResult = withContext(Dispatchers.IO) {
val validation = validateRequest(request)
if (validation is ValidationResult.Invalid) {
return@withContext ToolResult.Failure(
invocationId = request.invocationId,
reason = validation.reason,
recoverable = false,
)
}
val operation = request.parameters["operation"] as String
val pathString = request.parameters["path"] as String
val path = Paths.get(pathString).normalize().toAbsolutePath()
runCatching {
performOperation(operation, path, pathString, request)
}.getOrElse { e ->
handleExecutionException(e, request.invocationId, pathString)
}
}
private fun performOperation(
operation: String,
path: Path,
pathString: String,
request: ToolRequest,
): ToolResult = when (operation) {
"append" -> append(path, request)
"replace" -> replace(path, pathString, request)
"patch" -> patch(path, pathString, request)
else -> ToolResult.Failure(
invocationId = request.invocationId,
reason = "Unknown operation: $operation",
recoverable = false,
)
}
private fun append(path: Path, request: ToolRequest): ToolResult {
val content = request.parameters["content"] as String
Files.writeString(path, content, StandardOpenOption.APPEND)
return ToolResult.Success(
invocationId = request.invocationId,
output = "Content appended to ${path.toAbsolutePath()}",
)
}
private fun replace(path: Path, pathString: String, request: ToolRequest): ToolResult {
val target = request.parameters["target"] as String
val replacement = request.parameters["replacement"] as String
val currentContent = Files.readString(path)
val occurrences = currentContent.split(target).size - 1
return if (occurrences == 1) {
val newContent = currentContent.replace(target, replacement)
Files.writeString(path, newContent)
ToolResult.Success(
invocationId = request.invocationId,
output = "Target replaced in $pathString",
)
} else {
ToolResult.Failure(
invocationId = request.invocationId,
reason = if (occurrences == 0) {
"Target '$target' not found in $pathString"
} else {
"Target '$target' found $occurrences times, expected exactly once"
},
recoverable = false,
)
}
}
private fun patch(path: Path, pathString: String, request: ToolRequest): ToolResult {
val patchContent = request.parameters["patch"] as String
val parentDir = path.parent ?: Paths.get(".")
val fileName = path.fileName.toString()
val process = ProcessBuilder("patch", "-p1", fileName)
.directory(parentDir.toFile())
.redirectErrorStream(true)
.start()
return try {
process.outputStream.use { it.write(patchContent.toByteArray()) }
val output = process.inputStream.bufferedReader().use { it.readText() }
val exitCode = process.waitFor()
if (exitCode == 0) {
ToolResult.Success(
invocationId = request.invocationId,
output = "Patch applied successfully to $pathString: $output",
)
} else {
ToolResult.Failure(
invocationId = request.invocationId,
reason = "Patch failed with exit code $exitCode: $output",
recoverable = false,
)
}
} finally {
process.destroy()
}
}
private fun handleExecutionException(
e: Throwable,
invocationId: ToolInvocationId,
pathString: String,
): ToolResult = when (e) {
is CancellationException -> throw e
is IOException -> ToolResult.Failure(
invocationId = invocationId,
reason = "IO error: ${e.message}",
recoverable = false,
)
is SecurityException -> ToolResult.Failure(
invocationId = invocationId,
reason = "Access denied: $pathString, ${e.message}",
recoverable = false,
)
else -> ToolResult.Failure(
invocationId = invocationId,
reason = e.message ?: "Unknown error occurred",
recoverable = false,
)
}
}
@@ -0,0 +1,97 @@
package com.correx.infrastructure.tools.filesystem
import com.correx.core.approvals.Tier
import com.correx.core.events.events.ToolRequest
import com.correx.core.tools.contract.Tool
import com.correx.core.tools.contract.ToolCapability
import com.correx.core.tools.contract.ToolExecutor
import com.correx.core.tools.contract.ToolResult
import com.correx.core.tools.contract.ValidationResult
import kotlinx.coroutines.CancellationException
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext
import java.io.IOException
import java.nio.file.Files
import java.nio.file.InvalidPathException
import java.nio.file.Path
import java.nio.file.Paths
class FileReadTool(
allowedPaths: Set<Path> = emptySet(),
) : Tool, ToolExecutor {
private val normalizedAllowedPaths: Set<Path> = allowedPaths.map { it.normalize().toAbsolutePath() }.toSet()
override val name: String = "file_read"
override val tier: Tier = Tier.T1
override val requiredCapabilities: Set<ToolCapability> = setOf(ToolCapability.FILE_READ)
override fun validateRequest(request: ToolRequest): ValidationResult =
(request.parameters["path"] as? String)?.let { pathString ->
runCatching {
val path = Paths.get(pathString).normalize().toAbsolutePath()
val isAllowed = normalizedAllowedPaths.isEmpty() || path in normalizedAllowedPaths
|| normalizedAllowedPaths.any { path.startsWith(it) }
ValidationResult.Valid.takeIf { isAllowed }
?: ValidationResult.Invalid("Path '$pathString' is not in the allowed list.")
}.getOrElse {
when (it) {
is InvalidPathException -> ValidationResult.Invalid("Invalid path format: ${it.message}")
is IOException -> ValidationResult.Invalid("IO error: ${it.message}")
is SecurityException -> ValidationResult.Invalid("Security error: ${it.message}")
else -> ValidationResult.Invalid(it.message ?: "Unknown error occured")
}
}
} ?: ValidationResult.Invalid("Missing or invalid 'path' parameter. Expected String.")
override suspend fun execute(request: ToolRequest): ToolResult = withContext(Dispatchers.IO) {
validateRequest(request).run {
takeIf { it is ValidationResult.Valid }?.let {
val pathString = request.parameters["path"] as String
val path = Paths.get(pathString).normalize().toAbsolutePath()
if (Files.exists(path)) {
readStringCatching(request, path, pathString)
} else {
ToolResult.Failure(
invocationId = request.invocationId,
reason = "File not found: $pathString",
recoverable = false,
)
}
} ?: ToolResult.Failure(
invocationId = request.invocationId,
reason = (this as ValidationResult.Invalid).reason,
recoverable = false,
)
}
}
private fun readStringCatching(request: ToolRequest, path: Path, pathString: String): ToolResult =
runCatching {
ToolResult.Success(
invocationId = request.invocationId,
output = Files.readString(path),
)
}.getOrElse {
when (it) {
is CancellationException -> throw it
is SecurityException -> ToolResult.Failure(
invocationId = request.invocationId,
reason = "Access denied: $pathString, ${it.message}",
recoverable = false,
)
is IOException -> ToolResult.Failure(
invocationId = request.invocationId,
reason = "IO error: ${it.message}",
recoverable = false,
)
else -> ToolResult.Failure(
invocationId = request.invocationId,
reason = it.message ?: "Unknown error occurred while reading file",
recoverable = false,
)
}
}
}
@@ -0,0 +1,172 @@
package com.correx.infrastructure.tools.filesystem
import com.correx.core.approvals.Tier
import com.correx.core.events.events.ToolRequest
import com.correx.core.events.types.ToolInvocationId
import com.correx.core.tools.contract.FileAffectingTool
import com.correx.core.tools.contract.Tool
import com.correx.core.tools.contract.ToolCapability
import com.correx.core.tools.contract.ToolExecutor
import com.correx.core.tools.contract.ToolResult
import com.correx.core.tools.contract.ValidationResult
import kotlinx.coroutines.CancellationException
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext
import java.io.IOException
import java.nio.file.Files
import java.nio.file.InvalidPathException
import java.nio.file.Path
import java.nio.file.Paths
class FileWriteTool(
allowedPaths: Set<Path> = emptySet(),
) : Tool, FileAffectingTool, ToolExecutor {
private val normalizedAllowedPaths: Set<Path> = allowedPaths.map { it.normalize().toAbsolutePath() }.toSet()
override val name: String = "file_write"
override val tier: Tier = Tier.T2
override val requiredCapabilities: Set<ToolCapability> = setOf(ToolCapability.FILE_WRITE)
override fun affectedPaths(request: ToolRequest): Set<Path> {
val pathString = request.parameters["path"] as? String ?: return emptySet()
return setOf(Paths.get(pathString).normalize().toAbsolutePath())
}
override fun validateRequest(request: ToolRequest): ValidationResult {
val operation = request.parameters["operation"] as? String
val pathString = request.parameters["path"] as? String
return when {
operation == null ->
ValidationResult.Invalid("Missing 'operation' parameter. Expected 'write' or 'delete'.")
pathString == null ->
ValidationResult.Invalid("Missing 'path' parameter. Expected String.")
else -> runCatching {
val path = Paths.get(pathString).normalize().toAbsolutePath()
checkPathAllowed(path, pathString, operation, request)
}.getOrElse { e ->
mapExceptionToValidationResult(e)
}
}
}
private fun checkPathAllowed(
path: Path,
pathString: String,
operation: String,
request: ToolRequest,
): ValidationResult {
return when {
normalizedAllowedPaths.isEmpty() -> ValidationResult.Invalid("No paths are allowed.")
!isPathAllowed(path) ->
ValidationResult.Invalid("Path '$pathString' is not in the allowed list.")
operation == "write" && !request.parameters.containsKey("content") ->
ValidationResult.Invalid("Missing 'content' parameter for 'write' operation.")
operation != "write" && operation != "delete" ->
ValidationResult.Invalid("Unknown operation: $operation")
else -> ValidationResult.Valid
}
}
private fun isPathAllowed(path: Path): Boolean =
normalizedAllowedPaths.any { allowedPath ->
path.startsWith(allowedPath)
}
private fun mapExceptionToValidationResult(e: Throwable): ValidationResult =
when (e) {
is InvalidPathException -> ValidationResult.Invalid("Invalid path format: ${e.message}")
is IOException -> ValidationResult.Invalid("IO error: ${e.message}")
is SecurityException -> ValidationResult.Invalid("Security error: ${e.message}")
else -> ValidationResult.Invalid(e.message ?: "Unknown error occurred")
}
override suspend fun execute(request: ToolRequest): ToolResult = withContext(Dispatchers.IO) {
val validation = validateRequest(request)
if (validation is ValidationResult.Invalid) {
return@withContext ToolResult.Failure(
invocationId = request.invocationId,
reason = validation.reason,
recoverable = false,
)
}
val operation = request.parameters["operation"] as String
val pathString = request.parameters["path"] as String
val path = Paths.get(pathString).normalize().toAbsolutePath()
runCatching {
performOperation(operation, path, pathString, request)
}.getOrElse { e ->
handleExecutionException(e, request.invocationId, pathString)
}
}
private fun performOperation(
operation: String,
path: Path,
pathString: String,
request: ToolRequest,
): ToolResult = when (operation) {
"write" -> {
val content = request.parameters["content"] as String
Files.writeString(path, content)
ToolResult.Success(
invocationId = request.invocationId,
output = "File written successfully to $pathString",
)
}
"delete" -> {
if (Files.exists(path)) {
Files.deleteIfExists(path)
ToolResult.Success(
invocationId = request.invocationId,
output = "File deleted successfully: $pathString",
)
} else {
ToolResult.Failure(
invocationId = request.invocationId,
reason = "File not found: $pathString",
recoverable = false,
)
}
}
else -> ToolResult.Failure(
invocationId = request.invocationId,
reason = "Unknown operation: $operation",
recoverable = false,
)
}
private fun handleExecutionException(
e: Throwable,
invocationId: ToolInvocationId,
pathString: String,
): ToolResult = when (e) {
is CancellationException -> throw e
is IOException -> ToolResult.Failure(
invocationId = invocationId,
reason = "IO error: ${e.message}",
recoverable = false,
)
is SecurityException -> ToolResult.Failure(
invocationId = invocationId,
reason = "Access denied: $pathString, ${e.message}",
recoverable = false,
)
else -> ToolResult.Failure(
invocationId = invocationId,
reason = e.message ?: "Unknown error occurred",
recoverable = false,
)
}
}
@@ -0,0 +1,234 @@
package com.correx.infrastructure.tools.filesystem
import com.correx.core.events.types.SessionId
import com.correx.core.events.types.StageId
import com.correx.core.events.types.ToolInvocationId
import com.correx.core.events.events.ToolRequest
import com.correx.core.tools.contract.ToolResult
import com.correx.core.tools.contract.ValidationResult
import kotlinx.coroutines.runBlocking
import org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.Assertions.assertFalse
import org.junit.jupiter.api.Assertions.assertTrue
import org.junit.jupiter.api.Test
import java.nio.file.Files
import java.nio.file.Path
import java.util.UUID
class FileEditToolTest {
private val sessionId = SessionId(UUID.randomUUID().toString())
private val stageId = StageId(UUID.randomUUID().toString())
private val invocationId = ToolInvocationId(UUID.randomUUID().toString())
private fun createRequest(
parameters: Map<String, String>,
toolName: String = "file_edit",
): ToolRequest {
return ToolRequest(
invocationId = invocationId,
sessionId = sessionId,
stageId = stageId,
toolName = toolName,
parameters = parameters,
)
}
@Test
fun `execute patch success`() = runBlocking {
val tempDir = Files.createTempDirectory("file_edit_patch")
val filePath = tempDir.resolve("test.txt")
Files.writeString(filePath, "line1\nline2\nline3\n")
val tool = FileEditTool(allowedPaths = setOf(tempDir))
// Patch to replace line2 with line2_modified
val patchContent = """
--- a/test.txt
+++ b/test.txt
@@ -2 +2 @@
-line2
+line2_modified
""".trimIndent() + "\n"
val request = createRequest(
mapOf(
"operation" to "patch",
"path" to filePath.toString(),
"patch" to patchContent,
),
)
val result = tool.execute(request)
assertTrue(result is ToolResult.Success, "Expected Success but got $result")
val content = Files.readString(filePath)
assertEquals("line1\nline2_modified\nline3\n", content)
}
@Test
fun `validateRequest returns Invalid for disallowed path`() = runBlocking {
val tempDir = Files.createTempDirectory("file_edit_test")
val otherDir = Files.createTempDirectory("other_dir")
val tool = FileEditTool(allowedPaths = setOf(tempDir))
val request = createRequest(
mapOf(
"operation" to "append",
"path" to otherDir.resolve("file.txt").toString(),
"content" to "hello",
),
)
val result = tool.validateRequest(request)
assertTrue(result is ValidationResult.Invalid)
assertTrue((result as ValidationResult.Invalid).reason.contains("is not in the allowed list"))
}
@Test
fun `validateRequest returns Invalid for non-existent file`() = runBlocking {
val tempDir = Files.createTempDirectory("file_edit_test")
val tool = FileEditTool(allowedPaths = setOf(tempDir))
val request = createRequest(
mapOf(
"operation" to "append",
"path" to tempDir.resolve("non_existent.txt").toString(),
"content" to "hello",
),
)
val result = tool.validateRequest(request)
assertTrue(result is ValidationResult.Invalid)
assertTrue((result as ValidationResult.Invalid).reason.contains("File not found"))
}
@Test
fun `validateRequest returns Invalid for missing parameters`() = runBlocking {
val tempDir = Files.createTempDirectory("file_edit_test")
val filePath = tempDir.resolve("test.txt")
Files.writeString(filePath, "content")
val tool = FileEditTool(allowedPaths = setOf(tempDir))
// Missing operation
val req1 = createRequest(mapOf("path" to filePath.toString()))
assertTrue(tool.validateRequest(req1) is ValidationResult.Invalid)
// Missing content for append
val req2 = createRequest(mapOf("operation" to "append", "path" to filePath.toString()))
assertTrue(tool.validateRequest(req2) is ValidationResult.Invalid)
// Missing target/replacement for replace
val req3 = createRequest(mapOf("operation" to "replace", "path" to filePath.toString(), "target" to "content"))
assertTrue(tool.validateRequest(req3) is ValidationResult.Invalid)
// Missing patch for patch
val req4 = createRequest(mapOf("operation" to "patch", "path" to filePath.toString()))
assertTrue(tool.validateRequest(req4) is ValidationResult.Invalid)
}
@Test
fun `execute append success`() = runBlocking {
val tempDir = Files.createTempDirectory("file_edit_append")
val filePath = tempDir.resolve("test.txt")
Files.writeString(filePath, "hello")
val tool = FileEditTool(allowedPaths = setOf(tempDir))
val request = createRequest(
mapOf(
"operation" to "append",
"path" to filePath.toString(),
"content" to " world",
),
)
val result = tool.execute(request)
assertTrue(result is ToolResult.Success)
assertEquals("hello world", Files.readString(filePath))
}
@Test
fun `execute replace success`() = runBlocking {
val tempDir = Files.createTempDirectory("file_edit_replace")
val filePath = tempDir.resolve("test.txt")
Files.writeString(filePath, "the quick brown fox")
val tool = FileEditTool(allowedPaths = setOf(tempDir))
val request = createRequest(
mapOf(
"operation" to "replace",
"path" to filePath.toString(),
"target" to "brown",
"replacement" to "red",
),
)
val result = tool.execute(request)
assertTrue(result is ToolResult.Success)
assertEquals("the quick red fox", Files.readString(filePath))
}
@Test
fun `execute replace failure zero matches`() = runBlocking {
val tempDir = Files.createTempDirectory("file_edit_replace_fail")
val filePath = tempDir.resolve("test.txt")
Files.writeString(filePath, "the quick brown fox")
val tool = FileEditTool(allowedPaths = setOf(tempDir))
val request = createRequest(
mapOf(
"operation" to "replace",
"path" to filePath.toString(),
"target" to "green",
"replacement" to "red",
),
)
val result = tool.execute(request)
assertTrue(result is ToolResult.Failure)
val failure = result as ToolResult.Failure
assertTrue(failure.reason.contains("not found"))
}
@Test
fun `execute replace failure multiple matches`() = runBlocking {
val tempDir = Files.createTempDirectory("file_edit_replace_fail_multi")
val filePath = tempDir.resolve("test.txt")
Files.writeString(filePath, "fox fox fox")
val tool = FileEditTool(allowedPaths = setOf(tempDir))
val request = createRequest(
mapOf(
"operation" to "replace",
"path" to filePath.toString(),
"target" to "fox",
"replacement" to "dog",
),
)
val result = tool.execute(request)
assertTrue(result is ToolResult.Failure)
val failure = result as ToolResult.Failure
assertTrue(failure.reason.contains("found 3 times"))
}
@Test
fun `execute patch failure`() = runBlocking {
val tempDir = Files.createTempDirectory("file_edit_patch_fail")
val filePath = tempDir.resolve("test.txt")
Files.writeString(filePath, "original")
val tool = FileEditTool(allowedPaths = setOf(tempDir))
val patchContent = """
--- test.txt
+++ test.txt
@@ -1 +1 @@
-nonexistent
+something
""".trimIndent()
val request = createRequest(
mapOf(
"operation" to "patch",
"path" to filePath.toString(),
"patch" to patchContent,
),
)
val result = tool.execute(request)
assertTrue(result is ToolResult.Failure)
val failure = result as ToolResult.Failure
assertTrue(failure.reason.contains("Patch failed"))
}
}
@@ -0,0 +1,112 @@
package com.correx.infrastructure.tools.filesystem
import com.correx.core.events.types.SessionId
import com.correx.core.events.types.StageId
import com.correx.core.events.types.ToolInvocationId
import com.correx.core.tools.contract.ToolResult
import com.correx.core.tools.contract.ValidationResult
import com.correx.core.events.events.ToolRequest
import kotlinx.coroutines.runBlocking
import org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.Assertions.assertFalse
import org.junit.jupiter.api.Assertions.assertTrue
import org.junit.jupiter.api.Test
import java.nio.file.Files
import java.util.UUID
class FileReadToolTest {
private val sessionId = SessionId(UUID.randomUUID().toString())
private val stageId = StageId(UUID.randomUUID().toString())
private val invocationId = ToolInvocationId(UUID.randomUUID().toString())
private fun createRequest(path: String): ToolRequest {
return ToolRequest(
invocationId = invocationId,
sessionId = sessionId,
stageId = stageId,
toolName = "file_read",
parameters = mapOf("path" to path),
)
}
@Test
fun `validateRequest returns Valid for allowed path`() = runBlocking {
val tool = FileReadTool(allowedPaths = emptySet())
val tempFile = Files.createTempFile("test_allowed", ".txt")
val request = createRequest(tempFile.toString())
assertEquals(ValidationResult.Valid, tool.validateRequest(request))
}
@Test
fun `validateRequest returns Invalid for disallowed path`() = runBlocking {
val tempFile = Files.createTempFile("test_disallowed", ".txt")
val tool = FileReadTool(allowedPaths = setOf(tempFile.parent.resolve("subdir")))
val request = createRequest(tempFile.toString())
val result = tool.validateRequest(request)
assertTrue(result is ValidationResult.Invalid)
assertTrue((result as ValidationResult.Invalid).reason.contains("not in the allowed list"))
}
@Test
fun `validateRequest returns Invalid for missing path parameter`() = runBlocking {
val tool = FileReadTool()
val request = ToolRequest(
invocationId = invocationId,
sessionId = sessionId,
stageId = stageId,
toolName = "file_read",
parameters = emptyMap(),
)
val result = tool.validateRequest(request)
assertTrue(result is ValidationResult.Invalid)
assertEquals(
"Missing or invalid 'path' parameter. Expected String.",
(result as ValidationResult.Invalid).reason,
)
}
@Test
fun `execute returns Success for valid file`() = runBlocking {
val content = "Hello, world!"
val tempFile = Files.createTempFile("test_content", ".txt")
Files.writeString(tempFile, content)
val tool = FileReadTool(allowedPaths = emptySet())
val request = createRequest(tempFile.toString())
val result = tool.execute(request)
assertTrue(result is ToolResult.Success)
val success = result as ToolResult.Success
assertEquals(content, success.output)
assertEquals(invocationId, success.invocationId)
}
@Test
fun `execute returns Failure for non-existent file`() = runBlocking {
val tempDir = Files.createTempDirectory("test_dir")
val nonExistentFile = tempDir.resolve("missing.txt")
val tool = FileReadTool(allowedPaths = setOf(nonExistentFile))
val request = createRequest(nonExistentFile.toString())
val result = tool.execute(request)
assertTrue(result is ToolResult.Failure)
val failure = result as ToolResult.Failure
assertTrue(failure.reason.contains("File not found"))
assertFalse(failure.recoverable)
}
@Test
fun `execute returns Failure for disallowed path`() = runBlocking {
val tempFile = Files.createTempFile("test_disallowed_exec", ".txt")
val tool = FileReadTool(allowedPaths = setOf(tempFile.parent.resolve("other")))
val request = createRequest(tempFile.toString())
val result = tool.execute(request)
// It should fail at validation stage
assertTrue(result is ToolResult.Failure)
val failure = result as ToolResult.Failure
assertTrue(failure.reason.contains("is not in the allowed list"))
}
}
@@ -0,0 +1,238 @@
package com.correx.infrastructure.tools.filesystem
import com.correx.core.events.types.SessionId
import com.correx.core.events.types.StageId
import com.correx.core.events.types.ToolInvocationId
import com.correx.core.tools.contract.ToolResult
import com.correx.core.tools.contract.ValidationResult
import com.correx.core.events.events.ToolRequest
import kotlinx.coroutines.runBlocking
import org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.Assertions.assertFalse
import org.junit.jupiter.api.Assertions.assertTrue
import org.junit.jupiter.api.Test
import java.nio.file.Files
import java.nio.file.Path
import java.util.UUID
class FileWriteToolTest {
private val sessionId = SessionId(UUID.randomUUID().toString())
private val stageId = StageId(UUID.randomUUID().toString())
private val invocationId = ToolInvocationId(UUID.randomUUID().toString())
private fun createRequest(
parameters: Map<String, String>,
toolName: String = "file_write",
): ToolRequest {
return ToolRequest(
invocationId = invocationId,
sessionId = sessionId,
stageId = stageId,
toolName = toolName,
parameters = parameters,
)
}
@Test
fun `validateRequest returns Valid for allowed path and valid operation`() = runBlocking {
val tempDir = Files.createTempDirectory("file_write_test")
val tool = FileWriteTool(allowedPaths = setOf(tempDir))
val request = createRequest(
mapOf(
"operation" to "write",
"path" to tempDir.resolve("file.txt").toString(),
"content" to "hello",
),
)
assertEquals(ValidationResult.Valid, tool.validateRequest(request))
}
@Test
fun `validateRequest returns Invalid for disallowed path`() = runBlocking {
val tempDir = Files.createTempDirectory("file_write_test")
val otherDir = Files.createTempDirectory("other_dir")
val tool = FileWriteTool(allowedPaths = setOf(tempDir))
val request = createRequest(
mapOf(
"operation" to "write",
"path" to otherDir.resolve("file.txt").toString(),
"content" to "hello",
),
)
val result = tool.validateRequest(request)
assertTrue(result is ValidationResult.Invalid)
assertTrue((result as ValidationResult.Invalid).reason.contains("is not in the allowed list"))
}
@Test
fun `validateRequest returns Invalid for missing operation`() = runBlocking {
val tempDir = Files.createTempDirectory("file_write_test")
val tool = FileWriteTool(allowedPaths = setOf(tempDir))
val request = createRequest(
mapOf(
"path" to tempDir.resolve("file.txt").toString(),
"content" to "hello",
),
)
val result = tool.validateRequest(request)
assertTrue(result is ValidationResult.Invalid)
assertEquals(
"Missing 'operation' parameter. Expected 'write' or 'delete'.",
(result as ValidationResult.Invalid).reason,
)
}
@Test
fun `validateRequest returns Invalid for missing path`() = runBlocking {
val tempDir = Files.createTempDirectory("file_write_test")
val tool = FileWriteTool(allowedPaths = setOf(tempDir))
val request = createRequest(
mapOf(
"operation" to "write",
"content" to "hello",
),
)
val result = tool.validateRequest(request)
assertTrue(result is ValidationResult.Invalid)
assertEquals("Missing 'path' parameter. Expected String.", (result as ValidationResult.Invalid).reason)
}
@Test
fun `validateRequest returns Invalid for missing content on write`() = runBlocking {
val tempDir = Files.createTempDirectory("file_write_test")
val tool = FileWriteTool(allowedPaths = setOf(tempDir))
val request = createRequest(
mapOf(
"operation" to "write",
"path" to tempDir.resolve("file.txt").toString(),
),
)
val result = tool.validateRequest(request)
assertTrue(result is ValidationResult.Invalid)
assertEquals("Missing 'content' parameter for 'write' operation.", (result as ValidationResult.Invalid).reason)
}
@Test
fun `validateRequest returns Invalid for unknown operation`() = runBlocking {
val tempDir = Files.createTempDirectory("file_write_test")
val tool = FileWriteTool(allowedPaths = setOf(tempDir))
val request = createRequest(
mapOf(
"operation" to "unknown",
"path" to tempDir.resolve("file.txt").toString(),
),
)
val result = tool.validateRequest(request)
assertTrue(result is ValidationResult.Invalid)
assertEquals("Unknown operation: unknown", (result as ValidationResult.Invalid).reason)
}
@Test
fun `execute returns Success for write operation`() = runBlocking {
val tempDir = Files.createTempDirectory("file_write_test")
val tool = FileWriteTool(allowedPaths = setOf(tempDir))
val filePath = tempDir.resolve("test.txt")
val content = "Hello, FileWriteTool!"
val request = createRequest(
mapOf(
"operation" to "write",
"path" to filePath.toString(),
"content" to content,
),
)
val result = tool.execute(request)
assertTrue(result is ToolResult.Success)
val success = result as ToolResult.Success
assertEquals("File written successfully to ${filePath}", success.output)
assertEquals(invocationId, success.invocationId)
assertEquals(content, Files.readString(filePath))
}
@Test
fun `execute returns Success for delete operation`() = runBlocking {
val tempDir = Files.createTempDirectory("file_write_test")
val tool = FileWriteTool(allowedPaths = setOf(tempDir))
val filePath = tempDir.resolve("to_delete.txt")
Files.writeString(filePath, "content to delete")
val request = createRequest(
mapOf(
"operation" to "delete",
"path" to filePath.toString(),
),
)
println("TempDir: $tempDir")
println("FilePath: $filePath")
println("PathString from request: ${request.parameters["path"]}")
val result = tool.execute(request)
println("Result: $result")
println("Exists after: ${Files.exists(filePath)}")
assertTrue(result is ToolResult.Success)
val success = result as ToolResult.Success
assertEquals("File deleted successfully: $filePath", success.output)
assertEquals(invocationId, success.invocationId)
assertFalse(Files.exists(filePath))
}
@Test
fun `execute returns Failure for deleting non-existent file`() = runBlocking {
val tempDir = Files.createTempDirectory("file_write_test")
val tool = FileWriteTool(allowedPaths = setOf(tempDir))
val filePath = tempDir.resolve("non_existent.txt")
val request = createRequest(
mapOf(
"operation" to "delete",
"path" to filePath.toString(),
),
)
val result = tool.execute(request)
assertTrue(result is ToolResult.Failure)
val failure = result as ToolResult.Failure
assertEquals("File not found: $filePath", failure.reason)
assertFalse(failure.recoverable)
}
@Test
fun `execute returns Failure for disallowed path`() = runBlocking {
val tempDir = Files.createTempDirectory("file_write_test")
val otherDir = Files.createTempDirectory("other_dir")
val tool = FileWriteTool(allowedPaths = setOf(tempDir))
val filePath = otherDir.resolve("illegal.txt")
val request = createRequest(
mapOf(
"operation" to "write",
"path" to filePath.toString(),
"content" to "illegal",
),
)
val result = tool.execute(request)
assertTrue(result is ToolResult.Failure)
val failure = result as ToolResult.Failure
assertTrue(failure.reason.contains("is not in the allowed list"))
}
@Test
fun `validateRequest returns Valid for existing file in allowed directory`() = runBlocking {
val tempDir = Files.createTempDirectory("file_write_test_existing")
val filePath = tempDir.resolve("existing.txt")
Files.writeString(filePath, "content")
val tool = FileWriteTool(allowedPaths = setOf(tempDir))
val request = createRequest(
mapOf(
"operation" to "delete",
"path" to filePath.toString(),
),
)
assertEquals(ValidationResult.Valid, tool.validateRequest(request))
}
}
@@ -0,0 +1,19 @@
package com.correx.infrastructure.tools
import com.correx.core.tools.contract.Tool
import com.correx.core.tools.registry.ToolRegistry
class DefaultToolRegistry private constructor(
private val tools: Map<String, Tool>,
) : ToolRegistry {
override fun resolve(name: String): Tool? = tools[name]
override fun all(): List<Tool> = tools.values.toList()
companion object {
fun build(vararg tools: Tool): ToolRegistry =
DefaultToolRegistry(tools.associateBy { it.name })
fun build(tools: List<Tool>): ToolRegistry =
DefaultToolRegistry(tools.associateBy { it.name })
}
}
@@ -0,0 +1,26 @@
package com.correx.infrastructure.tools
import com.correx.core.events.events.ToolRequest
import com.correx.core.tools.contract.ToolExecutor
import com.correx.core.tools.contract.ToolResult
import com.correx.core.tools.registry.ToolRegistry
class DispatchingToolExecutor(
private val registry: ToolRegistry
) : ToolExecutor {
override suspend fun execute(request: ToolRequest): ToolResult {
val tool = registry.resolve(request.toolName)
?: return ToolResult.Failure(
invocationId = request.invocationId,
reason = "tool not found: ${request.toolName}",
recoverable = false
)
return (tool as? ToolExecutor)?.execute(request)
?: ToolResult.Failure(
invocationId = request.invocationId,
reason = "tool ${request.toolName} does not support execution",
recoverable = false
)
}
}
@@ -0,0 +1,267 @@
package com.correx.infrastructure.tools
import com.correx.core.approvals.Tier
import com.correx.core.approvals.domain.ApprovalEngine
import com.correx.core.approvals.model.ApprovalContext
import com.correx.core.approvals.model.ApprovalGrant
import com.correx.core.approvals.model.ApprovalScopeIdentity
import com.correx.core.approvals.model.DomainApprovalRequest
import com.correx.core.events.EventDispatcher
import com.correx.core.events.events.ApprovalGrantCreatedEvent
import com.correx.core.events.events.EventPayload
import com.correx.core.events.events.StoredEvent
import com.correx.core.events.events.ToolExecutionCompletedEvent
import com.correx.core.events.events.ToolExecutionFailedEvent
import com.correx.core.events.events.ToolExecutionRejectedEvent
import com.correx.core.events.events.ToolExecutionStartedEvent
import com.correx.core.events.events.ToolReceipt
import com.correx.core.events.events.ToolRequest
import com.correx.core.events.stores.EventStore
import com.correx.core.events.types.ApprovalRequestId
import com.correx.core.events.types.SessionId
import com.correx.core.events.types.StageId
import com.correx.core.events.types.ToolInvocationId
import com.correx.core.events.types.ValidationReportId
import com.correx.core.sessions.ApprovalMode
import com.correx.core.tools.contract.FileAffectingTool
import com.correx.core.tools.contract.Tool
import com.correx.core.tools.contract.ToolExecutor
import com.correx.core.tools.contract.ToolResult
import com.correx.core.tools.registry.ToolRegistry
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext
import kotlinx.datetime.Clock
import java.io.IOException
import java.nio.file.Files
import java.nio.file.Path
import java.nio.file.StandardCopyOption
import java.util.*
@SuppressWarnings("TooManyFunctions", "LongParameterList")
class SandboxedToolExecutor(
private val delegate: ToolExecutor,
private val registry: ToolRegistry,
private val approvalEngine: ApprovalEngine,
private val eventStore: EventStore,
private val eventDispatcher: EventDispatcher,
private val workDir: Path = Path.of("/tmp/correx-sandbox"),
) : ToolExecutor {
override suspend fun execute(request: ToolRequest): ToolResult = withContext(Dispatchers.IO) {
val sessionId = request.sessionId
val invocationId = request.invocationId
val toolName = request.toolName
val startTime = System.currentTimeMillis()
// 1. resolve tool
val tool = registry.resolve(toolName)
?: return@withContext ToolResult.Failure(
invocationId = invocationId,
reason = "tool not found: $toolName",
recoverable = false,
)
// 2. approval check for T2-T4
val requiresApproval = when (tool.tier) {
Tier.T0, Tier.T1 -> false
Tier.T2, Tier.T3, Tier.T4 -> true
}
if (requiresApproval) {
val sessionEvents = eventStore.read(sessionId)
val decision = evaluateApproval(tool, sessionId, request.stageId, sessionEvents)
if (!decision.isApproved) {
emitRejected(sessionId, invocationId, toolName, tool.tier, decision.reason ?: "approval denied")
return@withContext ToolResult.Failure(
invocationId = invocationId,
reason = decision.reason ?: "approval denied",
recoverable = false,
)
}
}
// 3. emit started
emitStarted(sessionId, invocationId, toolName)
// 4. create working dir
val workingDir = workDir.resolve(sessionId.value).resolve(invocationId.value)
Files.createDirectories(workingDir)
// 5. compute affected paths once (A3: avoid double-invocation divergence)
val affectedPaths: Set<Path> = if (tool is FileAffectingTool) tool.affectedPaths(request) else emptySet()
// 6. backup affected files
val backupMap: Map<Path, Path> = try {
backupAffectedFiles(affectedPaths, workingDir)
} catch (e: IOException) {
return@withContext ToolResult.Failure(
invocationId = invocationId,
reason = "backup failed: ${e.message}",
recoverable = false,
)
}
// 7. delegate
val durationMs = { System.currentTimeMillis() - startTime }
when (val result = delegate.execute(request)) {
is ToolResult.Success -> {
restoreOrClean(backupMap, success = true)
cleanWorkingDir(workingDir)
emitCompleted(sessionId, invocationId, toolName, result, tool, affectedPaths, durationMs())
result
}
is ToolResult.Failure -> {
restoreOrClean(backupMap, success = false)
cleanWorkingDir(workingDir)
emitFailed(sessionId, invocationId, toolName, result.reason)
result
}
}
}
// --- approval ---
private fun evaluateApproval(
tool: Tool,
sessionId: SessionId,
stageId: StageId,
sessionEvents: List<StoredEvent>,
) = approvalEngine.evaluate(
request = DomainApprovalRequest(
id = ApprovalRequestId(UUID.randomUUID().toString()),
tier = tool.tier,
validationReportId = ValidationReportId(UUID.randomUUID().toString()),
riskSummaryId = null,
timestamp = Clock.System.now(),
causationId = null,
correlationId = null,
),
context = ApprovalContext(
identity = ApprovalScopeIdentity(sessionId, stageId, null),
mode = ApprovalMode.PROMPT,
),
grants = extractGrants(sessionEvents),
now = Clock.System.now(),
)
private fun extractGrants(events: List<StoredEvent>): List<ApprovalGrant> =
events
.mapNotNull { it.payload as? ApprovalGrantCreatedEvent }
.map { e ->
ApprovalGrant(
id = e.grantId,
scope = e.scope,
permittedTiers = e.permittedTiers,
reason = e.reason,
timestamp = Clock.System.now(),
expiresAt = e.expiresAt,
)
}
// --- backup/restore ---
/**
* Returns a map of originalPath -> backupPath for all affected files.
* Throws IOException if any backup fails — caller handles this.
*/
private fun backupAffectedFiles(
affectedPaths: Collection<Path>,
workingDir: Path,
): Map<Path, Path> = buildMap {
for (original in affectedPaths) {
// A1: skip files that don't yet exist (new-file tools have nothing to back up)
if (!Files.exists(original)) continue
val backupPath = workingDir.resolve("${original.fileName}.bak")
Files.copy(original, backupPath, StandardCopyOption.REPLACE_EXISTING)
put(original, backupPath)
}
}
/**
* On success: delete backups.
* On failure: restore originals from backups.
*/
private fun restoreOrClean(backupMap: Map<Path, Path>, success: Boolean) {
for ((original, backup) in backupMap) {
try {
if (success) {
Files.deleteIfExists(backup)
} else {
Files.move(backup, original, StandardCopyOption.REPLACE_EXISTING)
}
} catch (_: IOException) {
// best effort
}
}
}
private fun cleanWorkingDir(workingDir: Path) {
try {
Files.walk(workingDir)
.sorted(Comparator.reverseOrder())
.forEach { Files.deleteIfExists(it) }
} catch (_: IOException) {
// best effort
}
}
// --- event emission ---
private suspend fun emitRejected(
sessionId: SessionId,
invocationId: ToolInvocationId,
toolName: String,
tier: Tier,
reason: String,
) = emit(sessionId, ToolExecutionRejectedEvent(invocationId, sessionId, toolName, tier, reason))
private suspend fun emitStarted(
sessionId: SessionId,
invocationId: ToolInvocationId,
toolName: String,
) = emit(sessionId, ToolExecutionStartedEvent(invocationId, sessionId, toolName))
private suspend fun emitCompleted(
sessionId: SessionId,
invocationId: ToolInvocationId,
toolName: String,
result: ToolResult.Success,
tool: Tool,
affectedPaths: Set<Path>,
durationMs: Long,
) {
val affectedEntities = affectedPaths.map { it.toString() }
emit(
sessionId,
ToolExecutionCompletedEvent(
invocationId = invocationId,
sessionId = sessionId,
toolName = toolName,
receipt = ToolReceipt(
invocationId = invocationId,
toolName = toolName,
exitCode = result.exitCode,
outputSummary = result.output,
structuredOutput = result.metadata,
affectedEntities = affectedEntities,
durationMs = durationMs,
tier = tool.tier,
timestamp = Clock.System.now(),
),
),
)
}
private suspend fun emitFailed(
sessionId: SessionId,
invocationId: ToolInvocationId,
toolName: String,
reason: String,
) = emit(sessionId, ToolExecutionFailedEvent(invocationId, sessionId, toolName, reason))
private suspend fun emit(sessionId: SessionId, payload: EventPayload) {
eventDispatcher.emit(payload, sessionId)
}
}
@@ -0,0 +1,46 @@
package com.correx.infrastructure.tools
import com.correx.core.tools.contract.Tool
import com.correx.infrastructure.tools.filesystem.FileEditTool
import com.correx.infrastructure.tools.filesystem.FileReadTool
import com.correx.infrastructure.tools.filesystem.FileWriteTool
import com.correx.infrastructure.tools.shell.ShellTool
import java.nio.file.Path
data class ToolConfig(
val shell: ShellConfig = ShellConfig(),
val fileRead: FileReadConfig = FileReadConfig(),
val fileWrite: FileWriteConfig = FileWriteConfig(),
val fileEdit: FileEditConfig = FileEditConfig(),
)
data class FileReadConfig(val enabled: Boolean = false, val allowedPaths: Set<Path> = emptySet())
data class FileWriteConfig(val enabled: Boolean = false, val allowedPaths: Set<Path> = emptySet())
data class FileEditConfig(val enabled: Boolean = false, val allowedPaths: Set<Path> = emptySet())
data class ShellConfig(val enabled: Boolean = false, val allowedExecutables: Set<String> = emptySet())
/**
* Extension function to convert [ToolConfig] into a list of [Tool] implementations.
*/
fun ToolConfig.buildTools(): List<Tool> = buildList {
if (shell.enabled) {
add(ShellTool(
allowedExecutables = shell.allowedExecutables
))
}
if (fileRead.enabled) {
add(FileReadTool(
allowedPaths = fileRead.allowedPaths
))
}
if (fileWrite.enabled) {
add(FileWriteTool(
allowedPaths = fileWrite.allowedPaths
))
}
if (fileEdit.enabled) {
add(FileEditTool(
allowedPaths = fileEdit.allowedPaths
))
}
}
@@ -0,0 +1,93 @@
package com.correx.infrastructure.tools.shell
import com.correx.core.approvals.Tier
import com.correx.core.events.events.ToolRequest
import com.correx.core.tools.contract.Tool
import com.correx.core.tools.contract.ToolCapability
import com.correx.core.tools.contract.ToolExecutor
import com.correx.core.tools.contract.ToolResult
import com.correx.core.tools.contract.ValidationResult
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.TimeoutCancellationException
import kotlinx.coroutines.async
import kotlinx.coroutines.withContext
import kotlinx.coroutines.withTimeout
class ShellTool(
private val allowedExecutables: Set<String> = emptySet(),
private val timeoutMs: Long = 30_000L,
) : Tool, ToolExecutor {
override val name: String = "shell"
override val tier: Tier = Tier.T2
override val requiredCapabilities: Set<ToolCapability> =
setOf(ToolCapability.SHELL_EXEC, ToolCapability.PROCESS_SPAWN)
override fun validateRequest(request: ToolRequest): ValidationResult {
val argv = (request.parameters["argv"] as? List<*>)?.filterIsInstance<String>()
return when {
argv.isNullOrEmpty() ->
ValidationResult.Invalid("Missing or empty 'argv' parameter. Expected List<String>.")
argv[0] !in allowedExecutables ->
ValidationResult.Invalid("Executable '${argv[0]}' is not in the allowed list.")
else -> ValidationResult.Valid
}
}
override suspend fun execute(request: ToolRequest): ToolResult = withContext(Dispatchers.IO) {
validateRequest(request).run {
takeIf { this is ValidationResult.Valid }?.let {
val argv = (request.parameters["argv"] as? List<*>)?.filterIsInstance<String>()
val process = ProcessBuilder(argv).apply {
environment().clear()
}.start()
runCatching {
runCmd(request, process)
}.getOrElse {
process.destroyForcibly()
if (it is TimeoutCancellationException) {
ToolResult.Failure(
invocationId = request.invocationId,
reason = "Process timed out after ${timeoutMs}ms",
recoverable = false,
)
} else {
ToolResult.Failure(
invocationId = request.invocationId,
reason = it.message ?: "Unknown error occurred during execution",
recoverable = false,
)
}
}
} ?: ToolResult.Failure(
invocationId = request.invocationId,
reason = (this as ValidationResult.Invalid).reason,
recoverable = false,
)
}
}
private suspend fun runCmd(request: ToolRequest, process: Process): ToolResult = withTimeout(timeoutMs) {
val stdoutDeferred = async { process.inputStream.bufferedReader().use { it.readText() } }
val stderrDeferred = async { process.errorStream.bufferedReader().use { it.readText() } }
val exitCode = process.waitFor()
val stdout = stdoutDeferred.await()
val stderr = stderrDeferred.await()
val metadata = if (stderr.isNotBlank()) mapOf("stderr" to stderr) else emptyMap()
if (exitCode != 0) {
ToolResult.Failure(
invocationId = request.invocationId,
reason = "Process exited with code $exitCode${if (stderr.isNotBlank()) ": $stderr" else ""}",
recoverable = false,
)
} else {
ToolResult.Success(
invocationId = request.invocationId,
output = stdout,
exitCode = exitCode,
metadata = metadata,
)
}
}
}
@@ -0,0 +1,115 @@
package com.correx.infrastructure.tools.shell
import com.correx.core.events.events.ToolRequest
import com.correx.core.events.types.SessionId
import com.correx.core.events.types.StageId
import com.correx.core.events.types.ToolInvocationId
import com.correx.core.tools.contract.ToolResult
import com.correx.core.tools.contract.ValidationResult
import kotlinx.coroutines.runBlocking
import org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.Assertions.assertFalse
import org.junit.jupiter.api.Assertions.assertTrue
import org.junit.jupiter.api.Test
import java.util.UUID
class ShellToolTest {
private val sessionId = SessionId(UUID.randomUUID().toString())
private val stageId = StageId(UUID.randomUUID().toString())
private val invocationId = ToolInvocationId(UUID.randomUUID().toString())
private fun createRequest(argv: List<String>): ToolRequest {
return ToolRequest(
invocationId = invocationId,
sessionId = sessionId,
stageId = stageId,
toolName = "shell",
parameters = mapOf("argv" to argv),
)
}
@Test
fun `validateRequest returns Valid for allowed executable`() = runBlocking {
val tool = ShellTool(allowedExecutables = setOf("echo"))
val request = createRequest(listOf("echo", "hello"))
assertEquals(ValidationResult.Valid, tool.validateRequest(request))
}
@Test
fun `validateRequest returns Invalid for disallowed executable`() = runBlocking {
val tool = ShellTool(allowedExecutables = setOf("echo"))
val request = createRequest(listOf("ls"))
val result = tool.validateRequest(request)
assertTrue(result is ValidationResult.Invalid)
assertEquals("Executable 'ls' is not in the allowed list.", (result as ValidationResult.Invalid).reason)
}
@Test
fun `validateRequest returns Invalid for empty argv`() = runBlocking {
val tool = ShellTool(allowedExecutables = setOf("echo"))
val request = createRequest(emptyList())
val result = tool.validateRequest(request)
assertTrue(result is ValidationResult.Invalid)
assertEquals(
"Missing or empty 'argv' parameter. Expected List<String>.",
(result as ValidationResult.Invalid).reason,
)
}
@Test
fun `validateRequest returns Invalid for null argv`() = runBlocking {
val tool = ShellTool(allowedExecutables = setOf("echo"))
val request = ToolRequest(
invocationId = invocationId,
sessionId = sessionId,
stageId = stageId,
toolName = "shell",
parameters = emptyMap(),
)
val result = tool.validateRequest(request)
assertTrue(result is ValidationResult.Invalid)
assertEquals(
"Missing or empty 'argv' parameter. Expected List<String>.",
(result as ValidationResult.Invalid).reason,
)
}
@Test
fun `execute returns Success for exit code 0`() = runBlocking {
val tool = ShellTool(allowedExecutables = setOf("echo"))
val request = createRequest(listOf("echo", "hello"))
val result = tool.execute(request)
assertTrue(result is ToolResult.Success)
val success = result as ToolResult.Success
assertTrue(success.output.startsWith("hello"))
assertEquals(0, success.exitCode)
assertEquals(invocationId, success.invocationId)
}
@Test
fun `execute returns Failure with non-zero exit code for failed process`() = runBlocking {
val tool = ShellTool(allowedExecutables = setOf("false"))
val request = createRequest(listOf("false"))
val result = tool.execute(request)
assertTrue(result is ToolResult.Failure)
val failure = result as ToolResult.Failure
assertTrue(failure.reason.contains("exited with code 1"))
assertEquals(invocationId, failure.invocationId)
}
@Test
fun `execute returns Failure on timeout`() = runBlocking {
// Using sleep command to simulate timeout
val tool = ShellTool(allowedExecutables = setOf("sleep"), timeoutMs = 100L)
val request = createRequest(listOf("sleep", "1"))
val result = tool.execute(request)
assertTrue(result is ToolResult.Failure)
val failure = result as ToolResult.Failure
assertEquals("Process timed out after 100ms", failure.reason)
assertFalse(failure.recoverable)
}
}