From d12d64e4fe6237b8bff403afb7f2cf4220ae01c3 Mon Sep 17 00:00:00 2001 From: kami Date: Wed, 24 Jun 2026 16:16:03 +0000 Subject: [PATCH] feat(workflow): reject execution plans that reference unknown tools MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A freestyle stage's tools are LLM-authored in the execution_plan. At run time an unresolvable tool name is silently dropped (resolve -> null -> mapNotNull), so the stage runs toolless — e.g. a misspelled task_update means the implementation stage quietly loses task tracking, invisible without a live run. ExecutionPlanCompiler now takes the registered-tool universe and rejects a plan that names an unknown tool, with a message identifying the tool and stage. FreestyleDriver already turns a compile failure into a rejection the architect retries. The param defaults to empty (validation skipped) so existing callers are unaffected; Main feeds it toolRegistry.all(). Co-Authored-By: Claude Opus 4.8 --- .../kotlin/com/correx/apps/server/Main.kt | 2 +- .../workflow/ExecutionPlanCompiler.kt | 23 +++++++++ .../workflow/ExecutionPlanCompilerTest.kt | 49 +++++++++++++++++++ 3 files changed, 73 insertions(+), 1 deletion(-) diff --git a/apps/server/src/main/kotlin/com/correx/apps/server/Main.kt b/apps/server/src/main/kotlin/com/correx/apps/server/Main.kt index 428d0482..849a9866 100644 --- a/apps/server/src/main/kotlin/com/correx/apps/server/Main.kt +++ b/apps/server/src/main/kotlin/com/correx/apps/server/Main.kt @@ -457,7 +457,7 @@ fun main() { val freestyleDriver = FreestyleDriver( eventStore = eventStore, - compiler = ExecutionPlanCompiler(artifactKindRegistry), + compiler = ExecutionPlanCompiler(artifactKindRegistry, toolRegistry.all().map { it.name }.toSet()), planContent = { sid -> orchestrator.validatedArtifactContent(sid, ArtifactId("execution_plan")) }, config = defaultOrchestrationConfig, runPhase2 = { sid, graph, cfg -> orchestrator.run(sid, graph, cfg) }, diff --git a/infrastructure/workflow/src/main/kotlin/com/correx/infrastructure/workflow/ExecutionPlanCompiler.kt b/infrastructure/workflow/src/main/kotlin/com/correx/infrastructure/workflow/ExecutionPlanCompiler.kt index 7d0f33f9..0e1177cd 100644 --- a/infrastructure/workflow/src/main/kotlin/com/correx/infrastructure/workflow/ExecutionPlanCompiler.kt +++ b/infrastructure/workflow/src/main/kotlin/com/correx/infrastructure/workflow/ExecutionPlanCompiler.kt @@ -17,6 +17,12 @@ private const val TERMINAL = "done" class ExecutionPlanCompiler( private val registry: ArtifactKindRegistry, + // Names of every registered tool. A stage that references a tool the runtime can't resolve + // is dropped silently at run time (mapNotNull), so the stage runs toolless — invisible + // without a live run. Validating here turns that into a compile-time rejection (which the + // freestyle driver feeds back for a retry). Empty = the caller didn't supply the tool + // universe, so validation is skipped (preserves the bare `ExecutionPlanCompiler(registry)`). + private val knownTools: Set = emptySet(), ) { private val mapper = JsonMapper.builder() .addModule(kotlinModule()) @@ -27,6 +33,7 @@ class ExecutionPlanCompiler( val plan = runCatching { mapper.readValue(planJson) } .getOrElse { throw WorkflowValidationException("Failed to parse execution_plan JSON: ${it.message}") } if (plan.stages.isEmpty()) throw WorkflowValidationException("execution_plan has no stages") + validateTools(plan) val stageMap = plan.stages.associate { s -> // `produces` is the unique slot name referenced by needs/edges; `kind` selects the @@ -81,6 +88,22 @@ class ExecutionPlanCompiler( return WorkflowGraph(id = workflowId, stages = stageMap, transitions = transitions, start = start) } + /** + * Every tool a stage names must be a registered tool, else the runtime silently drops it + * (resolve → null → mapNotNull) and the stage runs without it. Skipped when [knownTools] is + * empty (the caller didn't supply the tool universe). + */ + private fun validateTools(plan: ExecutionPlanModel) { + if (knownTools.isEmpty()) return + val offenders = plan.stages.flatMap { s -> s.tools.filter { it !in knownTools }.map { s.id to it } } + if (offenders.isEmpty()) return + val detail = offenders.joinToString(", ") { (stage, tool) -> "'$tool' in stage '$stage'" } + throw WorkflowValidationException( + "execution_plan references unknown tool(s): $detail — valid tools: " + + knownTools.sorted().joinToString(", "), + ) + } + /** * A field-equals edge can only ever fire if the producing stage's kind schema declares * that field — the kind's schema is also the LLM response format, so an undeclared field diff --git a/infrastructure/workflow/src/test/kotlin/com/correx/infrastructure/workflow/ExecutionPlanCompilerTest.kt b/infrastructure/workflow/src/test/kotlin/com/correx/infrastructure/workflow/ExecutionPlanCompilerTest.kt index 6d2b24da..c0f9ed18 100644 --- a/infrastructure/workflow/src/test/kotlin/com/correx/infrastructure/workflow/ExecutionPlanCompilerTest.kt +++ b/infrastructure/workflow/src/test/kotlin/com/correx/infrastructure/workflow/ExecutionPlanCompilerTest.kt @@ -273,4 +273,53 @@ class ExecutionPlanCompilerTest { fun `malformed JSON throws WorkflowValidationException`() { assertThrows { compiler.compile("not-json", "bad-workflow") } } + + // A compiler that knows the tool universe, so plan tool names are validated. + private val toolAware = ExecutionPlanCompiler( + registry, + knownTools = setOf("ShellTool", "file_write", "task_context", "task_update"), + ) + + @Test + fun `unknown tool name throws WorkflowValidationException naming the offender`() { + // The runtime would silently drop a misspelled tool; here it fails to compile instead. + val bad = validPlan.replace("\"tools\": [\"ShellTool\"]", "\"tools\": [\"task_updates\"]") + val ex = assertThrows { toolAware.compile(bad, "bad-workflow") } + assertEquals(true, ex.message?.contains("task_updates")) + assertEquals(true, ex.message?.contains("analyse")) + } + + @Test + fun `plan threading registered task tools compiles`() { + val plan = """ + { + "goal": "implement and review a tracked task", + "stages": [ + { + "id": "implement", "prompt": "claim auth-1, build, submit", "produces": "patch", + "needs": [], "tools": ["file_write", "task_context", "task_update"] + }, + { + "id": "review", "prompt": "review and complete auth-1", "produces": "patch", + "needs": ["patch"], "tools": ["task_context", "task_update"] + } + ], + "edges": [ + { "from": "implement", "to": "review", "condition": { "type": "always_true" } }, + { "from": "review", "to": "done", "condition": { "type": "always_true" } } + ] + } + """.trimIndent() + val graph = toolAware.compile(plan, "tracked-workflow") + val implement = graph.stages.values.first { it.metadata["promptInline"] == "claim auth-1, build, submit" } + assertTrue(implement.allowedTools.containsAll(setOf("task_context", "task_update"))) + } + + @Test + fun `tool validation is skipped when the tool universe is unknown`() { + // The default compiler has no knownTools, so an arbitrary tool name still compiles. + val plan = validPlan.replace("\"tools\": [\"ShellTool\"]", "\"tools\": [\"anything_goes\"]") + val graph = compiler.compile(plan, "unvalidated-workflow") + assertEquals(2, graph.stages.size) + } }