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 ad3144fa..738e9caf 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 @@ -116,6 +116,7 @@ import com.correx.core.sessions.Session import com.correx.core.tools.compression.ToolOutputContext import com.correx.core.tools.contract.FileAffectingTool import com.correx.core.tools.contract.Tool +import com.correx.core.tools.contract.ToolCapability import com.correx.core.tools.contract.ToolExecutor import com.correx.core.tools.contract.ToolResult import com.correx.core.tools.registry.ToolRegistry @@ -153,6 +154,7 @@ import kotlin.coroutines.cancellation.CancellationException private const val MAX_TOOL_ROUNDS = 10 private const val STAGE_COMPLETE_TOOL = "stage_complete" +private const val READ_BEFORE_WRITE_CODE = "READ_BEFORE_WRITE" private const val OUTPUT_SUMMARY_LIMIT = 500 private const val REPO_MAP_INJECT_TOP_K = 30 @@ -444,6 +446,27 @@ abstract class SessionOrchestrator( fun owesFileWrite(): Boolean = fileWrittenSlots.isNotEmpty() && fileWrittenSlots.any { artifactContentCache["${sessionId.value}:${it.name.value}"].isNullOrBlank() } + // Gate: if the stage declares require_task_decompose, block stage_complete until at least one + // task_decompose (or task_create) call has completed successfully in this stage. Checked via + // ToolInvocationRequestedEvent + ToolExecutionCompletedEvent in the session event log. + // ponytail: scans events on every stage_complete attempt; negligible cost (in-memory list, rare call). + fun owesTaskDecompose(): Boolean { + if (stageConfig.metadata["requireTaskDecompose"] != "true") return false + val events = eventStore.read(sessionId) + val requestedIds = events + .mapNotNull { it.payload as? ToolInvocationRequestedEvent } + .filter { it.stageId == stageId && it.toolName in setOf("task_decompose", "task_create") } + .map { it.invocationId } + .toSet() + if (requestedIds.isEmpty()) return true + val completedIds = events + .mapNotNull { it.payload as? ToolExecutionCompletedEvent } + .filter { it.toolName in setOf("task_decompose", "task_create") && it.invocationId in requestedIds } + .map { it.invocationId } + .toSet() + return completedIds.isEmpty() + } + // Append a corrective tool-result and re-run inference (bounded by MAX_TOOL_ROUNDS). // Returns the new result rather than mutating inferenceResult, to preserve smart casts. suspend fun pushBack(nudge: String): InferenceResult { @@ -504,6 +527,12 @@ abstract class SessionOrchestrator( "a single JSON message (no tool call) matching the provided schema.", ) owesFileWrite() -> inferenceResult = pushBack(writeNudge) + owesTaskDecompose() -> inferenceResult = pushBack( + "ERROR: stage_complete is not allowed yet — this stage requires a successful " + + "task_decompose (or task_create) call before it can complete. Call task_decompose " + + "with a parent epic and DEPENDS_ON-linked children representing the full task graph. " + + "Do not call stage_complete until task_decompose has returned successfully.", + ) else -> break } continue @@ -636,10 +665,16 @@ abstract class SessionOrchestrator( capabilities = tool?.requiredCapabilities ?: emptySet(), ), ) - if (toolCall.function.name != STAGE_COMPLETE_TOOL && + val deniedByStage = toolCall.function.name != STAGE_COMPLETE_TOOL && toolCall.function.name !in stageConfig.allowedTools - ) { - val denyReason = "tool '${toolCall.function.name}' is not permitted in stage '${stageId.value}'" + val deniedByReadOnly = !deniedByStage && isReadOnlyMode(sessionId) && + (tool?.requiredCapabilities?.contains(ToolCapability.FILE_WRITE) == true) + if (deniedByStage || deniedByReadOnly) { + val denyReason = if (deniedByReadOnly) { + "read-before-write: '${toolCall.function.name}' requires reading the file first — use file_read" + } else { + "tool '${toolCall.function.name}' is not permitted in stage '${stageId.value}'" + } emit( sessionId, ToolExecutionRejectedEvent( @@ -996,6 +1031,37 @@ abstract class SessionOrchestrator( .fold(projection.initial()) { acc, event -> projection.apply(acc, event) } } + /** + * Returns true when the agent must be offered only read-only tools on the next inference turn. + * Active from the moment a READ_BEFORE_WRITE block lands until a FILE_READ completion follows it. + * Derived purely from the event log — no separate flag stored. + */ + private fun isReadOnlyMode(sessionId: SessionId): Boolean { + var blocked = false + val pendingReadIds = mutableSetOf() + for (event in eventStore.read(sessionId)) { + when (val p = event.payload) { + is ToolInvocationRequestedEvent -> + if (p.sessionId == sessionId && ToolCapability.FILE_READ in p.capabilities) + pendingReadIds += p.invocationId.value + is ToolCallAssessedEvent -> + if (p.sessionId == sessionId && p.disposition == RiskAction.BLOCK && + p.issues.any { it.code == READ_BEFORE_WRITE_CODE } + ) { + blocked = true + pendingReadIds.clear() + } + is ToolExecutionCompletedEvent -> + if (p.sessionId == sessionId && p.invocationId.value in pendingReadIds) { + blocked = false + pendingReadIds.clear() + } + else -> {} + } + } + return blocked + } + private suspend fun buildSchemaEntries( responseFormat: ResponseFormat, stageId: StageId, @@ -1680,6 +1746,10 @@ abstract class SessionOrchestrator( responseFormat = responseFormat, tools = stageConfig.allowedTools .mapNotNull { effectives.registry?.resolve(it) } + .filter { tool -> + // ponytail: filter write tools while read-before-write block is active; restored once a read completes + !isReadOnlyMode(sessionId) || ToolCapability.FILE_WRITE !in tool.requiredCapabilities + } .map { tool -> ToolDefinition( function = ToolFunction( diff --git a/examples/workflows/prompts/task_planner.md b/examples/workflows/prompts/task_planner.md new file mode 100644 index 00000000..74a0353b --- /dev/null +++ b/examples/workflows/prompts/task_planner.md @@ -0,0 +1,57 @@ +You are a **Task Planner**. Your sole job is to understand the user's request (in the decision +history above), sweep the relevant parts of the codebase, and emit a task graph via `task_decompose`. +You do not design, implement, or review anything. You stop after the task graph is created. + +## Step 1 — Check for existing work + +Before deriving anything new, run `task_search` with keywords from the request. If an existing task +already covers the goal, load it with `task_context` and name it in your summary. Do not duplicate it. + +## Step 2 — Sweep the codebase (read-only) + +Use `file_read` (also lists a directory when given a directory path), `ShellTool` (grep, find, ls) +to locate the files, modules, and patterns the request touches. Read broadly enough to understand: +- what already exists +- what the request changes or adds +- what depends on what (dependency seams — things that must land before others) + +The repo map and L3 context are already injected above; use them as your starting index, then +drill in with file_read/grep for specifics. + +## Step 3 — Emit the task graph (REQUIRED) + +You **must** call `task_decompose` before calling `stage_complete`. The kernel enforces this — +`stage_complete` is blocked until `task_decompose` returns successfully. + +Shape the graph by the natural seams you found: + +- **Dependency seams**: things that must land before others → DEPENDS_ON links between children. +- **Independent review/handoff points**: pieces worth shipping or reviewing on their own → separate children. +- **Single coherent unit** with no seams → use `task_create` instead (one task, no graph). + +For `task_decompose`, provide: +- `project`: the project key (e.g. `correx`). +- `parent`: an umbrella epic with `title` + `goal` describing the whole request. +- `tasks`: the children, each with `title`, `goal`, `acceptance_criteria` (concrete, checkable), + `affected_paths` (workspace-relative globs), and `depends_on` (refs to other tasks in this batch + that must finish first). + +Keep the graph small. Over-splitting is worse than under-splitting — a session works one task at a +time, and each child is a future claim. Aim for 2–5 children unless the request genuinely demands +more. + +## Step 4 — Clarify only if genuinely blocked + +If something ambiguously blocks a plan (a fork only the user can resolve, a missing decision), add +a `questions` array to your final summary message **before** calling `task_decompose`. Each entry: +- `prompt` (required): the question. +- `options` (optional): suggested answers as strings. +- `header` (optional): 1–2 word label. + +Ask nothing you can answer by reading the code. This is the rare case — when in doubt, make the +reasonable call and decompose. + +## Step 5 — Call stage_complete + +After `task_decompose` (or `task_create`) returns successfully, call `stage_complete`. The workflow +ends. No execution, no handoff, no further stages. diff --git a/examples/workflows/task_planning.toml b/examples/workflows/task_planning.toml new file mode 100644 index 00000000..969463a7 --- /dev/null +++ b/examples/workflows/task_planning.toml @@ -0,0 +1,21 @@ +id = "task_planning" +description = "Sweep the codebase, understand the request, emit a task graph via task_decompose — then stop. No execution." +start = "planner" + +# Single stage: read the request + codebase, emit a task graph, done. +# require_task_decompose = true is a deterministic kernel gate: stage_complete is blocked +# until task_decompose (or task_create for a trivial single-task request) returns successfully. +# The model cannot skip task creation — this is enforced by the orchestrator, not just the prompt. +[[stages]] +id = "planner" +prompt = "prompts/task_planner.md" +allowed_tools = ["file_read", "ShellTool", "task_search", "task_context", "task_decompose", "task_create"] +require_task_decompose = true +token_budget = 16384 +max_retries = 2 + +[[transitions]] +id = "planner-to-done" +from = "planner" +to = "done" +condition_type = "always_true" diff --git a/infrastructure/workflow/src/main/kotlin/com/correx/infrastructure/workflow/TomlWorkflowLoader.kt b/infrastructure/workflow/src/main/kotlin/com/correx/infrastructure/workflow/TomlWorkflowLoader.kt index a51a3de4..e3ac9e73 100644 --- a/infrastructure/workflow/src/main/kotlin/com/correx/infrastructure/workflow/TomlWorkflowLoader.kt +++ b/infrastructure/workflow/src/main/kotlin/com/correx/infrastructure/workflow/TomlWorkflowLoader.kt @@ -49,6 +49,7 @@ private data class StageSection( @param:JsonProperty("inject_artifact_kinds") val injectArtifactKinds: Boolean = false, @param:JsonProperty("ground_references") val groundReferences: Boolean = false, @param:JsonProperty("brief_echo") val briefEcho: Boolean = false, + @param:JsonProperty("require_task_decompose") val requireTaskDecompose: Boolean = false, ) // condition fields flattened into the transition row @@ -125,6 +126,7 @@ class TomlWorkflowLoader( if (s.injectArtifactKinds) put("injectArtifactKinds", "true") if (s.groundReferences) put("groundReferences", "true") if (s.briefEcho) put("briefEcho", "true") + if (s.requireTaskDecompose) put("requireTaskDecompose", "true") }, ) } diff --git a/testing/integration/src/test/kotlin/ToolCallGateTest.kt b/testing/integration/src/test/kotlin/ToolCallGateTest.kt index ae7f6e72..900201e8 100644 --- a/testing/integration/src/test/kotlin/ToolCallGateTest.kt +++ b/testing/integration/src/test/kotlin/ToolCallGateTest.kt @@ -14,9 +14,12 @@ import kotlinx.datetime.Clock import com.correx.core.artifacts.DefaultArtifactReducer import com.correx.core.events.events.ApprovalRequestedEvent import com.correx.core.events.events.ToolCallAssessedEvent +import com.correx.core.events.events.ToolExecutionCompletedEvent import com.correx.core.events.events.ToolExecutionRejectedEvent import com.correx.core.events.events.ToolInvocationRequestedEvent import com.correx.core.events.events.ToolRequest +import com.correx.core.toolintent.rules.ReadBeforeWriteRule +import com.correx.core.tools.contract.ParamRole import com.correx.core.events.execution.RetryPolicy import com.correx.core.events.risk.RiskAction import com.correx.core.events.types.ProviderId @@ -410,6 +413,168 @@ class ToolCallGateTest { ) } + /** + * Read-before-write narrowing: after the READ_BEFORE_WRITE rule blocks a write, the next + * inference turn must only offer read tools; once a file_read completes, write tools return. + */ + @Test + fun `read-before-write block narrows tools to read-only then restores writes after a read`(): Unit = runBlocking { + val targetPath = "/work/Foo.kt" + + // Registry with both a file_write and a file_read tool, paramRoles declare the path param. + val fileWriteTool = object : Tool { + override val name = "file_write" + override val description = "write" + override val parametersSchema: JsonObject = buildJsonObject {} + override val tier = Tier.T1 + override val requiredCapabilities: Set = setOf(ToolCapability.FILE_WRITE) + override val paramRoles: Map = mapOf("path" to ParamRole.PATH) + override fun validateRequest(request: ToolRequest): ValidationResult = ValidationResult.Valid + } + val fileReadTool = object : Tool { + override val name = "file_read" + override val description = "read" + override val parametersSchema: JsonObject = buildJsonObject {} + override val tier = Tier.T1 + override val requiredCapabilities: Set = setOf(ToolCapability.FILE_READ) + override val paramRoles: Map = mapOf("path" to ParamRole.PATH) + override fun validateRequest(request: ToolRequest): ValidationResult = ValidationResult.Valid + } + val twoToolRegistry = object : ToolRegistry { + override fun resolve(name: String): Tool? = when (name) { + "file_write" -> fileWriteTool + "file_read" -> fileReadTool + else -> null + } + override fun all(): List = listOf(fileWriteTool, fileReadTool) + } + + // Provider: turn 1 → write, turn 2 → read, turn 3 → write, turn 4 → stop. + // Records the tool names offered on each inference request. + val offeredToolsPerTurn = mutableListOf>() + var turnCount = 0 + val provider = object : InferenceProvider { + override val id = ProviderId("rbw-test") + override val name = "rbw-test" + override val tokenizer: Tokenizer = MockTokenizer() + override suspend fun infer(request: InferenceRequest): InferenceResponse { + offeredToolsPerTurn += request.tools.map { it.function.name }.toSet() + turnCount++ + val toolCall = when (turnCount) { + 1 -> ToolCallRequest(id = "tc-write-1", function = ToolCallFunction("file_write", """{"path":"$targetPath"}""")) + 2 -> ToolCallRequest(id = "tc-read-1", function = ToolCallFunction("file_read", """{"path":"$targetPath"}""")) + 3 -> ToolCallRequest(id = "tc-write-2", function = ToolCallFunction("file_write", """{"path":"$targetPath"}""")) + else -> null + } + return if (toolCall != null) { + InferenceResponse(request.requestId, "", FinishReason.ToolCall, TokenUsage(1, 1), 0, listOf(toolCall)) + } else { + InferenceResponse(request.requestId, "done", FinishReason.Stop, TokenUsage(1, 1), 0) + } + } + override suspend fun healthCheck(): ProviderHealth = ProviderHealth.Healthy + override fun capabilities(): Set = setOf(CapabilityScore(ModelCapability.General, 1.0)) + } + + // WorldProbe: the target file always exists (triggers the rule) and resolves to itself. + val fakeProbe = object : com.correx.core.toolintent.WorldProbe { + override fun exists(path: java.nio.file.Path): Boolean = true + override fun resolveReal(path: java.nio.file.Path): java.nio.file.Path = path.toAbsolutePath().normalize() + } + + val eventStore = InMemoryEventStore() + val artifactStore = NoopArtifactStore() + val assessor = ToolCallAssessor(listOf(ReadBeforeWriteRule())) + val policy = WorkspacePolicy(workspace) + + val repositories = OrchestratorRepositories( + eventStore = eventStore, + inferenceRepository = InferenceRepository(object : EventReplayer { + override fun rebuild(sessionId: SessionId) = InferenceState() + }), + orchestrationRepository = OrchestrationRepository( + DefaultEventReplayer(eventStore, OrchestrationProjector(DefaultOrchestrationReducer())), + ), + sessionRepository = DefaultSessionRepository( + DefaultEventReplayer(eventStore, SessionProjector(DefaultSessionReducer())), + ), + artifactRepository = LiveArtifactRepository(eventStore, DefaultArtifactReducer()), + approvalRepository = DefaultApprovalRepository( + DefaultEventReplayer(eventStore, ApprovalProjector(DefaultApprovalReducer())), + ), + ) + val engines = OrchestratorEngines( + transitionResolver = DefaultTransitionResolver { _, _ -> true }, + contextPackBuilder = ContextFixtures.simpleBuilder(), + inferenceRouter = object : InferenceRouter { + override suspend fun route(stageId: StageId, requiredCapabilities: Set): InferenceProvider = provider + }, + validationPipeline = ValidationPipeline(validators = listOf(cyclePolicyMissingValidator())), + approvalEngine = DefaultApprovalEngine(), + riskAssessor = DefaultRiskAssessor(), + toolExecutor = object : ToolExecutor { + override suspend fun execute(request: ToolRequest): ToolResult = + ToolResult.Success(request.invocationId, "ok", 0) + }, + toolRegistry = twoToolRegistry, + toolCallAssessor = assessor, + workspacePolicy = policy, + worldProbe = fakeProbe, + ) + val orchestrator = DefaultSessionOrchestrator( + repositories = repositories, + engines = engines, + retryCoordinator = DefaultRetryCoordinator(eventStore), + artifactStore = artifactStore, + decisionJournalRepository = DefaultDecisionJournalRepository( + DefaultEventReplayer(eventStore, DecisionJournalProjector(DefaultDecisionJournalReducer())), + ), + ) + + val sessionId = SessionId("rbw-narrow") + val graph = WorkflowGraph( + id = "rbw-test", + stages = mapOf( + StageId("A") to StageConfig(allowedTools = setOf("file_write", "file_read")), + ), + transitions = setOf( + TransitionEdge(TransitionId("t1"), StageId("A"), StageId("done"), condition = { true }), + ), + start = StageId("A"), + ) + orchestrator.run(sessionId, graph, OrchestrationConfig(retryPolicy = RetryPolicy(maxAttempts = 1, backoffMs = 0))) + + // Turn 1: model asked for file_write → blocked by READ_BEFORE_WRITE. + // Turn 2: offered tools must NOT include file_write (read-only mode active). + assertTrue(offeredToolsPerTurn.size >= 2, "expected at least 2 inference turns, got ${offeredToolsPerTurn.size}") + val turn2Tools = offeredToolsPerTurn[1] + assertTrue("file_write" !in turn2Tools, "turn 2 must not offer file_write (read-only mode), got: $turn2Tools") + assertTrue("file_read" in turn2Tools, "turn 2 must offer file_read, got: $turn2Tools") + + // Turn 3: after the file_read completes, write tools are restored. + assertTrue(offeredToolsPerTurn.size >= 3, "expected at least 3 inference turns, got ${offeredToolsPerTurn.size}") + val turn3Tools = offeredToolsPerTurn[2] + assertTrue("file_write" in turn3Tools, "turn 3 must offer file_write (read completed), got: $turn3Tools") + + // Verify the READ_BEFORE_WRITE block event was recorded (turn 1 block). + val events = eventStore.read(sessionId) + assertTrue( + events.any { e -> + val p = e.payload as? ToolCallAssessedEvent + p != null && p.issues.any { it.code == "READ_BEFORE_WRITE" } + }, + "expected a READ_BEFORE_WRITE ToolCallAssessedEvent", + ) + // Turn 3's file_write must succeed (executor runs, ToolExecutionCompletedEvent for file_write). + assertTrue( + events.any { e -> + val p = e.payload as? ToolExecutionCompletedEvent + p != null && p.toolName == "file_write" + }, + "expected file_write to complete successfully on turn 3", + ) + } + @Test fun `empty allowedTools denies all domain tools and emits ToolExecutionRejectedEvent`(): Unit = runBlocking { val executor = RecordingExecutor()