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
@@ -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))
@@ -27,4 +27,5 @@ data class OrchestratorEngines(
val toolCallAssessor: ToolCallAssessor? = null,
val workspacePolicy: WorkspacePolicy? = null,
val worldProbe: WorldProbe = FileSystemWorldProbe(),
val workspaceToolRegistryProvider: WorkspaceToolRegistryProvider? = null,
)
@@ -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(
@@ -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<ToolCallRequest>,
stageConfig: StageConfig,
effectives: RunEffectives,
): List<ContextEntry> {
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(
@@ -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
}