feat(workflow): typed produces slots — StageConfig uses TypedArtifactSlot, TOML accepts {name, kind} objects

core/transitions/src/main/kotlin/com/correx/core/transitions/graph/StageConfig.kt

@@ -1,11 +1,10 @@
 package com.correx.core.transitions.graph
 
+import com.correx.core.artifacts.kind.TypedArtifactSlot
 import com.correx.core.events.types.ArtifactId
 import com.correx.core.inference.GenerationConfig
 import com.correx.core.inference.ModelCapability
-import kotlinx.serialization.Serializable
 
-@Serializable
 data class StageConfig(
     val requiredCapabilities: Set<ModelCapability> = emptySet(),
     val tokenBudget: Int = 4096,
@@ -16,7 +15,7 @@ data class StageConfig(
     ),
     val allowedTools: Set<String> = emptySet(),
     val maxRetries: Int = 3,
-    val produces: Set<ArtifactId> = emptySet(),
+    val produces: List<TypedArtifactSlot> = emptyList(),
     val needs: Set<ArtifactId> = emptySet(),
     val metadata: Map<String, String> = emptyMap(),
 )

infrastructure/workflow/build.gradle

@@ -7,6 +7,7 @@ dependencies {
     implementation(project(":core:transitions"))
     implementation(project(":core:inference"))
     implementation(project(":core:events"))
+    implementation(project(":core:artifacts"))
     implementation("com.fasterxml.jackson.dataformat:jackson-dataformat-toml:2.17.0")
     implementation("com.fasterxml.jackson.module:jackson-module-kotlin:2.17.0")
     testImplementation "org.junit.jupiter:junit-jupiter"

infrastructure/workflow/src/main/kotlin/com/correx/infrastructure/workflow/TomlWorkflowLoader.kt

@@ -1,5 +1,8 @@
 package com.correx.infrastructure.workflow
 
+import com.correx.core.artifacts.kind.ArtifactKindRegistry
+import com.correx.core.artifacts.kind.DefaultArtifactKindRegistry
+import com.correx.core.artifacts.kind.TypedArtifactSlot
 import com.correx.core.events.types.ArtifactId
 import com.correx.core.events.types.StageId
 import com.correx.core.events.types.TransitionId
@@ -21,11 +24,16 @@ private data class WorkflowFile(
     val transitions: List<TransitionSection> = emptyList(),
 )
 
+private data class ProducesEntry(
+    val name: String = "",
+    val kind: String = "",
+)
+
 private data class StageSection(
     val id: String = "",
     val prompt: String? = null,
-    val systemPrompt: String? = null,
+    @JsonProperty("system_prompt") val systemPrompt: String? = null,
-    val produces: List<String> = emptyList(),
+    val produces: List<ProducesEntry> = emptyList(),
     val needs: List<String> = emptyList(),
     @JsonProperty("allowed_tools") val allowedTools: List<String> = emptyList(),
     @JsonProperty("token_budget") val tokenBudget: Int = 4096,
@@ -50,7 +58,9 @@ private val mapper = TomlMapper.builder()
     .disable(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES)
     .build()
 
-class TomlWorkflowLoader : WorkflowLoader {
+class TomlWorkflowLoader(
+    private val registry: ArtifactKindRegistry = DefaultArtifactKindRegistry(),
+) : WorkflowLoader {
     override fun load(path: Path): WorkflowGraph {
         val raw = path.readText()
         val file = runCatching { mapper.readValue<WorkflowFile>(raw) }
@@ -62,14 +72,18 @@ class TomlWorkflowLoader : WorkflowLoader {
         val startId = StageId(start)
         val stageMap = stages.associate { s ->
             StageId(s.id) to StageConfig(
-                produces = s.produces.map { ArtifactId(it) }.toSet(),
+                produces = s.produces.map { entry ->
+                    val resolvedKind = registry.get(entry.kind)
+                        ?: error("Unknown artifact kind '${entry.kind}' in stage '${s.id}'")
+                    TypedArtifactSlot(name = ArtifactId(entry.name), kind = resolvedKind)
+                },
                 needs = s.needs.map { ArtifactId(it) }.toSet(),
                 allowedTools = s.allowedTools.toSet(),
                 tokenBudget = s.tokenBudget,
                 maxRetries = s.maxRetries,
                 metadata = buildMap {
                     s.prompt?.let { put("prompt", it) }
-                        ?: s.systemPrompt?.let { put("systemPrompt", it) }
+                    s.systemPrompt?.let { put("systemPrompt", it) }
                 },
             )
         }
@@ -77,7 +91,7 @@ class TomlWorkflowLoader : WorkflowLoader {
         val declaredIds = stageMap.keys.map { it.value }.toSet()
         validate(start, declaredIds, transitions)
 
-        val allProduced = stageMap.values.flatMap { it.produces }.toSet()
+        val allProduced = stageMap.values.flatMap { it.produces }.map { it.name }.toSet()
         stageMap.forEach { (stageId, config) ->
             config.needs.forEach { needed ->
                 if (needed !in allProduced) {

infrastructure/workflow/src/test/kotlin/com/correx/infrastructure/workflow/TomlWorkflowLoaderTest.kt

@@ -17,7 +17,7 @@ class TomlWorkflowLoaderTest {
         [[stages]]
         id = "collect"
         prompt = "prompts/collect.md"
-        produces = ["system_stats"]
+        produces = [{ name = "system_stats", kind = "file_written" }]
         allowed_tools = ["ShellTool"]
         token_budget = 2048
         max_retries = 2
@@ -26,7 +26,7 @@ class TomlWorkflowLoaderTest {
         id = "report"
         prompt = "prompts/report.md"
         needs = ["system_stats"]
-        produces = ["report"]
+        produces = [{ name = "report", kind = "file_written" }]
         token_budget = 2048
 
         [[transitions]]
@@ -53,10 +53,10 @@ class TomlWorkflowLoaderTest {
         assertEquals(2, graph.transitions.size)
 
         val collectStage = graph.stages[graph.start]!!
-        assertEquals(setOf("system_stats"), collectStage.produces.map { it.value }.toSet())
+        assertEquals(setOf("system_stats"), collectStage.produces.map { it.name.value }.toSet())
         assertEquals("prompts/collect.md", collectStage.metadata["prompt"])
 
-        val reportStage = graph.stages.values.first { it.produces.any { a -> a.value == "report" } }
+        val reportStage = graph.stages.values.first { it.produces.any { a -> a.name.value == "report" } }
         assertEquals(setOf("system_stats"), reportStage.needs.map { it.value }.toSet())
     }
 
@@ -90,4 +90,11 @@ class TomlWorkflowLoaderTest {
         val path = Files.createTempFile("workflow", ".toml").also { it.writeText(bad) }
         assertThrows<WorkflowValidationException> { loader.load(path) }
     }
+
+    @Test
+    fun `unknown artifact kind throws error`() {
+        val bad = validToml.replace("kind = \"file_written\"", "kind = \"unknown_kind\"")
+        val path = Files.createTempFile("workflow", ".toml").also { it.writeText(bad) }
+        assertThrows<IllegalStateException> { loader.load(path) }
+    }
 }