feat(kernel): inject needs-artifact content into stage context

Stages declaring needs=[...] now receive each needed artifact's validated JSON (from artifactContentCache) as an L1/USER neededArtifact context entry, so a one-shot subagent sees its inputs. Also emit llm-emitted artifact lifecycle events (Created/Validating/Validated) on successful validation via emitLlmArtifacts, closing a latent gap where verifyProduces would fail for stages producing only an llm-emitted artifact. Events are emitted only after validation passes, preserving invariant #7.
2026-06-08 01:12:38 +04:00
parent 638f57709d
commit 63e5bcea93
2 changed files with 271 additions and 1 deletions
@@ -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<ContextEntry> {
+        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,
@@ -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<ContextEntry>()
+
+    private val capturingBuilder = object : ContextPackBuilder {
+        override fun build(
+            id: ContextPackId,
+            sessionId: SessionId,
+            stageId: StageId,
+            entries: List<ContextEntry>,
+            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<InferenceState> {
+            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<ModelCapability>,
+            ) = 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")
+    }
+}