correx/testing/integration/src/test/kotlin/ToolCallGateTest.kt

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<ToolCapability> = 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<ToolCapability> = 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<Tool> = 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<InferenceRequest>()
        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<CapabilityScore> = setOf(CapabilityScore(ModelCapability.General, 1.0))
    }

    private fun ruleReturning(action: RiskAction): ToolCallRule = object : ToolCallRule {
        override fun appliesTo(capabilities: Set<ToolCapability>) = 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<DefaultSessionOrchestrator, InMemoryEventStore, ToolCallingProvider> {
        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<ModelCapability>): InferenceProvider =
                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 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<String> = 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}",
        )
    }
}