diff --git a/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/DefaultSessionOrchestrator.kt b/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/DefaultSessionOrchestrator.kt index 82cd0cc2..32bd47bf 100644 --- a/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/DefaultSessionOrchestrator.kt +++ b/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/DefaultSessionOrchestrator.kt @@ -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)) { diff --git a/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/OrchestrationConfig.kt b/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/OrchestrationConfig.kt index 1c12dcd1..2750e26e 100644 --- a/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/OrchestrationConfig.kt +++ b/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/OrchestrationConfig.kt @@ -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, ) \ No newline at end of file diff --git a/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/OrchestratorEngines.kt b/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/OrchestratorEngines.kt index 85e0bd93..9e7441b2 100644 --- a/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/OrchestratorEngines.kt +++ b/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/OrchestratorEngines.kt @@ -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, ) diff --git a/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/ReplayOrchestrator.kt b/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/ReplayOrchestrator.kt index 7e8588a0..d41fda54 100644 --- a/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/ReplayOrchestrator.kt +++ b/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/ReplayOrchestrator.kt @@ -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( diff --git a/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/SessionOrchestrator.kt b/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/SessionOrchestrator.kt index 5d9254b0..b132bd63 100644 --- a/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/SessionOrchestrator.kt +++ b/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/SessionOrchestrator.kt @@ -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, + ): List { + 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 { + 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) {