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:
@@ -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"))
|
||||
}
|
||||
@@ -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<CapabilityScore> = 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<DefaultSessionOrchestrator, InMemoryEventStore> {
|
||||
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<ModelCapability>) = provider
|
||||
}
|
||||
|
||||
val approvalRepository = DefaultApprovalRepository(
|
||||
DefaultEventReplayer(eventStore, ApprovalProjector(DefaultApprovalReducer())),
|
||||
)
|
||||
|
||||
val repositories = OrchestratorRepositories(
|
||||
eventStore = eventStore,
|
||||
inferenceRepository = InferenceRepository(object : EventReplayer<InferenceState> {
|
||||
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<ApprovalGrant>,
|
||||
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<com.correx.core.tools.contract.Tool>())
|
||||
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<com.correx.core.tools.contract.Tool>())
|
||||
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<com.correx.core.tools.contract.Tool>())
|
||||
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<com.correx.core.tools.contract.Tool>())
|
||||
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}",
|
||||
)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user