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 dbd42090..5b2afd9d 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 @@ -306,8 +306,9 @@ abstract class SessionOrchestrator( ), ) val repoMapEntries = buildRepoMapEntries(sessionId) + val needsEntries = buildNeedsArtifactEntries(sessionId, stageConfig) var accumulatedEntries = - systemPrompt + journalEntries + repoMapEntries + schemaEntries + promptEntries + steeringEntries + systemPrompt + journalEntries + repoMapEntries + needsEntries + schemaEntries + promptEntries + steeringEntries val contextPack = contextPackBuilder.build( id = ContextPackId(UUID.randomUUID().toString()), sessionId = sessionId, @@ -384,6 +385,7 @@ abstract class SessionOrchestrator( when (val outcome = mapValidationOutcome(sessionId, stageId, validationCtx)) { is StageExecutionResult.Success -> { emitToolArtifacts(sessionId, stageId, stageConfig) + emitLlmArtifacts(sessionId, stageId, stageConfig) verifyProduces(sessionId, stageId, stageConfig) } @@ -843,6 +845,26 @@ abstract class SessionOrchestrator( ) } + private suspend fun buildNeedsArtifactEntries( + sessionId: SessionId, + stageConfig: StageConfig, + ): List { + return stageConfig.needs.mapNotNull { needed -> + val cached = artifactContentCache["${sessionId.value}:${needed.value}"] + ?.takeIf { it.isNotBlank() } ?: return@mapNotNull null + val content = "## Input artifact: ${needed.value}\n$cached" + ContextEntry( + id = ContextEntryId(UUID.randomUUID().toString()), + layer = ContextLayer.L1, + content = content, + sourceType = "neededArtifact", + sourceId = needed.value, + tokenEstimate = estimateTokens(content), + role = EntryRole.USER, + ) + } + } + private suspend fun emitToolArtifacts( sessionId: SessionId, stageId: StageId, @@ -857,6 +879,20 @@ abstract class SessionOrchestrator( } } + private suspend fun emitLlmArtifacts( + sessionId: SessionId, + stageId: StageId, + stageConfig: StageConfig, + ) { + stageConfig.produces + .filter { it.kind.llmEmitted } + .forEach { slot -> + emit(sessionId, ArtifactCreatedEvent(slot.name, sessionId, stageId, schemaVersion = 1)) + emit(sessionId, ArtifactValidatingEvent(slot.name, sessionId, stageId)) + emit(sessionId, ArtifactValidatedEvent(slot.name, sessionId, stageId)) + } + } + private fun verifyProduces( sessionId: SessionId, stageId: StageId, diff --git a/testing/integration/src/test/kotlin/NeedsArtifactInjectionTest.kt b/testing/integration/src/test/kotlin/NeedsArtifactInjectionTest.kt new file mode 100644 index 00000000..d81b106d --- /dev/null +++ b/testing/integration/src/test/kotlin/NeedsArtifactInjectionTest.kt @@ -0,0 +1,234 @@ +import com.correx.core.approvals.ApprovalProjector +import com.correx.core.approvals.DefaultApprovalReducer +import com.correx.core.approvals.DefaultApprovalRepository +import com.correx.core.approvals.domain.DefaultApprovalEngine +import com.correx.core.artifacts.DefaultArtifactReducer +import com.correx.core.artifacts.kind.ConfigArtifactKind +import com.correx.core.artifacts.kind.JsonSchema +import com.correx.core.artifacts.kind.TypedArtifactSlot +import com.correx.core.context.builder.ContextPackBuilder +import com.correx.core.context.model.CompressionMetadata +import com.correx.core.context.model.ContextEntry +import com.correx.core.context.model.ContextPack +import com.correx.core.context.model.TokenBudget +import com.correx.core.events.execution.RetryPolicy +import com.correx.core.events.types.ArtifactId +import com.correx.core.events.types.ContextPackId +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.InferenceRepository +import com.correx.core.inference.InferenceState +import com.correx.core.inference.ModelCapability +import com.correx.core.journal.DecisionJournalProjector +import com.correx.core.journal.DefaultDecisionJournalReducer +import com.correx.core.journal.DefaultDecisionJournalRepository +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.DefaultSessionRepository +import com.correx.core.sessions.projections.replay.DefaultEventReplayer +import com.correx.core.sessions.projections.replay.EventReplayer +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.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.InferenceFixtures +import com.correx.testing.fixtures.WorkflowFixtures +import com.correx.testing.fixtures.cyclePolicyMissingValidator +import com.correx.testing.fixtures.inference.MockInferenceProvider +import com.correx.testing.fixtures.transitions.TransitionFixtures +import com.correx.testing.kernel.MockSessionEventReplayer +import kotlinx.coroutines.runBlocking +import org.junit.jupiter.api.Assertions.assertFalse +import org.junit.jupiter.api.Assertions.assertTrue +import org.junit.jupiter.api.Test + +/** + * Integration tests verifying that a stage's declared [StageConfig.needs] artifact ids are + * injected into its context pack as "neededArtifact" entries, sourced from artifactContentCache. + */ +class NeedsArtifactInjectionTest { + + /** + * A capturing [ContextPackBuilder] that records every entry list it sees so tests + * can assert on what was assembled for each stage. + */ + private val capturedEntries = mutableListOf() + + private val capturingBuilder = object : ContextPackBuilder { + override fun build( + id: ContextPackId, + sessionId: SessionId, + stageId: StageId, + entries: List, + budget: TokenBudget, + ): ContextPack { + capturedEntries.addAll(entries) + return ContextPack( + id = id, + sessionId = sessionId, + stageId = stageId, + layers = emptyMap(), + budgetUsed = 0, + budgetLimit = budget.limit, + ) + } + } + + private val eventStore = InMemoryEventStore() + private val sessionReplayer = MockSessionEventReplayer() + private val sessionRepository = DefaultSessionRepository(sessionReplayer) + private val orchestrationRepository = OrchestrationRepository( + DefaultEventReplayer(eventStore, OrchestrationProjector(DefaultOrchestrationReducer())), + ) + private val inferenceRepository = InferenceRepository( + object : EventReplayer { + override fun rebuild(sessionId: SessionId) = InferenceState() + }, + ) + private val approvalRepository = DefaultApprovalRepository( + DefaultEventReplayer(eventStore, ApprovalProjector(DefaultApprovalReducer())), + ) + private val decisionJournalRepository = DefaultDecisionJournalRepository( + DefaultEventReplayer(eventStore, DecisionJournalProjector(DefaultDecisionJournalReducer())), + ) + private val artifactStore = NoopArtifactStore() + private val retryCoordinator = DefaultRetryCoordinator(eventStore) + + private val repositories = OrchestratorRepositories( + eventStore = eventStore, + inferenceRepository = inferenceRepository, + orchestrationRepository = orchestrationRepository, + sessionRepository = sessionRepository, + artifactRepository = LiveArtifactRepository(eventStore, DefaultArtifactReducer()), + approvalRepository = approvalRepository, + ) + + private fun buildOrchestrator(inferenceRouter: com.correx.core.inference.InferenceRouter) = + DefaultSessionOrchestrator( + repositories = repositories, + engines = OrchestratorEngines( + transitionResolver = TransitionFixtures.simpleResolver(), + contextPackBuilder = capturingBuilder, + inferenceRouter = inferenceRouter, + validationPipeline = ValidationPipeline(validators = listOf(cyclePolicyMissingValidator())), + approvalEngine = DefaultApprovalEngine(), + riskAssessor = DefaultRiskAssessor(), + ), + retryCoordinator = retryCoordinator, + artifactStore = artifactStore, + decisionJournalRepository = decisionJournalRepository, + ) + + private val defaultConfig = OrchestrationConfig( + retryPolicy = RetryPolicy(maxAttempts = 1, backoffMs = 0), + ) + + /** + * Stage A produces an llm-emitted artifact ("analysis") so the orchestrator stores the + * inference response text in artifactContentCache. Stage B declares needs=[analysis] and + * should receive a "neededArtifact" context entry with that content. + */ + @Test + fun `stage with needs receives artifact content in context entries`(): Unit = runBlocking { + val sessionId = SessionId("needs-test-1") + val stageA = StageId("A") + val stageB = StageId("B") + val artifactId = ArtifactId("analysis") + val artifactJson = """{"summary":"important findings"}""" + + val analysisKind = ConfigArtifactKind( + id = "analysis", + schema = JsonSchema(type = "object", properties = emptyMap()), + llmEmitted = true, + ) + + val stageC = StageId("C") + val graph = WorkflowGraph( + id = "needs-injection-test", + stages = mapOf( + stageA to StageConfig( + produces = listOf(TypedArtifactSlot(name = artifactId, kind = analysisKind)), + ), + stageB to StageConfig( + needs = setOf(artifactId), + ), + stageC to StageConfig(), + ), + transitions = setOf( + TransitionEdge(id = TransitionId("t1"), from = stageA, to = stageB, condition = { true }), + TransitionEdge(id = TransitionId("t2"), from = stageB, to = stageC, condition = { true }), + ), + start = stageA, + ) + + // Stage A returns the artifact JSON so it gets cached; stages B and C return generic text. + val router = object : com.correx.core.inference.InferenceRouter { + private var callCount = 0 + override suspend fun route( + stageId: StageId, + requiredCapabilities: Set, + ) = MockInferenceProvider( + fixedResponse = if (callCount++ == 0) artifactJson else "stage response", + ) + } + + buildOrchestrator(router).run(sessionId, graph, defaultConfig) + + val needsEntries = capturedEntries.filter { it.sourceType == "neededArtifact" } + assertTrue(needsEntries.isNotEmpty(), "Expected at least one neededArtifact entry") + val analysisEntry = needsEntries.find { it.sourceId == "analysis" } + assertFalse(analysisEntry == null, "Expected a neededArtifact entry with sourceId=analysis") + assertTrue( + analysisEntry!!.content.contains(artifactJson), + "Entry content should include artifact JSON, got: ${analysisEntry.content}", + ) + } + + @Test + fun `stage with empty needs adds no neededArtifact entries`(): Unit = runBlocking { + val sessionId = SessionId("needs-test-2") + val graph = WorkflowFixtures.simpleGraph() + + buildOrchestrator(InferenceFixtures.fixedRouter()).run(sessionId, graph, defaultConfig) + + val needsEntries = capturedEntries.filter { it.sourceType == "neededArtifact" } + assertTrue(needsEntries.isEmpty(), "No neededArtifact entries expected for stages with empty needs") + } + + @Test + fun `missing cache content for declared need is skipped without crash`(): Unit = runBlocking { + val sessionId = SessionId("needs-test-3") + val stageA = StageId("A") + val stageB = StageId("B") + val missingArtifactId = ArtifactId("missing_artifact") + + val graph = WorkflowGraph( + id = "missing-needs-test", + stages = mapOf( + stageA to StageConfig(), + stageB to StageConfig(needs = setOf(missingArtifactId)), + ), + transitions = setOf( + TransitionEdge(id = TransitionId("t1"), from = stageA, to = stageB, condition = { true }), + ), + start = stageA, + ) + + // No llm-emitted artifact in stageA → cache never populated for missingArtifactId + buildOrchestrator(InferenceFixtures.fixedRouter()).run(sessionId, graph, defaultConfig) + + val needsEntries = capturedEntries.filter { it.sourceType == "neededArtifact" } + assertTrue(needsEntries.isEmpty(), "Missing cache content should produce no neededArtifact entries") + } +}