feat: capture file pre/post-image to CAS and emit FileWrittenEvent

This commit is contained in:
2026-05-31 10:34:30 +04:00
parent 0e688f3daf
commit b778c06f6c
4 changed files with 189 additions and 0 deletions
@@ -105,6 +105,7 @@ fun main() {
registry = toolRegistry,
eventDispatcher = EventDispatcher(eventStore),
workDir = sandboxRoot,
artifactStore = artifactStore,
)
logToolInfo(sandboxRoot, workingDir, shellAllowedExecutables, toolRegistry)
@@ -135,11 +135,13 @@ object InfrastructureModule {
registry: ToolRegistry,
eventDispatcher: EventDispatcher,
workDir: Path,
artifactStore: ArtifactStore? = null,
): ToolExecutor = SandboxedToolExecutor(
delegate = DispatchingToolExecutor(registry),
registry = registry,
eventDispatcher = eventDispatcher,
workDir = workDir,
artifactStore = artifactStore,
)
fun createNoopEmbedder(dimension: Int = 1536): Embedder = NoopEmbedder(dimension)
@@ -1,7 +1,9 @@
package com.correx.infrastructure.tools
import com.correx.core.artifactstore.ArtifactStore
import com.correx.core.events.EventDispatcher
import com.correx.core.events.events.EventPayload
import com.correx.core.events.events.FileWrittenEvent
import com.correx.core.events.events.ToolExecutionCompletedEvent
import com.correx.core.events.events.ToolExecutionFailedEvent
import com.correx.core.events.events.ToolExecutionStartedEvent
@@ -28,8 +30,11 @@ class SandboxedToolExecutor(
private val registry: ToolRegistry,
private val eventDispatcher: EventDispatcher,
private val workDir: Path = Path.of("/tmp/correx-sandbox"),
private val artifactStore: ArtifactStore? = null,
) : ToolExecutor {
private data class PreImage(val existed: Boolean, val hash: String?)
override suspend fun execute(request: ToolRequest): ToolResult = withContext(Dispatchers.IO) {
val sessionId = request.sessionId
val invocationId = request.invocationId
@@ -54,6 +59,9 @@ class SandboxedToolExecutor(
// 5. compute affected paths once (A3: avoid double-invocation divergence)
val affectedPaths: Set<Path> = if (tool is FileAffectingTool) tool.affectedPaths(request) else emptySet()
// 5b. capture pre-image bytes into CAS for event-sourced reversibility (invariants #1/#5)
val preImages: Map<Path, PreImage> = capturePreImages(affectedPaths)
// 6. backup affected files
val backupMap: Map<Path, Path> = try {
backupAffectedFiles(affectedPaths, workingDir)
@@ -74,6 +82,7 @@ class SandboxedToolExecutor(
restoreOrClean(backupMap, success = true)
cleanWorkingDir(workingDir)
emitCompleted(sessionId, invocationId, toolName, result, tool, affectedPaths, durationMs(), diff)
emitFileMutations(sessionId, invocationId, affectedPaths, preImages)
result
}
@@ -215,4 +224,46 @@ class SandboxedToolExecutor(
private suspend fun emit(sessionId: SessionId, payload: EventPayload) {
eventDispatcher.emit(payload, sessionId)
}
// --- event-sourced reversibility (CAS pre/post images) ---
private suspend fun capturePreImages(paths: Set<Path>): Map<Path, PreImage> {
if (artifactStore == null) return emptyMap()
return paths.associateWith { path ->
val existed = Files.exists(path)
PreImage(existed = existed, hash = if (existed) storeBytes(path) else null)
}
}
private suspend fun storeBytes(path: Path): String? {
val store = artifactStore ?: return null
val bytes = runCatching { Files.readAllBytes(path) }.getOrNull() ?: return null
return store.put(bytes).value
}
/** Emits one [FileWrittenEvent] per affected path, carrying the CAS hashes that make the
* mutation reversible. Post-image is read after the swap; null ⇒ the path was deleted. */
private suspend fun emitFileMutations(
sessionId: SessionId,
invocationId: ToolInvocationId,
affectedPaths: Set<Path>,
preImages: Map<Path, PreImage>,
) {
if (artifactStore == null) return
for (path in affectedPaths) {
val pre = preImages[path]
emit(
sessionId,
FileWrittenEvent(
invocationId = invocationId,
sessionId = sessionId,
path = path.toString(),
preImageHash = pre?.hash,
postImageHash = storeBytes(path),
preExisted = pre?.existed ?: false,
timestampMs = Clock.System.now().toEpochMilliseconds(),
),
)
}
}
}
@@ -0,0 +1,135 @@
package com.correx.infrastructure.tools
import com.correx.core.artifactstore.ArtifactStore
import com.correx.core.events.EventDispatcher
import com.correx.core.events.events.EventPayload
import com.correx.core.events.events.FileWrittenEvent
import com.correx.core.events.events.NewEvent
import com.correx.core.events.events.StoredEvent
import com.correx.core.events.events.ToolRequest
import com.correx.core.events.stores.EventStore
import com.correx.core.events.types.ArtifactId
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.Tool
import com.correx.core.tools.contract.ToolResult
import com.correx.core.tools.registry.ToolRegistry
import com.correx.infrastructure.tools.filesystem.FileWriteTool
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.emptyFlow
import kotlinx.coroutines.runBlocking
import org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.Assertions.assertNotNull
import org.junit.jupiter.api.Assertions.assertNull
import org.junit.jupiter.api.Assertions.assertTrue
import org.junit.jupiter.api.Test
import java.nio.file.Files
class SandboxedToolExecutorFileMutationTest {
/** Content-addressed fake: identical bytes → identical id (mirrors CAS dedup). */
private class FakeArtifactStore : ArtifactStore {
val blobs = linkedMapOf<String, ByteArray>()
override suspend fun put(bytes: ByteArray): ArtifactId {
val id = ArtifactId(bytes.toString(Charsets.UTF_8).hashCode().toString())
blobs[id.value] = bytes
return id
}
override suspend fun get(id: ArtifactId): ByteArray? = blobs[id.value]
override suspend fun flushBefore(commit: suspend () -> Unit) = commit()
}
private class CapturingEventStore : EventStore {
val payloads = mutableListOf<EventPayload>()
override suspend fun append(event: NewEvent): StoredEvent {
payloads += event.payload
val seq = payloads.size.toLong()
return StoredEvent(event.metadata, seq, seq, event.payload)
}
override suspend fun appendAll(events: List<NewEvent>): List<StoredEvent> = events.map { append(it) }
override fun read(sessionId: SessionId): List<StoredEvent> = emptyList()
override fun readFrom(sessionId: SessionId, fromSequence: Long): List<StoredEvent> = emptyList()
override fun lastSequence(sessionId: SessionId): Long? = null
override fun subscribe(sessionId: SessionId): Flow<StoredEvent> = emptyFlow()
override fun subscribeAll(): Flow<StoredEvent> = emptyFlow()
override suspend fun lastGlobalSequence(): Long = payloads.size.toLong()
override fun allEvents(): Sequence<StoredEvent> = emptySequence()
override fun allSessionIds(): Set<SessionId> = emptySet()
}
private class SingleToolRegistry(private val tool: Tool) : ToolRegistry {
override fun resolve(name: String): Tool? = if (name == tool.name) tool else null
override fun all(): List<Tool> = listOf(tool)
}
private fun writeRequest(path: String, content: String) = ToolRequest(
ToolInvocationId("inv"), SessionId("s"), StageId("st"), "file_write",
mapOf("operation" to "write", "path" to path, "content" to content),
)
private fun executor(tool: FileWriteTool, store: FakeArtifactStore, events: CapturingEventStore) =
SandboxedToolExecutor(
delegate = tool,
registry = SingleToolRegistry(tool),
eventDispatcher = EventDispatcher(events),
workDir = Files.createTempDirectory("sbx-work"),
artifactStore = store,
)
@Test
fun `overwrite emits FileWrittenEvent with distinct pre and post images in CAS`(): Unit = runBlocking {
val dir = Files.createTempDirectory("sbx").toRealPath()
val target = dir.resolve("f.txt")
Files.writeString(target, "OLD")
val tool = FileWriteTool(allowedPaths = setOf(dir))
val store = FakeArtifactStore()
val events = CapturingEventStore()
val result = executor(tool, store, events).execute(writeRequest(target.toString(), "NEW"))
assertTrue(result is ToolResult.Success)
val fw = events.payloads.filterIsInstance<FileWrittenEvent>().single()
assertEquals(target.toString(), fw.path)
assertTrue(fw.preExisted)
assertNotNull(fw.preImageHash)
assertNotNull(fw.postImageHash)
assertTrue(fw.preImageHash != fw.postImageHash)
assertEquals("OLD", store.get(ArtifactId(fw.preImageHash!!))!!.toString(Charsets.UTF_8))
assertEquals("NEW", store.get(ArtifactId(fw.postImageHash!!))!!.toString(Charsets.UTF_8))
}
@Test
fun `new file write records preExisted false and null pre-image`(): Unit = runBlocking {
val dir = Files.createTempDirectory("sbx-new").toRealPath()
val target = dir.resolve("new.txt")
val tool = FileWriteTool(allowedPaths = setOf(dir))
val store = FakeArtifactStore()
val events = CapturingEventStore()
executor(tool, store, events).execute(writeRequest(target.toString(), "HELLO"))
val fw = events.payloads.filterIsInstance<FileWrittenEvent>().single()
assertEquals(false, fw.preExisted)
assertNull(fw.preImageHash)
assertEquals("HELLO", store.get(ArtifactId(fw.postImageHash!!))!!.toString(Charsets.UTF_8))
}
@Test
fun `no artifact store means no FileWrittenEvent (backward compatible)`(): Unit = runBlocking {
val dir = Files.createTempDirectory("sbx-nostore").toRealPath()
val target = dir.resolve("f.txt")
val tool = FileWriteTool(allowedPaths = setOf(dir))
val events = CapturingEventStore()
val exec = SandboxedToolExecutor(
delegate = tool,
registry = SingleToolRegistry(tool),
eventDispatcher = EventDispatcher(events),
workDir = Files.createTempDirectory("sbx-nostore-work"),
)
exec.execute(writeRequest(target.toString(), "X"))
assertTrue(events.payloads.filterIsInstance<FileWrittenEvent>().isEmpty())
}
}