feat(kernel): read-only mode after read-before-write + mandatory task-decompose gate

Two orchestrator tool-gating behaviors (both touch SessionOrchestrator): - Read-only switch: when a tool call is BLOCKed with READ_BEFORE_WRITE, derive a read-only mode from the event log (block until a later FILE_READ completes) and withhold FILE_WRITE tools from the next inference turn(s), so a weak model is pushed down the read-then-write path instead of looping on the blocked write. - require_task_decompose stage flag (TomlWorkflowLoader) + owesTaskDecompose guard: a stage so marked cannot stage_complete until a task_decompose/task_create has succeeded this stage. New task_planning workflow + prompt are the first consumer (decompose-only: swoop codebase, emit a parent + DEPENDS_ON task graph, stop).
2026-06-28 20:43:18 +04:00
parent 201b599472
commit 1cb7fec677
5 changed files with 318 additions and 3 deletions
@@ -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<String>()
+        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(
@@ -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.
@@ -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"
@@ -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")
                },
            )
        }
@@ -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<ToolCapability> = setOf(ToolCapability.FILE_WRITE)
+            override val paramRoles: Map<String, ParamRole> = 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<ToolCapability> = setOf(ToolCapability.FILE_READ)
+            override val paramRoles: Map<String, ParamRole> = 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<Tool> = 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<Set<String>>()
+        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<CapabilityScore> = 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<InferenceState> {
+                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<ModelCapability>): 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()