epic-12: after epic audit and init commit
This commit is contained in:
@@ -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"
|
||||
}
|
||||
@@ -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"
|
||||
}
|
||||
+26
@@ -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
|
||||
}
|
||||
+15
@@ -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(),
|
||||
)
|
||||
+11
@@ -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)
|
||||
+36
@@ -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
|
||||
}
|
||||
+10
@@ -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"
|
||||
}
|
||||
+208
@@ -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
|
||||
}
|
||||
}
|
||||
+45
@@ -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>)
|
||||
+120
@@ -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", "")) }
|
||||
}
|
||||
}
|
||||
+26
@@ -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
|
||||
}
|
||||
+39
@@ -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
|
||||
}
|
||||
}
|
||||
+38
@@ -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() }
|
||||
}
|
||||
+32
@@ -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
|
||||
+80
@@ -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 })
|
||||
}
|
||||
}
|
||||
+77
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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"
|
||||
}
|
||||
+68
@@ -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
|
||||
}
|
||||
}
|
||||
+215
@@ -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"))
|
||||
)
|
||||
}
|
||||
+23
@@ -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
|
||||
}
|
||||
}
|
||||
}
|
||||
+8
@@ -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()
|
||||
}
|
||||
+23
@@ -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")
|
||||
)
|
||||
}
|
||||
}
|
||||
+12
@@ -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)
|
||||
}
|
||||
}
|
||||
+28
@@ -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")
|
||||
)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
@@ -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,
|
||||
)
|
||||
}
|
||||
@@ -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
|
||||
@@ -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')
|
||||
}
|
||||
+231
@@ -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,
|
||||
)
|
||||
}
|
||||
}
|
||||
+97
@@ -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,
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
+172
@@ -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,
|
||||
)
|
||||
}
|
||||
}
|
||||
+234
@@ -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"))
|
||||
}
|
||||
}
|
||||
+112
@@ -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"))
|
||||
}
|
||||
}
|
||||
+238
@@ -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))
|
||||
}
|
||||
}
|
||||
+19
@@ -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 })
|
||||
}
|
||||
}
|
||||
+26
@@ -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
|
||||
)
|
||||
}
|
||||
}
|
||||
+267
@@ -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
|
||||
))
|
||||
}
|
||||
}
|
||||
+93
@@ -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,
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
+115
@@ -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)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user