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.Tier import com.correx.core.approvals.UserSteering import com.correx.core.approvals.domain.DefaultApprovalEngine import com.correx.core.approvals.model.ApprovalContext import com.correx.core.approvals.model.ApprovalDecision import com.correx.core.approvals.model.ApprovalScopeIdentity import com.correx.core.sessions.ApprovalMode import kotlinx.datetime.Clock import com.correx.core.artifacts.DefaultArtifactReducer import com.correx.core.events.events.ApprovalRequestedEvent import com.correx.core.events.events.ToolCallAssessedEvent import com.correx.core.events.events.ToolExecutionRejectedEvent import com.correx.core.events.events.ToolInvocationRequestedEvent import com.correx.core.events.events.ToolRequest 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.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.ToolCallAssessment import com.correx.core.toolintent.ToolCallAssessmentInput import com.correx.core.toolintent.ToolCallAssessor import com.correx.core.toolintent.ToolCallRule import com.correx.core.toolintent.WorkspacePolicy import com.correx.core.tools.compression.CompressionRule.StripBlankLines import com.correx.core.tools.compression.DeclarativeCompressor import com.correx.core.tools.compression.OutputCompressionSpec import com.correx.core.tools.contract.Tool import com.correx.core.tools.contract.ToolCapability import com.correx.core.tools.contract.ToolExecutor import com.correx.core.tools.contract.ToolResult import com.correx.core.tools.contract.ValidationResult 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.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.launch import kotlinx.coroutines.runBlocking import kotlinx.coroutines.withTimeout import kotlinx.coroutines.yield import kotlinx.serialization.json.JsonObject import kotlinx.serialization.json.buildJsonObject import java.nio.file.Path import java.util.concurrent.atomic.AtomicBoolean import kotlin.test.Test import kotlin.test.assertEquals import kotlin.test.assertNotNull import kotlin.test.assertNull import kotlin.test.assertTrue @Suppress("LongMethod") class ToolCallGateTest { private inner class FakeFileWriteTool(override val tier: Tier = Tier.T1) : Tool { override val name = "file_write" override val description = "fake file write" override val parametersSchema: JsonObject = buildJsonObject {} override val requiredCapabilities: Set = setOf(ToolCapability.FILE_WRITE) override fun validateRequest(request: ToolRequest): ValidationResult = ValidationResult.Valid } /** A tool that declares an output compressor (strips blank lines) — the Tier-A seam under test. */ private inner class FakeCompressingTool(override val tier: Tier = Tier.T1) : Tool { override val name = "file_write" override val description = "fake compressing tool" override val parametersSchema: JsonObject = buildJsonObject {} override val requiredCapabilities: Set = setOf(ToolCapability.FILE_WRITE) override val outputCompressor = DeclarativeCompressor(OutputCompressionSpec(listOf(StripBlankLines))) override fun validateRequest(request: ToolRequest): ValidationResult = ValidationResult.Valid } private inner class FakeToolRegistry(private val tool: Tool) : ToolRegistry { override fun resolve(name: String): Tool? = if (name == tool.name) tool else null override fun all(): List = listOf(tool) } private inner class RecordingExecutor : ToolExecutor { val executeCalled = AtomicBoolean(false) override suspend fun execute(request: ToolRequest): ToolResult { executeCalled.set(true) return ToolResult.Success(invocationId = request.invocationId, output = "ok", exitCode = 0) } } private inner class ToolCallingProvider(private val toolName: String) : InferenceProvider { override val id = ProviderId("tool-caller") override val name = "tool-caller" override val tokenizer: Tokenizer = MockTokenizer() val requests = mutableListOf() private var callCount = 0 override suspend fun infer(request: InferenceRequest): InferenceResponse { requests += request callCount++ return if (callCount == 1) { InferenceResponse( requestId = request.requestId, text = "", finishReason = FinishReason.ToolCall, tokensUsed = TokenUsage(10, 5), latencyMs = 0, toolCalls = listOf( ToolCallRequest( id = "tc-1", function = ToolCallFunction(name = toolName, arguments = "{}"), ), ), ) } 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 ruleReturning(action: RiskAction): ToolCallRule = object : ToolCallRule { override fun appliesTo(capabilities: Set) = true override fun assess(input: ToolCallAssessmentInput) = ToolCallAssessment(disposition = action) } private val workspace = Path.of("/work") private fun buildOrchestrator( executor: ToolExecutor, tool: Tool, assessorRule: ToolCallRule?, ): Triple { val eventStore = InMemoryEventStore() val artifactStore = NoopArtifactStore() val toolRegistry = FakeToolRegistry(tool) val assessor = assessorRule?.let { ToolCallAssessor(listOf(it)) } val policy = if (assessor != null) WorkspacePolicy(workspace) else null val provider = ToolCallingProvider(tool.name) val inferenceRouter = object : InferenceRouter { override suspend fun route(stageId: StageId, requiredCapabilities: Set): InferenceProvider = 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 engines = OrchestratorEngines( transitionResolver = DefaultTransitionResolver { _, _ -> true }, contextPackBuilder = ContextFixtures.simpleBuilder(), inferenceRouter = inferenceRouter, validationPipeline = ValidationPipeline(validators = listOf(cyclePolicyMissingValidator())), approvalEngine = DefaultApprovalEngine(), riskAssessor = DefaultRiskAssessor(), toolExecutor = executor, toolRegistry = toolRegistry, toolCallAssessor = assessor, workspacePolicy = policy, ) val orchestrator = DefaultSessionOrchestrator( repositories = repositories, engines = engines, retryCoordinator = DefaultRetryCoordinator(eventStore), artifactStore = artifactStore, ) return Triple(orchestrator, eventStore, provider) } private fun singleStageGraph(allowedTools: Set = setOf("file_write")): WorkflowGraph = WorkflowGraph( id = "gate-test", stages = mapOf( StageId("A") to StageConfig(allowedTools = allowedTools), ), transitions = setOf( TransitionEdge(TransitionId("t1"), StageId("A"), StageId("done"), condition = { true }), ), start = StageId("A"), ) @Test fun `BLOCK rule prevents executor from being called and emits rejected event`(): Unit = runBlocking { val executor = RecordingExecutor() val (orchestrator, eventStore, _) = buildOrchestrator(executor, FakeFileWriteTool(), ruleReturning(RiskAction.BLOCK)) val sessionId = SessionId("gate-block") val config = OrchestrationConfig(retryPolicy = RetryPolicy(maxAttempts = 1, backoffMs = 0)) orchestrator.run(sessionId, singleStageGraph(), config) val events = eventStore.read(sessionId) assertTrue(!executor.executeCalled.get(), "executor must NOT be called when BLOCK") assertNotNull(events.find { it.payload is ToolCallAssessedEvent }, "ToolCallAssessedEvent must be emitted") assertNotNull(events.find { it.payload is ToolExecutionRejectedEvent }, "ToolExecutionRejectedEvent must be emitted") } @Test fun `PROMPT_USER rule on T2 tool triggers approval with plane2 risk summary`(): Unit = runBlocking { val executor = RecordingExecutor() // Use T2 so the engine does not auto-approve (PROMPT mode auto-approves up to T1 only) val (orchestrator, eventStore, _) = buildOrchestrator( executor, FakeFileWriteTool(Tier.T2), ruleReturning(RiskAction.PROMPT_USER), ) val sessionId = SessionId("gate-prompt") val config = OrchestrationConfig(retryPolicy = RetryPolicy(maxAttempts = 1, backoffMs = 0)) // Run in background; orchestrator will pause waiting for approval val job = launch { orchestrator.run(sessionId, singleStageGraph(), config) } // Wait for ApprovalRequestedEvent withTimeout(5_000) { while (eventStore.read(sessionId).none { it.payload is ApprovalRequestedEvent }) { yield() } } val approval = eventStore.read(sessionId).firstNotNullOfOrNull { it.payload as? ApprovalRequestedEvent } assertNotNull(approval, "ApprovalRequestedEvent must be emitted for PROMPT_USER plane-2 rule on T1 tool") assertEquals(RiskAction.PROMPT_USER, approval?.riskSummary?.recommendedAction) job.cancel() job.join() } @Test fun `null assessor means executor is called normally (regression guard)`(): Unit = runBlocking { val executor = RecordingExecutor() val (orchestrator, eventStore, _) = buildOrchestrator(executor, FakeFileWriteTool(), assessorRule = null) val sessionId = SessionId("gate-null") val config = OrchestrationConfig(retryPolicy = RetryPolicy(maxAttempts = 1, backoffMs = 0)) orchestrator.run(sessionId, singleStageGraph(), config) val events = eventStore.read(sessionId) assertTrue(executor.executeCalled.get(), "executor must be called when assessor is null") assertNull(events.find { it.payload is ToolCallAssessedEvent }, "ToolCallAssessedEvent must NOT be emitted when assessor is null") } @Test fun `approving with a steering note injects it into the next inference context`(): Unit = runBlocking { val executor = RecordingExecutor() // T2 so PROMPT mode does not auto-approve — the orchestrator pauses for a human decision. val (orchestrator, eventStore, provider) = buildOrchestrator( executor, FakeFileWriteTool(Tier.T2), ruleReturning(RiskAction.PROMPT_USER), ) val sessionId = SessionId("gate-steer") val config = OrchestrationConfig(retryPolicy = RetryPolicy(maxAttempts = 1, backoffMs = 0)) val job = launch { orchestrator.run(sessionId, singleStageGraph(), config) } val request = withTimeout(5_000) { var req: ApprovalRequestedEvent? = null while (req == null) { req = eventStore.read(sessionId).firstNotNullOfOrNull { it.payload as? ApprovalRequestedEvent } if (req == null) yield() } req } val note = "add the current distro to the output" orchestrator.submitApprovalDecision( request.requestId, ApprovalDecision( id = null, requestId = request.requestId, outcome = ApprovalOutcome.APPROVED, state = ApprovalStatus.COMPLETED, tier = Tier.T2, contextSnapshot = ApprovalContext( identity = ApprovalScopeIdentity(sessionId, StageId("A"), projectId = null), mode = ApprovalMode.PROMPT, ), resolutionTimestamp = Clock.System.now(), reason = note, userSteering = UserSteering(text = note, sessionId = sessionId, timestamp = Clock.System.now()), ), ) // The approved call executes, then the loop re-infers — that second request must carry the note. withTimeout(5_000) { while (provider.requests.size < 2) yield() } job.cancel() job.join() assertTrue(executor.executeCalled.get(), "the approved tool must still execute") val steeringEntry = provider.requests[1].contextPack.layers.values.flatten().firstOrNull { it.sourceType == "steeringNote" && it.content.contains(note) } assertNotNull(steeringEntry, "approved steering note must be injected into the next inference context") } @Test fun `tool output is compressed on the path into context`(): Unit = runBlocking { val rawOutput = "line1\n\nline2\n\n\nline3" val executor = object : ToolExecutor { override suspend fun execute(request: ToolRequest): ToolResult = ToolResult.Success(invocationId = request.invocationId, output = rawOutput, exitCode = 0) } val (orchestrator, _, provider) = buildOrchestrator(executor, FakeCompressingTool(), assessorRule = null) val config = OrchestrationConfig(retryPolicy = RetryPolicy(maxAttempts = 1, backoffMs = 0)) orchestrator.run(SessionId("compress"), singleStageGraph(), config) // The second inference carries round-1's tool result in its context pack — compressed // (blank lines stripped), not raw. Raw output stays authoritative on the executor's // event path (invariant #6); the compressor only shapes the derived context entry. val secondPack = provider.requests[1].contextPack val toolEntry = secondPack.layers.values.flatten().first { it.sourceType == "toolResult" } assertEquals("line1\nline2\nline3", toolEntry.content) } @Test fun `tool not in allowedTools is not executed and emits ToolExecutionRejectedEvent`(): Unit = runBlocking { val executor = RecordingExecutor() // Stage only allows "allowed_tool"; the LLM returns a call to "file_write" (not in the set). val (orchestrator, eventStore, _) = buildOrchestrator(executor, FakeFileWriteTool(), assessorRule = null) val sessionId = SessionId("gate-stage-deny") val config = OrchestrationConfig(retryPolicy = RetryPolicy(maxAttempts = 1, backoffMs = 0)) val graph = singleStageGraph(allowedTools = setOf("allowed_tool")) orchestrator.run(sessionId, graph, config) val events = eventStore.read(sessionId) assertTrue(!executor.executeCalled.get(), "executor must NOT be called for a tool outside allowedTools") // Denial is recorded in the event log (invariant #5). val rejected = events.firstNotNullOfOrNull { it.payload as? ToolExecutionRejectedEvent } assertNotNull(rejected, "ToolExecutionRejectedEvent must be emitted when tool is not in allowedTools") assertEquals("file_write", rejected!!.toolName) assertTrue( rejected.reason.contains("not permitted"), "rejection reason should mention 'not permitted', got: ${rejected.reason}", ) // ToolInvocationRequestedEvent is still emitted so the request itself is in the log. assertTrue( events.any { (it.payload as? ToolInvocationRequestedEvent)?.toolName == "file_write" }, "ToolInvocationRequestedEvent must be emitted to record the denied request", ) } @Test fun `empty allowedTools denies all domain tools and emits ToolExecutionRejectedEvent`(): Unit = runBlocking { val executor = RecordingExecutor() // Stage has allowedTools = emptySet() — only stage_complete is permitted. val (orchestrator, eventStore, _) = buildOrchestrator(executor, FakeFileWriteTool(), assessorRule = null) val sessionId = SessionId("gate-empty-allowed") val config = OrchestrationConfig(retryPolicy = RetryPolicy(maxAttempts = 1, backoffMs = 0)) val graph = singleStageGraph(allowedTools = emptySet()) orchestrator.run(sessionId, graph, config) val events = eventStore.read(sessionId) assertTrue(!executor.executeCalled.get(), "executor must NOT be called when allowedTools is empty") val rejected = events.firstNotNullOfOrNull { it.payload as? ToolExecutionRejectedEvent } assertNotNull(rejected, "ToolExecutionRejectedEvent must be emitted when allowedTools is empty") assertEquals("file_write", rejected!!.toolName) assertTrue( rejected.reason.contains("not permitted"), "rejection reason should mention 'not permitted', got: ${rejected.reason}", ) } }