diff --git a/apps/server/src/main/kotlin/com/correx/apps/server/Main.kt b/apps/server/src/main/kotlin/com/correx/apps/server/Main.kt index 1a08fccf..7ebdcdcc 100644 --- a/apps/server/src/main/kotlin/com/correx/apps/server/Main.kt +++ b/apps/server/src/main/kotlin/com/correx/apps/server/Main.kt @@ -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, + 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() diff --git a/apps/server/src/main/kotlin/com/correx/apps/server/undo/SessionUndoService.kt b/apps/server/src/main/kotlin/com/correx/apps/server/undo/SessionUndoService.kt index 2d530d29..3d2f92bf 100644 --- a/apps/server/src/main/kotlin/com/correx/apps/server/undo/SessionUndoService.kt +++ b/apps/server/src/main/kotlin/com/correx/apps/server/undo/SessionUndoService.kt @@ -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, ) { suspend fun undo(sessionId: SessionId): UndoSummary { - val mutations = eventStore.read(sessionId) + val events = eventStore.read(sessionId) + + val sessionRoots: Set = 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() diff --git a/apps/server/src/test/kotlin/com/correx/apps/server/lifecycle/LifecycleTestSupport.kt b/apps/server/src/test/kotlin/com/correx/apps/server/lifecycle/LifecycleTestSupport.kt index 8f9886e6..f567d527 100644 --- a/apps/server/src/test/kotlin/com/correx/apps/server/lifecycle/LifecycleTestSupport.kt +++ b/apps/server/src/test/kotlin/com/correx/apps/server/lifecycle/LifecycleTestSupport.kt @@ -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 { diff --git a/apps/server/src/test/kotlin/com/correx/apps/server/undo/SessionUndoServiceTest.kt b/apps/server/src/test/kotlin/com/correx/apps/server/undo/SessionUndoServiceTest.kt index 03165c3e..5d8379aa 100644 --- a/apps/server/src/test/kotlin/com/correx/apps/server/undo/SessionUndoServiceTest.kt +++ b/apps/server/src/test/kotlin/com/correx/apps/server/undo/SessionUndoServiceTest.kt @@ -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)) + } } diff --git a/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/DefaultSessionOrchestrator.kt b/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/DefaultSessionOrchestrator.kt index c85e8313..b1689614 100644 --- a/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/DefaultSessionOrchestrator.kt +++ b/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/DefaultSessionOrchestrator.kt @@ -50,6 +50,7 @@ class DefaultSessionOrchestrator( graph: WorkflowGraph, config: OrchestrationConfig, ): WorkflowResult { + val effectives = effectivesFor(config) log.debug("[Orchestrator] session={} workflow={} start={}", sessionId.value, graph.id, graph.start.value) emitWorkflowStarted(sessionId, graph, config) @@ -57,14 +58,14 @@ class DefaultSessionOrchestrator( val enriched = base.enrich() // Execute the start stage before entering the step loop - return when (val result = enterStage(enriched, graph.start)) { - is StepResult.Continue -> step(result.ctx) + return when (val result = enterStage(enriched, graph.start, effectives)) { + is StepResult.Continue -> step(result.ctx, effectives) is StepResult.Terminal -> result.result } } @Suppress("LongMethod") - private tailrec suspend fun step(ctx: EnrichedExecutionContext): WorkflowResult { + private tailrec suspend fun step(ctx: EnrichedExecutionContext, effectives: RunEffectives): WorkflowResult { log.debug( "[Orchestrator] step session={} stage={} stageCount={}", ctx.sessionId.value, ctx.currentStageId.value, ctx.stageCount, @@ -106,8 +107,8 @@ class DefaultSessionOrchestrator( enriched.sessionId.value, enriched.currentStageId.value, decision::class.simpleName, ) return when (decision) { - is TransitionDecision.Move -> when (val r = executeMove(enriched, decision)) { - is StepResult.Continue -> step(r.ctx) + is TransitionDecision.Move -> when (val r = executeMove(enriched, decision, effectives)) { + is StepResult.Continue -> step(r.ctx, effectives) is StepResult.Terminal -> r.result } @@ -144,13 +145,14 @@ class DefaultSessionOrchestrator( graph: WorkflowGraph, config: OrchestrationConfig, ): WorkflowResult { + val effectives = effectivesFor(config) val stageId = orchestrationRepository.getState(sessionId).currentStageId ?: return WorkflowResult.Failed(sessionId, "resume: no currentStageId", retryExhausted = false) log.info("[Orchestrator] resuming session={} workflow={} stage={}", sessionId.value, graph.id, stageId.value) val base = ExecutionContext(graph, sessionId, 0, stageId, config, null, null) val enriched = base.enrich() - return when (val result = enterStage(enriched, stageId)) { - is StepResult.Continue -> step(result.ctx.copy(currentStageId = stageId)) + return when (val result = enterStage(enriched, stageId, effectives)) { + is StepResult.Continue -> step(result.ctx.copy(currentStageId = stageId), effectives) is StepResult.Terminal -> result.result } } @@ -191,6 +193,7 @@ class DefaultSessionOrchestrator( private suspend fun executeMove( ctx: EnrichedExecutionContext, decision: TransitionDecision.Move, + effectives: RunEffectives, ): StepResult { val nextStageId = decision.to @@ -201,7 +204,7 @@ class DefaultSessionOrchestrator( ) } - return when (val result = enterStage(ctx, nextStageId)) { + return when (val result = enterStage(ctx, nextStageId, effectives)) { is StepResult.Continue -> StepResult.Continue( result.ctx.copy( currentStageId = advanceStage(ctx.sessionId, ctx.currentStageId, decision), @@ -216,13 +219,14 @@ class DefaultSessionOrchestrator( private suspend fun enterStage( ctx: EnrichedExecutionContext, stageId: StageId, + effectives: RunEffectives, ): StepResult { log.debug("[Orchestrator] executeStage session=${ctx.sessionId.value} stage=${stageId.value}") while (true) { if (isCancelled(ctx.sessionId)) { return StepResult.Terminal(handleCancellation(ctx.sessionId, stageId)) } - when (val result = executeStage(ctx.sessionId, stageId, ctx.graph, ctx.session, ctx.config)) { + when (val result = executeStage(ctx.sessionId, stageId, ctx.graph, ctx.session, ctx.config, effectives)) { is StageExecutionResult.Success -> return StepResult.Continue(ctx.copy(stageCount = ctx.stageCount + 1)) diff --git a/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/OrchestratorEngines.kt b/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/OrchestratorEngines.kt index 4363278c..0ef4393d 100644 --- a/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/OrchestratorEngines.kt +++ b/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/OrchestratorEngines.kt @@ -27,4 +27,5 @@ data class OrchestratorEngines( val toolCallAssessor: ToolCallAssessor? = null, val workspacePolicy: WorkspacePolicy? = null, val worldProbe: WorldProbe = FileSystemWorldProbe(), + val workspaceToolRegistryProvider: WorkspaceToolRegistryProvider? = null, ) diff --git a/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/ReplayOrchestrator.kt b/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/ReplayOrchestrator.kt index 18e3ced7..5235adfc 100644 --- a/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/ReplayOrchestrator.kt +++ b/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/ReplayOrchestrator.kt @@ -128,7 +128,7 @@ class ReplayOrchestrator( session: Session, ): ReplayStepResult { log.debug("[Orchestrator] executeStage session=${ctx.sessionId.value} stage=${stageId.value}") - return when (val result = executeStage(ctx.sessionId, stageId, ctx.graph, session, ctx.config)) { + return when (val result = executeStage(ctx.sessionId, stageId, ctx.graph, session, ctx.config, effectivesFor(ctx.config))) { is StageExecutionResult.Success -> ReplayStepResult.Continue( ctx.copy(stageCount = ctx.stageCount + 1), ) @@ -151,6 +151,7 @@ class ReplayOrchestrator( stageConfig: StageConfig, timeoutMs: Long, responseFormat: ResponseFormat, + effectives: RunEffectives, ): InferenceResult = when (strategy) { is ReplayStrategy.SkipInference -> { // bypass router entirely — use recorded artifact @@ -179,7 +180,7 @@ class ReplayOrchestrator( }, ) } - else -> super.runInference(sessionId, stageId, contextPack, stageConfig, timeoutMs, responseFormat) + else -> super.runInference(sessionId, stageId, contextPack, stageConfig, timeoutMs, responseFormat, effectives) } override suspend fun mapValidationOutcome( diff --git a/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/SessionOrchestrator.kt b/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/SessionOrchestrator.kt index 8d6193e6..4ea9089a 100644 --- a/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/SessionOrchestrator.kt +++ b/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/SessionOrchestrator.kt @@ -143,6 +143,7 @@ abstract class SessionOrchestrator( private val toolCallAssessor: ToolCallAssessor? = engines.toolCallAssessor private val workspacePolicy: WorkspacePolicy? = engines.workspacePolicy private val worldProbe: WorldProbe = engines.worldProbe + private val workspaceToolRegistryProvider: WorkspaceToolRegistryProvider? = engines.workspaceToolRegistryProvider private val inferenceRepository: InferenceRepository = repositories.inferenceRepository internal val orchestrationRepository: OrchestrationRepository = repositories.orchestrationRepository protected open val tokenizer: Tokenizer? = null @@ -165,6 +166,24 @@ abstract class SessionOrchestrator( abstract suspend fun cancel(sessionId: SessionId) + // --- per-run effective registry + policy --- + + internal data class RunEffectives( + val registry: ToolRegistry?, + val executor: ToolExecutor?, + val policy: WorkspacePolicy?, + ) + + internal fun effectivesFor(config: OrchestrationConfig): RunEffectives { + val wsTools = config.workspace?.let { workspaceToolRegistryProvider?.forWorkspace(it) } + val registry = wsTools?.registry ?: toolRegistry + val executor = wsTools?.executor ?: toolExecutor + val policy = config.workspace + ?.let { WorkspacePolicy(it.workspaceRoot, it.privilegedLocations) } + ?: workspacePolicy + return RunEffectives(registry, executor, policy) + } + // --- stage execution --- @Suppress("CyclomaticComplexMethod") @@ -174,6 +193,7 @@ abstract class SessionOrchestrator( graph: WorkflowGraph, session: Session, config: OrchestrationConfig, + effectives: RunEffectives, ): StageExecutionResult { val stageConfig = requireNotNull(graph.stages[stageId]) { "Stage '${stageId.value}' not declared in workflow graph" @@ -272,7 +292,7 @@ abstract class SessionOrchestrator( var currentContext = contextPack var inferenceResult = runInference( - sessionId, stageId, currentContext, stageConfig, config.stageTimeoutMs, responseFormat, + sessionId, stageId, currentContext, stageConfig, config.stageTimeoutMs, responseFormat, effectives, ) var toolRounds = 0 while ( @@ -287,7 +307,7 @@ abstract class SessionOrchestrator( break } val toolEntries = dispatchToolCalls(sessionId, stageId, - inferenceResult.response.toolCalls, stageConfig) + inferenceResult.response.toolCalls, stageConfig, effectives) val fatalEntry = toolEntries.firstOrNull { it.content.startsWith("FATAL:") } if (fatalEntry != null) { emitProcessResultEvents(sessionId, stageId, stageConfig) @@ -314,7 +334,7 @@ abstract class SessionOrchestrator( ) emitContextTruncationIfNeeded(sessionId, stageId, currentContext) inferenceResult = runInference( - sessionId, stageId, currentContext, stageConfig, config.stageTimeoutMs, responseFormat, + sessionId, stageId, currentContext, stageConfig, config.stageTimeoutMs, responseFormat, effectives, ) toolRounds++ } @@ -326,7 +346,7 @@ abstract class SessionOrchestrator( val validationCtx = ValidationContext( graph = graph, sessionState = session.state, - availableTools = toolRegistry?.all()?.map { it.name }?.toSet() + availableTools = effectives.registry?.all()?.map { it.name }?.toSet() ) when (val outcome = mapValidationOutcome(sessionId, stageId, validationCtx)) { is StageExecutionResult.Success -> { @@ -353,8 +373,9 @@ abstract class SessionOrchestrator( stageId: StageId, toolCalls: List, stageConfig: StageConfig, + effectives: RunEffectives, ): List { - val executor = toolExecutor ?: return emptyList() + val executor = effectives.executor ?: return emptyList() val processResultSlots = stageConfig.produces.filter { it.kind.id == "process_result" } return toolCalls.flatMap { toolCall -> val invocationId = ToolInvocationId(UUID.randomUUID().toString()) @@ -372,7 +393,7 @@ abstract class SessionOrchestrator( toolName = toolCall.function.name, parameters = parameters, ) - val tool = toolRegistry?.resolve(toolCall.function.name) + val tool = effectives.registry?.resolve(toolCall.function.name) val tier = tool?.tier ?: Tier.T2 emit( sessionId, @@ -386,7 +407,7 @@ abstract class SessionOrchestrator( ), ) val plane2Risk: RiskSummary? = runPlane2Assessment( - sessionId, stageId, invocationId, toolCall.function.name, request, tool, + sessionId, stageId, invocationId, toolCall.function.name, request, tool, effectives, )?.let { assessment -> if (assessment.recommendedAction == RiskAction.BLOCK) { emit( @@ -607,9 +628,10 @@ abstract class SessionOrchestrator( toolName: String, request: ToolRequest, tool: Tool?, + effectives: RunEffectives, ): RiskSummary? { val assessor = toolCallAssessor ?: return null - val policy = workspacePolicy ?: return null + val policy = effectives.policy ?: return null val assessment = assessor.assess( ToolCallAssessmentInput( request = request, @@ -771,6 +793,7 @@ abstract class SessionOrchestrator( stageConfig: StageConfig, timeoutMs: Long, responseFormat: ResponseFormat = ResponseFormat.Text, + effectives: RunEffectives = RunEffectives(toolRegistry, toolExecutor, workspacePolicy), ): InferenceResult { val provider = inferenceRouter.route(stageId, stageConfig.requiredCapabilities, stageConfig.modelId) log.debug( @@ -786,7 +809,7 @@ abstract class SessionOrchestrator( generationConfig = stageConfig.generationConfig, responseFormat = responseFormat, tools = stageConfig.allowedTools - .mapNotNull { toolRegistry?.resolve(it) } + .mapNotNull { effectives.registry?.resolve(it) } .map { tool -> ToolDefinition( function = ToolFunction( diff --git a/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/WorkspaceToolRegistryProvider.kt b/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/WorkspaceToolRegistryProvider.kt new file mode 100644 index 00000000..60bec3fc --- /dev/null +++ b/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/WorkspaceToolRegistryProvider.kt @@ -0,0 +1,13 @@ +package com.correx.core.kernel.orchestration + +import com.correx.core.tools.contract.ToolExecutor +import com.correx.core.tools.registry.ToolRegistry + +data class WorkspaceTools( + val registry: ToolRegistry, + val executor: ToolExecutor, +) + +fun interface WorkspaceToolRegistryProvider { + fun forWorkspace(workspace: WorkspaceContext): WorkspaceTools +} diff --git a/testing/integration/build.gradle b/testing/integration/build.gradle index ca53504c..f5e9b574 100644 --- a/testing/integration/build.gradle +++ b/testing/integration/build.gradle @@ -25,4 +25,6 @@ dependencies { testImplementation(project(":testing:contracts")) testImplementation(testFixtures(project(":testing:contracts"))) testImplementation("org.jetbrains.kotlinx:kotlinx-coroutines-test:1.9.0") + testImplementation(project(":infrastructure:tools")) + testImplementation(project(":infrastructure:tools:filesystem")) } \ No newline at end of file diff --git a/testing/integration/src/test/kotlin/WorkspaceScopedToolRegistryTest.kt b/testing/integration/src/test/kotlin/WorkspaceScopedToolRegistryTest.kt new file mode 100644 index 00000000..cb46b94f --- /dev/null +++ b/testing/integration/src/test/kotlin/WorkspaceScopedToolRegistryTest.kt @@ -0,0 +1,375 @@ +import com.correx.core.approvals.ApprovalOutcome +import com.correx.core.approvals.ApprovalProjector +import com.correx.core.approvals.ApprovalStatus +import com.correx.core.approvals.DefaultApprovalReducer +import com.correx.core.approvals.DefaultApprovalRepository +import com.correx.core.approvals.domain.ApprovalEngine +import com.correx.core.approvals.model.ApprovalContext +import com.correx.core.approvals.model.ApprovalDecision +import com.correx.core.approvals.model.ApprovalGrant +import com.correx.core.approvals.model.DomainApprovalRequest +import com.correx.core.artifacts.DefaultArtifactReducer +import com.correx.core.events.events.ToolCallAssessedEvent +import com.correx.core.events.execution.RetryPolicy +import com.correx.core.events.risk.RiskAction +import com.correx.core.events.types.ProviderId +import com.correx.core.events.types.SessionId +import com.correx.core.events.types.StageId +import com.correx.core.events.types.TransitionId +import com.correx.core.inference.CapabilityScore +import com.correx.core.inference.FinishReason +import com.correx.core.inference.InferenceProvider +import com.correx.core.inference.InferenceRepository +import com.correx.core.inference.InferenceRequest +import com.correx.core.inference.InferenceResponse +import com.correx.core.inference.InferenceRouter +import com.correx.core.inference.InferenceState +import com.correx.core.inference.ModelCapability +import com.correx.core.inference.ProviderHealth +import com.correx.core.inference.Tokenizer +import com.correx.core.inference.ToolCallFunction +import com.correx.core.inference.ToolCallRequest +import com.correx.core.inference.TokenUsage +import com.correx.core.kernel.orchestration.DefaultOrchestrationReducer +import com.correx.core.kernel.orchestration.DefaultSessionOrchestrator +import com.correx.core.kernel.orchestration.OrchestrationConfig +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.sessions.DefaultSessionReducer +import com.correx.core.sessions.DefaultSessionRepository +import com.correx.core.sessions.SessionProjector +import com.correx.core.sessions.projections.replay.DefaultEventReplayer +import com.correx.core.sessions.projections.replay.EventReplayer +import com.correx.core.toolintent.ToolCallAssessor +import com.correx.core.toolintent.WorkspacePolicy +import com.correx.core.toolintent.rules.PathContainmentRule +import com.correx.core.tools.registry.ToolRegistry +import com.correx.core.transitions.graph.StageConfig +import com.correx.core.transitions.graph.TransitionEdge +import com.correx.core.transitions.graph.WorkflowGraph +import com.correx.core.transitions.resolution.DefaultTransitionResolver +import com.correx.core.validation.pipeline.ValidationPipeline +import com.correx.infrastructure.persistence.InMemoryEventStore +import com.correx.infrastructure.persistence.artifact.LiveArtifactRepository +import com.correx.infrastructure.tools.DefaultToolRegistry +import com.correx.infrastructure.tools.DispatchingToolExecutor +import com.correx.infrastructure.tools.FileEditConfig +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.buildTools +import com.correx.testing.contracts.fixtures.artifactstore.NoopArtifactStore +import com.correx.testing.fixtures.context.ContextFixtures +import com.correx.testing.fixtures.cyclePolicyMissingValidator +import com.correx.testing.fixtures.inference.MockTokenizer +import kotlinx.coroutines.runBlocking +import kotlinx.datetime.Instant +import java.nio.file.Files +import kotlin.test.Test +import kotlin.test.assertFalse +import kotlin.test.assertNotNull +import kotlin.test.assertTrue + +@Suppress("LongMethod") +class WorkspaceScopedToolRegistryTest { + + private inner class FileWriteCallingProvider( + private val targetPath: String, + private val content: String, + ) : InferenceProvider { + override val id = ProviderId("file-write-caller") + override val name = "file-write-caller" + override val tokenizer: Tokenizer = MockTokenizer() + private var callCount = 0 + + override suspend fun infer(request: InferenceRequest): InferenceResponse { + callCount++ + return if (callCount == 1) { + InferenceResponse( + requestId = request.requestId, + text = "", + finishReason = FinishReason.ToolCall, + tokensUsed = TokenUsage(10, 5), + latencyMs = 0, + toolCalls = listOf( + ToolCallRequest( + id = "tc-write", + function = ToolCallFunction( + name = "file_write", + arguments = """{"path":"$targetPath","operation":"write","content":"$content"}""", + ), + ), + ), + ) + } else { + InferenceResponse( + requestId = request.requestId, + text = "done", + finishReason = FinishReason.Stop, + tokensUsed = TokenUsage(10, 5), + latencyMs = 0, + ) + } + } + + override suspend fun healthCheck(): ProviderHealth = ProviderHealth.Healthy + override fun capabilities(): Set = setOf(CapabilityScore(ModelCapability.General, 1.0)) + } + + private fun buildToolsForWorkspace(workspace: WorkspaceContext): WorkspaceTools { + val registry = DefaultToolRegistry.build( + ToolConfig( + shell = ShellConfig(enabled = false), + fileRead = FileReadConfig(enabled = false), + fileWrite = FileWriteConfig( + enabled = true, + allowedPaths = workspace.allowedPaths, + workingDir = workspace.workingDir, + ), + fileEdit = FileEditConfig(enabled = false), + ).buildTools(), + ) + return WorkspaceTools(registry = registry, executor = DispatchingToolExecutor(registry)) + } + + private fun buildOrchestrator( + workspaceContext: WorkspaceContext?, + bootRegistry: ToolRegistry, + provider: InferenceProvider, + ): Pair { + val eventStore = InMemoryEventStore() + val artifactStore = NoopArtifactStore() + + val assessor = ToolCallAssessor(listOf(PathContainmentRule())) + val bootPolicyRoot = workspaceContext?.workspaceRoot + ?: java.nio.file.Path.of(System.getProperty("java.io.tmpdir")) + val bootPolicy = WorkspacePolicy(bootPolicyRoot) + + val inferenceRouter = object : InferenceRouter { + override suspend fun route(stageId: StageId, requiredCapabilities: Set) = provider + } + + val approvalRepository = DefaultApprovalRepository( + DefaultEventReplayer(eventStore, ApprovalProjector(DefaultApprovalReducer())), + ) + + val repositories = OrchestratorRepositories( + eventStore = eventStore, + inferenceRepository = InferenceRepository(object : EventReplayer { + override fun rebuild(sessionId: SessionId) = InferenceState() + }), + orchestrationRepository = OrchestrationRepository( + DefaultEventReplayer(eventStore, OrchestrationProjector(DefaultOrchestrationReducer())), + ), + sessionRepository = DefaultSessionRepository( + DefaultEventReplayer(eventStore, SessionProjector(DefaultSessionReducer())), + ), + artifactRepository = LiveArtifactRepository(eventStore, DefaultArtifactReducer()), + approvalRepository = approvalRepository, + ) + + val workspaceProvider = WorkspaceToolRegistryProvider { ws -> buildToolsForWorkspace(ws) } + + val autoApproveEngine = object : ApprovalEngine { + override fun evaluate( + request: DomainApprovalRequest, + context: ApprovalContext, + grants: List, + now: Instant, + ) = ApprovalDecision( + id = null, + requestId = request.id, + outcome = ApprovalOutcome.AUTO_APPROVED, + state = ApprovalStatus.COMPLETED, + tier = request.tier, + contextSnapshot = context, + resolutionTimestamp = now, + reason = "test-auto-approve", + ) + } + + val engines = OrchestratorEngines( + transitionResolver = DefaultTransitionResolver { _, _ -> true }, + contextPackBuilder = ContextFixtures.simpleBuilder(), + inferenceRouter = inferenceRouter, + validationPipeline = ValidationPipeline(validators = listOf(cyclePolicyMissingValidator())), + approvalEngine = autoApproveEngine, + riskAssessor = DefaultRiskAssessor(), + toolRegistry = bootRegistry, + toolExecutor = DispatchingToolExecutor(bootRegistry), + toolCallAssessor = assessor, + workspacePolicy = bootPolicy, + workspaceToolRegistryProvider = workspaceProvider, + ) + + return DefaultSessionOrchestrator( + repositories = repositories, + engines = engines, + retryCoordinator = DefaultRetryCoordinator(eventStore), + artifactStore = artifactStore, + ) to eventStore + } + + private fun singleStageGraph(): WorkflowGraph = WorkflowGraph( + id = "ws-test", + stages = mapOf( + StageId("A") to StageConfig(allowedTools = setOf("file_write")), + ), + transitions = setOf( + TransitionEdge(TransitionId("t1"), StageId("A"), StageId("done"), condition = { true }), + ), + start = StageId("A"), + ) + + @Test + fun `write inside workspace A succeeds and file exists`(): Unit = runBlocking { + val workspaceA = Files.createTempDirectory("ws-a") + val targetFile = workspaceA.resolve("output.txt") + + val workspace = WorkspaceContext( + workspaceRoot = workspaceA, + workingDir = workspaceA, + allowedPaths = setOf(workspaceA), + ) + val provider = FileWriteCallingProvider(targetFile.toString(), "hello workspace A") + val bootRegistry = DefaultToolRegistry.build(emptyList()) + val (orchestrator, _) = buildOrchestrator(workspace, bootRegistry, provider) + + val config = OrchestrationConfig( + retryPolicy = RetryPolicy(maxAttempts = 1, backoffMs = 0), + workspace = workspace, + ) + orchestrator.run(SessionId("ws-a-ok"), singleStageGraph(), config) + + assertTrue(Files.exists(targetFile), "file_write inside workspace A must create the file") + } + + @Test + fun `write outside workspace A is rejected by jail`(): Unit = runBlocking { + val workspaceA = Files.createTempDirectory("ws-a-jail") + val outsideDir = Files.createTempDirectory("ws-outside") + val targetFile = outsideDir.resolve("escape.txt") + + val workspace = WorkspaceContext( + workspaceRoot = workspaceA, + workingDir = workspaceA, + allowedPaths = setOf(workspaceA), + ) + val provider = FileWriteCallingProvider(targetFile.toString(), "should be blocked") + val bootRegistry = DefaultToolRegistry.build(emptyList()) + val (orchestrator, _) = buildOrchestrator(workspace, bootRegistry, provider) + + val config = OrchestrationConfig( + retryPolicy = RetryPolicy(maxAttempts = 1, backoffMs = 0), + workspace = workspace, + ) + orchestrator.run(SessionId("ws-a-jail"), singleStageGraph(), config) + + assertFalse(Files.exists(targetFile), "file_write outside workspace A must NOT create the file") + } + + @Test + fun `two sessions with different workspaces are independently jailed`(): Unit = runBlocking { + val workspaceA = Files.createTempDirectory("ws-aa") + val workspaceB = Files.createTempDirectory("ws-bb") + val fileA = workspaceA.resolve("a.txt") + val fileB = workspaceB.resolve("b.txt") + + val ctxA = WorkspaceContext( + workspaceRoot = workspaceA, + workingDir = workspaceA, + allowedPaths = setOf(workspaceA), + ) + val ctxB = WorkspaceContext( + workspaceRoot = workspaceB, + workingDir = workspaceB, + allowedPaths = setOf(workspaceB), + ) + + val bootRegistry = DefaultToolRegistry.build(emptyList()) + val (orchA, _) = buildOrchestrator(ctxA, bootRegistry, FileWriteCallingProvider(fileA.toString(), "session A")) + val (orchB, _) = buildOrchestrator(ctxB, bootRegistry, FileWriteCallingProvider(fileB.toString(), "session B")) + + val configA = OrchestrationConfig(retryPolicy = RetryPolicy(maxAttempts = 1, backoffMs = 0), workspace = ctxA) + val configB = OrchestrationConfig(retryPolicy = RetryPolicy(maxAttempts = 1, backoffMs = 0), workspace = ctxB) + + orchA.run(SessionId("session-a"), singleStageGraph(), configA) + orchB.run(SessionId("session-b"), singleStageGraph(), configB) + + assertTrue(Files.exists(fileA), "session A must write to workspace A") + assertTrue(Files.exists(fileB), "session B must write to workspace B") + assertFalse(Files.exists(workspaceB.resolve("a.txt")), "session A must NOT write to workspace B") + assertFalse(Files.exists(workspaceA.resolve("b.txt")), "session B must NOT write to workspace A") + } + + @Test + fun `null workspace uses boot registry — regression guard`(): Unit = runBlocking { + val bootWorkspace = Files.createTempDirectory("ws-boot") + val targetFile = bootWorkspace.resolve("boot.txt") + + val bootRegistry = DefaultToolRegistry.build( + ToolConfig( + shell = ShellConfig(enabled = false), + fileRead = FileReadConfig(enabled = false), + fileWrite = FileWriteConfig( + enabled = true, + allowedPaths = setOf(bootWorkspace), + workingDir = bootWorkspace, + ), + fileEdit = FileEditConfig(enabled = false), + ).buildTools(), + ) + + val provider = FileWriteCallingProvider(targetFile.toString(), "boot content") + val (orchestrator, _) = buildOrchestrator(null, bootRegistry, provider) + + val config = OrchestrationConfig( + retryPolicy = RetryPolicy(maxAttempts = 1, backoffMs = 0), + workspace = null, + ) + orchestrator.run(SessionId("null-workspace"), singleStageGraph(), config) + + assertTrue(Files.exists(targetFile), "null-workspace run must use boot registry — file must be written") + } + + @Test + fun `plane-2 assessment uses session workspace policy not boot policy`(): Unit = runBlocking { + val workspaceA = Files.createTempDirectory("ws-plane2") + val outsideDir = Files.createTempDirectory("ws-plane2-outside") + val targetFile = outsideDir.resolve("plane2.txt") + + val workspace = WorkspaceContext( + workspaceRoot = workspaceA, + workingDir = workspaceA, + allowedPaths = setOf(workspaceA), + ) + val provider = FileWriteCallingProvider(targetFile.toString(), "test") + val bootRegistry = DefaultToolRegistry.build(emptyList()) + val (orchestrator, eventStore) = buildOrchestrator(workspace, bootRegistry, provider) + + val config = OrchestrationConfig( + retryPolicy = RetryPolicy(maxAttempts = 1, backoffMs = 0), + workspace = workspace, + ) + val sessionId = SessionId("plane2-session") + orchestrator.run(sessionId, singleStageGraph(), config) + + val assessed = eventStore.read(sessionId) + .mapNotNull { it.payload as? ToolCallAssessedEvent } + .firstOrNull() + + assertNotNull(assessed, "ToolCallAssessedEvent must be emitted when path is outside session workspace") + assertTrue( + assessed.disposition != RiskAction.PROCEED, + "plane-2 disposition for out-of-workspace write must not be PROCEED, got ${assessed.disposition}", + ) + } +}