feat(core:kernel): tool call dispatch loop, tool definitions wiring, sandboxRoot and defaultSystemPromptPath in config

2026-05-18 21:39:50 +04:00
parent 15d1e09c44
commit 5c3a8fda63
5 changed files with 218 additions and 35 deletions
@@ -36,7 +36,7 @@ class DefaultSessionOrchestrator(
        graph: WorkflowGraph,
        config: OrchestrationConfig,
    ): WorkflowResult {
-        System.err.println("[Orchestrator] session=${sessionId.value} workflow=${graph.id} start=${graph.start.value}")
+        log.debug("[Orchestrator] session={} workflow={} start={}", sessionId.value, graph.id, graph.start.value)
        emitWorkflowStarted(sessionId, graph, config)

        val base = ExecutionContext(graph, sessionId, 0, graph.start, config, null, null)
@@ -50,9 +50,9 @@ class DefaultSessionOrchestrator(
    }

    private tailrec suspend fun step(ctx: EnrichedExecutionContext): WorkflowResult {
-        System.err.println(
-            "[Orchestrator] step session=${ctx.sessionId.value} stage=${ctx.currentStageId.value} " +
-            "stageCount=${ctx.stageCount}"
+        log.debug(
+            "[Orchestrator] step session={} stage={} stageCount={}",
+            ctx.sessionId.value, ctx.currentStageId.value, ctx.stageCount,
        )
        if (isCancelled(ctx.sessionId)) return handleCancellation(ctx.sessionId, ctx.currentStageId)

@@ -68,13 +68,15 @@ class DefaultSessionOrchestrator(

        val stageConfig = enriched.graph.stages[enriched.currentStageId]
        val stageArtifacts = stageConfig?.produces
-            ?.let { repositories.artifactRepository.getByIds(enriched.sessionId, it) }
+            ?.let { slots ->
+                repositories.artifactRepository.getByIds(enriched.sessionId, slots.map { it.name }.toSet())
+            }
            ?: emptyMap()

        val decision = resolveTransition(enriched.graph, enriched.sessionId, enriched.currentStageId, stageArtifacts)
-        System.err.println(
-            "[Orchestrator] transition session=${enriched.sessionId.value} " +
-            "stage=${enriched.currentStageId.value} decision=${decision::class.simpleName}"
+        log.debug(
+            "[Orchestrator] transition session={} stage={} decision={}",
+            enriched.sessionId.value, enriched.currentStageId.value, decision::class.simpleName,
        )
        return when (decision) {
            is TransitionDecision.Move -> when (val r = executeMove(enriched, decision)) {
@@ -2,9 +2,12 @@ package com.correx.core.kernel.orchestration

 import com.correx.core.events.execution.RetryPolicy
 import com.correx.core.kernel.execution.ReplayStrategy
+import java.nio.file.Path

 data class OrchestrationConfig(
    val retryPolicy: RetryPolicy = RetryPolicy(maxAttempts = 3),
    val replayStrategy: ReplayStrategy = ReplayStrategy.Full,
    val stageTimeoutMs: Long = 60_000L,
+    val sandboxRoot: Path = Path.of(System.getProperty("java.io.tmpdir"), "correx-sandbox"),
+    val defaultSystemPromptPath: String? = null,
 )
@@ -4,6 +4,8 @@ import com.correx.core.approvals.domain.ApprovalEngine
 import com.correx.core.context.builder.ContextPackBuilder
 import com.correx.core.inference.InferenceRouter
 import com.correx.core.risk.RiskAssessor
+import com.correx.core.tools.contract.ToolExecutor
+import com.correx.core.tools.registry.ToolRegistry
 import com.correx.core.transitions.evaluation.PromptResolver
 import com.correx.core.transitions.resolution.TransitionResolver
 import com.correx.core.validation.pipeline.ValidationPipeline
@@ -16,4 +18,6 @@ data class OrchestratorEngines(
    val approvalEngine: ApprovalEngine,
    val riskAssessor: RiskAssessor,
    val promptResolver: PromptResolver = PromptResolver { "" },
+    val toolExecutor: ToolExecutor? = null,
+    val toolRegistry: ToolRegistry? = null,
 )
@@ -10,8 +10,9 @@ import com.correx.core.events.types.InferenceRequestId
 import com.correx.core.events.types.SessionId
 import com.correx.core.events.types.StageId
 import com.correx.core.inference.GenerationConfig
-import com.correx.core.sessions.Session
 import com.correx.core.inference.InferenceRequest
+import com.correx.core.inference.ResponseFormat
+import com.correx.core.sessions.Session
 import com.correx.core.kernel.execution.ReplayStrategy
 import com.correx.core.kernel.execution.WorkflowResult
 import com.correx.core.kernel.replay.ReplayInferenceProvider
@@ -142,12 +143,14 @@ class ReplayOrchestrator(
        cancellations.getOrPut(sessionId) { AtomicBoolean(false) }.set(true)
    }

+    @Suppress("LongParameterList")
    override suspend fun runInference(
        sessionId: SessionId,
        stageId: StageId,
        contextPack: ContextPack,
        stageConfig: StageConfig,
        timeoutMs: Long,
+        responseFormat: ResponseFormat,
    ): InferenceResult = when (strategy) {
        is ReplayStrategy.SkipInference -> {
            // bypass router entirely — use recorded artifact
@@ -162,6 +165,7 @@ class ReplayOrchestrator(
                    maxTokens = 0,
                    seed = 0L,
                ),
+                responseFormat = responseFormat,
            )
            runCatching { replayProvider.infer(request) }
                .fold(
@@ -175,7 +179,7 @@ class ReplayOrchestrator(
                    },
                )
        }
-        else -> super.runInference(sessionId, stageId, contextPack, stageConfig, timeoutMs)
+        else -> super.runInference(sessionId, stageId, contextPack, stageConfig, timeoutMs, responseFormat)
    }

    override suspend fun mapValidationOutcome(
@@ -28,6 +28,7 @@ import com.correx.core.events.events.WorkflowCompletedEvent
 import com.correx.core.events.events.WorkflowFailedEvent
 import com.correx.core.events.events.WorkflowStartedEvent
 import com.correx.core.artifacts.ArtifactState
+import com.correx.core.artifacts.kind.JsonSchema
 import com.correx.core.events.stores.EventStore
 import com.correx.core.events.types.ApprovalDecisionId
 import com.correx.core.events.types.ApprovalRequestId
@@ -39,10 +40,22 @@ import com.correx.core.events.types.RiskSummaryId
 import com.correx.core.events.types.SessionId
 import com.correx.core.events.types.StageId
 import com.correx.core.events.types.ValidationReportId
+import com.correx.core.approvals.Tier
+import com.correx.core.events.events.ToolInvocationRequestedEvent
+import com.correx.core.events.events.ToolRequest
+import com.correx.core.events.types.ToolInvocationId
+import com.correx.core.inference.FinishReason
 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.ResponseFormat
+import com.correx.core.inference.ToolCallRequest
+import com.correx.core.inference.ToolDefinition
+import com.correx.core.inference.ToolFunction
+import com.correx.core.tools.contract.ToolExecutor
+import com.correx.core.tools.contract.ToolResult
+import com.correx.core.tools.registry.ToolRegistry
 import com.correx.core.kernel.execution.WorkflowResult
 import com.correx.core.risk.RiskAssessor
 import com.correx.core.risk.RiskContext
@@ -58,16 +71,19 @@ import com.correx.core.transitions.resolution.TransitionResolver
 import com.correx.core.validation.model.ValidationContext
 import com.correx.core.validation.pipeline.ValidationOutcome
 import com.correx.core.validation.pipeline.ValidationPipeline
-import org.slf4j.LoggerFactory
 import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.TimeoutCancellationException
 import kotlinx.coroutines.withTimeout
 import kotlinx.datetime.Clock
+import kotlinx.serialization.json.Json
+import org.slf4j.LoggerFactory
 import java.util.*
 import java.util.concurrent.*
 import java.util.concurrent.atomic.*
 import kotlin.coroutines.cancellation.CancellationException

+private const val MAX_TOOL_ROUNDS = 10
+
@SuppressWarnings(
    "ForbiddenComment",
    "UnusedParameter",
@@ -90,6 +106,8 @@ abstract class SessionOrchestrator(
    private val validationPipeline: ValidationPipeline = engines.validationPipeline
    private val riskAssessor: RiskAssessor = engines.riskAssessor
    private val promptResolver: PromptResolver = engines.promptResolver
+    private val toolExecutor: ToolExecutor? = engines.toolExecutor
+    private val toolRegistry: ToolRegistry? = engines.toolRegistry
    private val inferenceRepository: InferenceRepository = repositories.inferenceRepository
    private val artifactRepository = repositories.artifactRepository
    internal val orchestrationRepository: OrchestrationRepository = repositories.orchestrationRepository
@@ -107,6 +125,7 @@ abstract class SessionOrchestrator(

    // --- stage execution ---

+    @Suppress("CyclomaticComplexMethod")
    internal suspend fun executeStage(
        sessionId: SessionId,
        stageId: StageId,
@@ -129,6 +148,29 @@ abstract class SessionOrchestrator(
                    .getOrNull()
            }
            ?.takeIf { it.isNotBlank() }
+            ?.let { text ->
+                listOf(
+                    ContextEntry(
+                        id = ContextEntryId(UUID.randomUUID().toString()),
+                        layer = ContextLayer.L0,
+                        content = text,
+                        sourceType = "systemPrompt",
+                        sourceId = stageId.value,
+                        tokenEstimate = text.length / 4,
+                    ),
+                )
+            } ?: config.defaultSystemPromptPath
+            ?.let { path ->
+                runCatching { promptResolver.resolve(path) }
+                    .onFailure {
+                        log.error(
+                            "[SessionOrchestrator] stage=${stageId.value}: " +
+                                "failed to load default system prompt '$path': ${it.message}",
+                        )
+                    }
+                    .getOrNull()
+            }
+            ?.takeIf { it.isNotBlank() }
            ?.let { text ->
                listOf(
                    ContextEntry(
@@ -166,15 +208,49 @@ abstract class SessionOrchestrator(
                )
            } ?: emptyList()

+        val llmEmittedSlots = stageConfig.produces.filter { it.kind.llmEmitted }
+        require(llmEmittedSlots.size <= 1) {
+            "Stage '${stageId.value}' declares ${llmEmittedSlots.size} LLM-emitted artifact kinds; " +
+                "only 0 or 1 is supported"
+        }
+        val responseFormat = llmEmittedSlots.firstOrNull()
+            ?.let { ResponseFormat.Json(it.kind.deriveJsonSchema()) }
+            ?: ResponseFormat.Text
+
+        val schemaEntries = buildSchemaEntries(responseFormat, stageId)
+
        val contextPack = contextPackBuilder.build(
            id = ContextPackId(UUID.randomUUID().toString()),
            sessionId = sessionId,
            stageId = stageId,
-            entries = systemPrompt + promptEntries,
+            entries = systemPrompt + schemaEntries + promptEntries,
            budget = TokenBudget(limit = stageConfig.tokenBudget),
        )

-        val inferenceResult = runInference(sessionId, stageId, contextPack, stageConfig, config.stageTimeoutMs)
+        var currentContext = contextPack
+        var inferenceResult = runInference(
+            sessionId, stageId, currentContext, stageConfig, config.stageTimeoutMs, responseFormat,
+        )
+        var toolRounds = 0
+        while (
+            inferenceResult is InferenceResult.Success &&
+            inferenceResult.response.finishReason is FinishReason.ToolCall &&
+            toolRounds < MAX_TOOL_ROUNDS
+        ) {
+            val toolEntries = dispatchToolCalls(sessionId, stageId, inferenceResult.response.toolCalls)
+            val allEntries = currentContext.layers.values.flatten() + toolEntries
+            currentContext = contextPackBuilder.build(
+                id = ContextPackId(UUID.randomUUID().toString()),
+                sessionId = sessionId,
+                stageId = stageId,
+                entries = allEntries,
+                budget = TokenBudget(limit = stageConfig.tokenBudget),
+            )
+            inferenceResult = runInference(
+                sessionId, stageId, currentContext, stageConfig, config.stageTimeoutMs, responseFormat,
+            )
+            toolRounds++
+        }
        return when (inferenceResult) {
            is InferenceResult.Cancelled -> StageExecutionResult.Failure("CANCELLED", retryable = false)
            is InferenceResult.Failed -> StageExecutionResult.Failure(inferenceResult.reason, retryable = true)
@@ -189,14 +265,94 @@ abstract class SessionOrchestrator(
        }
    }

+    private suspend fun dispatchToolCalls(
+        sessionId: SessionId,
+        stageId: StageId,
+        toolCalls: List<ToolCallRequest>,
+    ): List<ContextEntry> {
+        val executor = toolExecutor ?: return emptyList()
+        return toolCalls.flatMap { toolCall ->
+            val invocationId = ToolInvocationId(UUID.randomUUID().toString())
+            val parameters = runCatching {
+                val element = Json.parseToJsonElement(toolCall.function.arguments)
+                val jsonObject = element as? kotlinx.serialization.json.JsonObject ?: return@runCatching emptyMap()
+                jsonObject.entries.associate { (k, v) -> k to v.toString().removeSurrounding("\"") as Any }
+            }.getOrElse { emptyMap() }
+            val request = ToolRequest(
+                invocationId = invocationId,
+                sessionId = sessionId,
+                stageId = stageId,
+                toolName = toolCall.function.name,
+                parameters = parameters,
+            )
+            val tier = toolRegistry?.resolve(toolCall.function.name)?.tier ?: Tier.T2
+            emit(
+                sessionId,
+                ToolInvocationRequestedEvent(
+                    invocationId = invocationId,
+                    sessionId = sessionId,
+                    stageId = stageId,
+                    toolName = toolCall.function.name,
+                    tier = tier,
+                    request = request,
+                ),
+            )
+            val result = executor.execute(request)
+            val sourceId = toolCall.id ?: invocationId.value
+            val assistantEntry = ContextEntry(
+                id = ContextEntryId(UUID.randomUUID().toString()),
+                layer = ContextLayer.L2,
+                sourceType = "assistantToolCall",
+                sourceId = sourceId,
+                content = Json.encodeToString(ToolCallRequest.serializer(), toolCall),
+                tokenEstimate = toolCall.function.arguments.length / 4,
+            )
+            val resultContent = when (result) {
+                is ToolResult.Success -> result.output
+                is ToolResult.Failure -> "ERROR: ${result.reason}"
+            }
+            val resultEntry = ContextEntry(
+                id = ContextEntryId(UUID.randomUUID().toString()),
+                layer = ContextLayer.L2,
+                sourceType = "toolResult",
+                sourceId = sourceId,
+                content = resultContent,
+                tokenEstimate = resultContent.length / 4,
+            )
+            listOf(assistantEntry, resultEntry)
+        }
+    }
+
+    private fun buildSchemaEntries(
+        responseFormat: ResponseFormat,
+        stageId: StageId,
+    ): List<ContextEntry> {
+        if (responseFormat !is ResponseFormat.Json) return emptyList()
+        val compactSchema = Json.encodeToString(JsonSchema.serializer(), responseFormat.schema)
+        val instruction = "Respond with a single JSON object matching this schema. " +
+            "Do not include markdown, code fences, or commentary outside the JSON. " +
+            "Schema: $compactSchema"
+        return listOf(
+            ContextEntry(
+                id = ContextEntryId(UUID.randomUUID().toString()),
+                layer = ContextLayer.L0,
+                content = instruction,
+                sourceType = "schemaInstruction",
+                sourceId = stageId.value,
+                tokenEstimate = instruction.length / 4,
+            ),
+        )
+    }
+
    private fun verifyProduces(
        sessionId: SessionId,
        stageId: StageId,
        stageConfig: StageConfig,
    ): StageExecutionResult {
        if (stageConfig.produces.isEmpty()) return StageExecutionResult.Success(emptyList())
-        val present = artifactRepository.getByIds(sessionId, stageConfig.produces).keys
-        val missing = stageConfig.produces - present
+        val producedIds = stageConfig.produces.map { it.name }.toSet()
+        val present = artifactRepository.getByIds(sessionId, producedIds).keys
+        val missing = producedIds - present
        if (missing.isEmpty()) return StageExecutionResult.Success(emptyList())
        val missingIds = missing.joinToString(", ") { it.value }
        log.error(
@@ -221,17 +377,19 @@ abstract class SessionOrchestrator(
        is ValidationOutcome.NeedsApproval -> handleApproval(sessionId, stageId, outcome)
    }

+    @Suppress("LongParameterList")
    internal open suspend fun runInference(
        sessionId: SessionId,
        stageId: StageId,
        contextPack: ContextPack,
        stageConfig: StageConfig,
        timeoutMs: Long,
+        responseFormat: ResponseFormat = ResponseFormat.Text,
    ): InferenceResult {
        val provider = inferenceRouter.route(stageId, stageConfig.requiredCapabilities)
-        System.err.println(
-            "[Orchestrator] inference session=${sessionId.value} stage=${stageId.value} " +
-                    "provider=${provider.id.value} timeoutMs=$timeoutMs",
+        log.debug(
+            "[Orchestrator] inference session={} stage={} provider={} timeoutMs={}",
+            sessionId.value, stageId.value, provider.id.value, timeoutMs,
        )
        val requestId = InferenceRequestId(UUID.randomUUID().toString())
        val request = InferenceRequest(
@@ -240,6 +398,18 @@ abstract class SessionOrchestrator(
            stageId = stageId,
            contextPack = contextPack,
            generationConfig = stageConfig.generationConfig,
+            responseFormat = responseFormat,
+            tools = stageConfig.allowedTools
+                .mapNotNull { toolRegistry?.resolve(it) }
+                .map { tool ->
+                    ToolDefinition(
+                        function = ToolFunction(
+                            name = tool.name,
+                            description = tool.description,
+                            parameters = tool.parametersSchema,
+                        ),
+                    )
+                },
        )

        // TODO: capture rendered prompt; currently contextpack.toString()
@@ -250,9 +420,9 @@ abstract class SessionOrchestrator(

        return try {
            val response = withTimeout(timeoutMs) { provider.infer(request) }
-            System.err.println(
-                "[Orchestrator] inference done session=${sessionId.value} stage=${stageId.value} " +
-                        "tokens=${response.tokensUsed} latencyMs=${response.latencyMs}",
+            log.debug(
+                "[Orchestrator] inference done session={} stage={} tokens={} latencyMs={}",
+                sessionId.value, stageId.value, response.tokensUsed, response.latencyMs,
            )

            if (isCancelled(sessionId)) return InferenceResult.Cancelled
@@ -343,9 +513,9 @@ abstract class SessionOrchestrator(
        terminalStageId: StageId,
        stageCount: Int,
    ): WorkflowResult.Completed {
-        System.err.println(
-            "[Orchestrator] COMPLETED session=${sessionId.value} " +
-                    "terminalStage=${terminalStageId.value} stages=$stageCount",
+        log.debug(
+            "[Orchestrator] COMPLETED session={} terminalStage={} stages={}",
+            sessionId.value, terminalStageId.value, stageCount,
        )
        emit(sessionId, WorkflowCompletedEvent(sessionId, terminalStageId, stageCount))
        cancellations.remove(sessionId)
@@ -358,9 +528,9 @@ abstract class SessionOrchestrator(
        reason: String,
        retryExhausted: Boolean,
    ): WorkflowResult.Failed {
-        System.err.println(
-            "[Orchestrator] FAILED session=${sessionId.value} stage=${stageId.value} " +
-                    "reason=$reason retryExhausted=$retryExhausted",
+        log.warn(
+            "[Orchestrator] FAILED session={} stage={} reason={} retryExhausted={}",
+            sessionId.value, stageId.value, reason, retryExhausted,
        )
        emit(sessionId, WorkflowFailedEvent(sessionId, stageId, reason, retryExhausted))
        cancellations.remove(sessionId)
@@ -371,7 +541,7 @@ abstract class SessionOrchestrator(
        sessionId: SessionId,
        stageId: StageId,
    ): WorkflowResult.Cancelled {
-        System.err.println("[Orchestrator] CANCELLED session=${sessionId.value} stage=${stageId.value}")
+        log.warn("[Orchestrator] CANCELLED session={} stage={}", sessionId.value, stageId.value)
        emit(sessionId, WorkflowFailedEvent(sessionId, stageId, "CANCELLED", retryExhausted = false))
        cancellations.remove(sessionId)
        return WorkflowResult.Cancelled(sessionId)
@@ -407,7 +577,7 @@ abstract class SessionOrchestrator(
        stageId: StageId,
        outcome: ValidationOutcome.NeedsApproval,
    ): StageExecutionResult {
-        System.err.println("[Orchestrator] approval required session=${sessionId.value} stage=${stageId.value}")
+        log.debug("[Orchestrator] approval required session={} stage={}", sessionId.value, stageId.value)
        emit(sessionId, OrchestrationPausedEvent(sessionId, stageId, "APPROVAL_PENDING"))
        val domainRequest = buildApprovalRequest(sessionId, stageId, outcome)

@@ -428,14 +598,14 @@ abstract class SessionOrchestrator(
        pendingApprovals[domainRequest.id] = deferred

        return try {
-            System.err.println(
-                "[Orchestrator] awaiting approval session=${sessionId.value} " +
-                        "requestId=${domainRequest.id.value}",
+            log.debug(
+                "[Orchestrator] awaiting approval session={} requestId={}",
+                sessionId.value, domainRequest.id.value,
            )
            val decision = deferred.await()
-            System.err.println(
-                "[Orchestrator] approval resolved session=${sessionId.value} " +
-                        "approved=${decision.isApproved}",
+            log.debug(
+                "[Orchestrator] approval resolved session={} approved={}",
+                sessionId.value, decision.isApproved,
            )
            emitDecisionResolved(sessionId, domainRequest, decision)
            if (decision.isApproved) {