feat(workspace): per-session workspace-scoped tools, policy, and undo

Compute the effective tool registry/executor and plane-2 WorkspacePolicy per run
from config.workspace (effectivesFor), threaded through the orchestrators; a null
workspace falls back to the boot instances byte-for-byte. Wire the concrete
WorkspaceToolRegistryProvider in Main (buildToolConfigForWorkspace). Session undo
unions the session's recorded workspace into its jail roots. (Axis 2 Phase A, tasks 3 + 7.)
This commit is contained in:
2026-06-02 19:57:51 +04:00
parent 7d7e524756
commit b2f60d09bb
11 changed files with 720 additions and 36 deletions
@@ -28,6 +28,9 @@ import com.correx.core.kernel.orchestration.OrchestrationProjector
import com.correx.core.kernel.orchestration.OrchestrationRepository
import com.correx.core.kernel.orchestration.OrchestratorEngines
import com.correx.core.kernel.orchestration.OrchestratorRepositories
import com.correx.core.kernel.orchestration.WorkspaceContext
import com.correx.core.kernel.orchestration.WorkspaceTools
import com.correx.core.kernel.orchestration.WorkspaceToolRegistryProvider
import com.correx.core.kernel.retry.DefaultRetryCoordinator
import com.correx.core.risk.DefaultRiskAssessor
import com.correx.core.router.l3.L3MetadataRehydrator
@@ -57,6 +60,7 @@ import com.correx.infrastructure.inference.commons.ResourceProbe
import com.correx.infrastructure.inference.commons.UnavailableProbe
import com.correx.core.inference.InferenceProvider
import com.correx.infrastructure.inference.llama.cpp.LlamaCppInferenceProvider
import com.correx.infrastructure.tools.DispatchingToolExecutor
import com.correx.infrastructure.tools.FileEditConfig
import com.correx.infrastructure.tools.FileReadConfig
import com.correx.infrastructure.tools.FileWriteConfig
@@ -188,6 +192,14 @@ fun main() {
),
)
val wsToolRegistryProvider = WorkspaceToolRegistryProvider { workspace ->
val wsRegistry = InfrastructureModule.createToolRegistry(
buildToolConfigForWorkspace(workspace, shellAllowedExecutables, toolsConfig),
)
val wsExecutor = DispatchingToolExecutor(wsRegistry)
WorkspaceTools(registry = wsRegistry, executor = wsExecutor)
}
val engines = OrchestratorEngines(
transitionResolver = DefaultTransitionResolver { condition, ctx -> condition.evaluate(ctx) },
contextPackBuilder = DefaultContextPackBuilder(DefaultContextCompressor()),
@@ -210,6 +222,7 @@ fun main() {
toolExecutor = toolExecutor,
toolCallAssessor = toolCallAssessor,
workspacePolicy = workspacePolicy,
workspaceToolRegistryProvider = wsToolRegistryProvider,
)
val orchestrator = DefaultSessionOrchestrator(
repositories = repositories,
@@ -239,10 +252,8 @@ fun main() {
)
val sessionUndoService = com.correx.apps.server.undo.SessionUndoService(
eventStore = eventStore,
reverser = com.correx.infrastructure.tools.filesystem.FileMutationReverser(
artifactStore = artifactStore,
allowedRoots = setOf(workspaceRoot, workingDir),
),
artifactStore = artifactStore,
bootRoots = setOf(workspaceRoot, workingDir),
)
val module = ServerModule(
orchestrator = orchestrator,
@@ -436,6 +447,32 @@ private fun buildToolConfig(
)
}
private fun buildToolConfigForWorkspace(
workspace: WorkspaceContext,
shellAllowedExecutables: Set<String>,
toolsConfig: com.correx.core.config.ToolsConfig,
): ToolConfig = ToolConfig(
shell = ShellConfig(
enabled = toolsConfig.shellEnabled,
allowedExecutables = shellAllowedExecutables,
workingDir = workspace.workingDir,
),
fileRead = FileReadConfig(
enabled = toolsConfig.fileReadEnabled,
allowedPaths = workspace.allowedPaths,
),
fileWrite = FileWriteConfig(
enabled = toolsConfig.fileWriteEnabled,
allowedPaths = workspace.allowedPaths,
workingDir = workspace.workingDir,
),
fileEdit = FileEditConfig(
enabled = toolsConfig.fileEditEnabled,
allowedPaths = workspace.allowedPaths,
workingDir = workspace.workingDir,
),
)
private fun DefaultProviderRegistry.asServerRegistry(): ProviderRegistry = object : ProviderRegistry {
override fun listAll() = this@asServerRegistry.listAll()
override suspend fun healthCheckAll() = this@asServerRegistry.healthCheckAll()
@@ -1,11 +1,14 @@
package com.correx.apps.server.undo
import com.correx.core.artifactstore.ArtifactStore
import com.correx.core.events.events.FileWrittenEvent
import com.correx.core.events.events.SessionWorkspaceBoundEvent
import com.correx.core.events.stores.EventStore
import com.correx.core.events.types.SessionId
import com.correx.infrastructure.tools.filesystem.FileMutationReverser
import com.correx.infrastructure.tools.filesystem.RevertResult
import kotlinx.serialization.Serializable
import java.nio.file.Path
@Serializable
data class UndoSummary(
@@ -21,13 +24,34 @@ data class UndoSummary(
* Reverses every file mutation a session performed, newest-first, using only the
* event log + CAS (the reverser). Newest-first so a path written multiple times in
* the session ends at its pre-session content.
*
* Jail roots for each undo = [bootRoots] the workspace recorded in the session's
* [SessionWorkspaceBoundEvent] (if present). When no workspace event exists the
* effective roots are identical to [bootRoots], preserving pre-existing behaviour.
*/
class SessionUndoService(
private val eventStore: EventStore,
private val reverser: FileMutationReverser,
private val artifactStore: ArtifactStore,
private val bootRoots: Set<Path>,
) {
suspend fun undo(sessionId: SessionId): UndoSummary {
val mutations = eventStore.read(sessionId)
val events = eventStore.read(sessionId)
val sessionRoots: Set<Path> = events
.mapNotNull { it.payload as? SessionWorkspaceBoundEvent }
.lastOrNull()
?.let { bound ->
buildSet {
add(Path.of(bound.workspaceRoot))
bound.allowedPaths.forEach { add(Path.of(it)) }
}
}
?: emptySet()
val effectiveRoots = bootRoots + sessionRoots
val reverser = FileMutationReverser(artifactStore, effectiveRoots)
val mutations = events
.mapNotNull { it.payload as? FileWrittenEvent }
.asReversed()
@@ -44,7 +44,6 @@ import com.correx.infrastructure.tools.FileReadConfig
import com.correx.infrastructure.tools.FileWriteConfig
import com.correx.infrastructure.tools.ShellConfig
import com.correx.infrastructure.tools.ToolConfig
import com.correx.infrastructure.tools.filesystem.FileMutationReverser
import io.ktor.client.HttpClient
import io.ktor.client.engine.mock.MockEngine
import io.ktor.client.engine.mock.respond
@@ -124,10 +123,8 @@ fun buildTestServerModule(
val sessionUndoService = SessionUndoService(
eventStore = eventStore,
reverser = FileMutationReverser(
artifactStore = artifactStore,
allowedRoots = setOf(tempDir),
),
artifactStore = artifactStore,
bootRoots = setOf(tempDir),
)
val providerRegistry: ProviderRegistry = object : ProviderRegistry {
@@ -5,13 +5,13 @@ import com.correx.core.events.events.EventMetadata
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.SessionWorkspaceBoundEvent
import com.correx.core.events.events.StoredEvent
import com.correx.core.events.stores.EventStore
import com.correx.core.events.types.ArtifactId
import com.correx.core.events.types.EventId
import com.correx.core.events.types.SessionId
import com.correx.core.events.types.ToolInvocationId
import com.correx.infrastructure.tools.filesystem.FileMutationReverser
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.emptyFlow
import kotlinx.coroutines.runBlocking
@@ -20,6 +20,7 @@ import org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.Assertions.assertFalse
import org.junit.jupiter.api.Test
import java.nio.file.Files
import java.nio.file.Path
class SessionUndoServiceTest {
@@ -83,8 +84,7 @@ class SessionUndoServiceTest {
timestampMs = 0L,
)
val fakeEventStore = FakeEventStore(listOf(storedEvent(event, sessionId, 1L)))
val reverser = FileMutationReverser(fakeStore, setOf(dir))
val service = SessionUndoService(fakeEventStore, reverser)
val service = SessionUndoService(fakeEventStore, fakeStore, setOf(dir))
val summary = service.undo(sessionId)
@@ -113,8 +113,7 @@ class SessionUndoServiceTest {
timestampMs = 0L,
)
val fakeEventStore = FakeEventStore(listOf(storedEvent(event, sessionId, 1L)))
val reverser = FileMutationReverser(fakeStore, setOf(dir))
val service = SessionUndoService(fakeEventStore, reverser)
val service = SessionUndoService(fakeEventStore, fakeStore, setOf(dir))
val summary = service.undo(sessionId)
@@ -124,4 +123,212 @@ class SessionUndoServiceTest {
assertEquals(0, summary.failed)
assertFalse(Files.exists(target))
}
/**
* When the session log contains a [SessionWorkspaceBoundEvent] the workspace root
* is unioned into the jail roots, so a file written under that workspace (which is
* NOT in the boot roots) is successfully reverted rather than rejected.
*
* The sibling assertion that proves the workspace exclusion: a service built with
* boot roots that do NOT include [workspaceDir] would return rejected=1 instead of
* reverted=1. We verify this via a separate call using an unrelated boot root.
*/
@Test
fun `undo uses session workspace root from SessionWorkspaceBoundEvent`(): Unit = runBlocking {
val bootDir = Files.createTempDirectory("undo-boot").toRealPath()
val workspaceDir = Files.createTempDirectory("undo-workspace").toRealPath()
val target = workspaceDir.resolve("ws-file.txt")
Files.writeString(target, "CURRENT")
val fakeStore = FakeArtifactStore()
val preImageHash = fakeStore.put("ORIGINAL".toByteArray()).value
val sessionId = SessionId("session-workspace")
val workspaceBound = SessionWorkspaceBoundEvent(
sessionId = sessionId,
workspaceRoot = workspaceDir.toString(),
allowedPaths = emptyList(),
)
val fileWritten = FileWrittenEvent(
invocationId = ToolInvocationId("inv-ws"),
sessionId = sessionId,
path = target.toString(),
preImageHash = preImageHash,
postImageHash = null,
preExisted = true,
timestampMs = 0L,
)
val eventList = listOf(
storedEvent(workspaceBound, sessionId, 1L),
storedEvent(fileWritten, sessionId, 2L),
)
val fakeEventStore = FakeEventStore(eventList)
// Boot roots do NOT include workspaceDir — the workspace event unlocks it.
val service = SessionUndoService(fakeEventStore, fakeStore, setOf(bootDir))
val summary = service.undo(sessionId)
assertEquals(1, summary.reverted)
assertEquals(0, summary.rejected, "workspace file must not be jailed out when bound event is present")
assertEquals("ORIGINAL", Files.readString(target))
// Prove that without the workspace event the revert would have been jailed:
// a service with boot-only roots (excluding workspaceDir) and no bound event rejects the file.
val storeNoWorkspace = FakeArtifactStore().also { it.blobs[preImageHash] = "ORIGINAL".toByteArray() }
val noWorkspaceEventStore = FakeEventStore(listOf(storedEvent(fileWritten, sessionId, 1L)))
val serviceBootOnly = SessionUndoService(noWorkspaceEventStore, storeNoWorkspace, setOf(bootDir))
val summaryBootOnly = serviceBootOnly.undo(sessionId)
assertEquals(1, summaryBootOnly.rejected, "boot-only roots must jail the workspace-dir file when no bound event exists")
}
/**
* A session with NO [SessionWorkspaceBoundEvent] uses boot roots exactly as before.
* Files inside boot roots are reverted; files outside are rejected.
*/
@Test
fun `undo with no SessionWorkspaceBoundEvent falls back to boot roots only`(): Unit = runBlocking {
val bootDir = Files.createTempDirectory("undo-boot-compat").toRealPath()
val outsideDir = Files.createTempDirectory("undo-outside").toRealPath()
val fakeStore = FakeArtifactStore()
val preImageHash = fakeStore.put("PREV".toByteArray()).value
val sessionId = SessionId("session-no-workspace")
// File inside boot root — should be reverted.
val insideTarget = bootDir.resolve("inside.txt")
Files.writeString(insideTarget, "NEW")
// File outside boot root and no bound event — should be rejected.
val outsideTarget = outsideDir.resolve("outside.txt")
Files.writeString(outsideTarget, "NEW")
val insideEvent = FileWrittenEvent(
invocationId = ToolInvocationId("inv-in"),
sessionId = sessionId,
path = insideTarget.toString(),
preImageHash = preImageHash,
postImageHash = null,
preExisted = true,
timestampMs = 0L,
)
val outsideEvent = FileWrittenEvent(
invocationId = ToolInvocationId("inv-out"),
sessionId = sessionId,
path = outsideTarget.toString(),
preImageHash = preImageHash,
postImageHash = null,
preExisted = true,
timestampMs = 0L,
)
val fakeEventStore = FakeEventStore(
listOf(storedEvent(insideEvent, sessionId, 1L), storedEvent(outsideEvent, sessionId, 2L)),
)
val service = SessionUndoService(fakeEventStore, fakeStore, setOf(bootDir))
val summary = service.undo(sessionId)
assertEquals(1, summary.reverted, "file inside boot root must be reverted")
assertEquals(1, summary.rejected, "file outside boot root must be rejected when no workspace event exists")
}
/**
* When [SessionWorkspaceBoundEvent.allowedPaths] carries additional paths beyond
* [SessionWorkspaceBoundEvent.workspaceRoot], those paths are also included in the
* effective jail roots.
*/
@Test
fun `undo includes allowedPaths from SessionWorkspaceBoundEvent`(): Unit = runBlocking {
val bootDir = Files.createTempDirectory("undo-boot-allowed").toRealPath()
val workspaceDir = Files.createTempDirectory("undo-ws-root").toRealPath()
val extraAllowedDir = Files.createTempDirectory("undo-extra-allowed").toRealPath()
val target = extraAllowedDir.resolve("extra.txt")
Files.writeString(target, "STALE")
val fakeStore = FakeArtifactStore()
val preImageHash = fakeStore.put("FRESH".toByteArray()).value
val sessionId = SessionId("session-allowed-paths")
val workspaceBound = SessionWorkspaceBoundEvent(
sessionId = sessionId,
workspaceRoot = workspaceDir.toString(),
allowedPaths = listOf(extraAllowedDir.toString()),
)
val fileWritten = FileWrittenEvent(
invocationId = ToolInvocationId("inv-extra"),
sessionId = sessionId,
path = target.toString(),
preImageHash = preImageHash,
postImageHash = null,
preExisted = true,
timestampMs = 0L,
)
val fakeEventStore = FakeEventStore(
listOf(storedEvent(workspaceBound, sessionId, 1L), storedEvent(fileWritten, sessionId, 2L)),
)
val service = SessionUndoService(fakeEventStore, fakeStore, setOf(bootDir))
val summary = service.undo(sessionId)
assertEquals(1, summary.reverted)
assertEquals(0, summary.rejected, "extra allowed path must not be jailed")
assertEquals("FRESH", Files.readString(target))
}
/**
* When multiple [SessionWorkspaceBoundEvent]s appear the LAST one wins.
*/
@Test
fun `undo uses last SessionWorkspaceBoundEvent when multiple are present`(): Unit = runBlocking {
val bootDir = Files.createTempDirectory("undo-boot-multi").toRealPath()
val firstWorkspace = Files.createTempDirectory("undo-ws-first").toRealPath()
val lastWorkspace = Files.createTempDirectory("undo-ws-last").toRealPath()
val target = lastWorkspace.resolve("last-ws-file.txt")
Files.writeString(target, "MODIFIED")
val fakeStore = FakeArtifactStore()
val preImageHash = fakeStore.put("ORIGINAL".toByteArray()).value
val sessionId = SessionId("session-multi-workspace")
val firstBound = SessionWorkspaceBoundEvent(
sessionId = sessionId,
workspaceRoot = firstWorkspace.toString(),
allowedPaths = emptyList(),
)
val lastBound = SessionWorkspaceBoundEvent(
sessionId = sessionId,
workspaceRoot = lastWorkspace.toString(),
allowedPaths = emptyList(),
)
val fileWritten = FileWrittenEvent(
invocationId = ToolInvocationId("inv-multi"),
sessionId = sessionId,
path = target.toString(),
preImageHash = preImageHash,
postImageHash = null,
preExisted = true,
timestampMs = 0L,
)
val fakeEventStore = FakeEventStore(
listOf(
storedEvent(firstBound, sessionId, 1L),
storedEvent(lastBound, sessionId, 2L),
storedEvent(fileWritten, sessionId, 3L),
),
)
val service = SessionUndoService(fakeEventStore, fakeStore, setOf(bootDir))
val summary = service.undo(sessionId)
assertEquals(1, summary.reverted)
assertEquals(0, summary.rejected, "file under last workspace must be reverted")
assertEquals("ORIGINAL", Files.readString(target))
}
}