fix: implement stage_complete tool for LLM-driven stage completion

- Add STAGE_COMPLETE_TOOL constant and meta-tool handler in executeStage
  loop: when LLM calls stage_complete, break out of tool dispatch loop
  and proceed to normal validation/transition
- Inject stage_complete tool definition in every inference request so
  the LLM always has the option to signal stage completion
- Replace expand-all diff with scrollable diff window in approval
  surface (ctrl+x toggles, up/down scrolls, shows line range)
This commit is contained in:
2026-05-29 17:40:51 +04:00
parent 7c55a53a9f
commit 2127824942
5 changed files with 41 additions and 11 deletions
@@ -27,6 +27,7 @@ fun approvalSurfaceWidget(state: TuiState): Paragraph {
val pendingDecision = state.pendingDecision val pendingDecision = state.pendingDecision
val diffExpanded = state.diffExpanded val diffExpanded = state.diffExpanded
val scrollOffset = state.diffScrollOffset
val lines = buildList<Line> { val lines = buildList<Line> {
if (approval == null) { if (approval == null) {
@@ -45,8 +46,11 @@ fun approvalSurfaceWidget(state: TuiState): Paragraph {
if (preview != null) { if (preview != null) {
val previewLines = preview.lines() val previewLines = preview.lines()
if (previewLines.isNotEmpty()) { if (previewLines.isNotEmpty()) {
val totalLines = previewLines.size
val visibleLines = if (diffExpanded) { val visibleLines = if (diffExpanded) {
previewLines val from = scrollOffset.coerceIn(0, maxOf(0, totalLines - 1))
val to = minOf(from + PREVIEW_MAX_LINES, totalLines)
previewLines.subList(from, to)
} else { } else {
previewLines.take(PREVIEW_MAX_LINES) previewLines.take(PREVIEW_MAX_LINES)
} }
@@ -69,12 +73,13 @@ fun approvalSurfaceWidget(state: TuiState): Paragraph {
add(Line.from(Span.raw(" $line"))) add(Line.from(Span.raw(" $line")))
} }
} }
if (!diffExpanded && previewLines.size > PREVIEW_MAX_LINES) { if (!diffExpanded && totalLines > PREVIEW_MAX_LINES) {
val remaining = previewLines.size - PREVIEW_MAX_LINES val remaining = totalLines - PREVIEW_MAX_LINES
add(Line.from(Span.styled(" (preview truncated, $remaining more lines — ctrl+x to expand)", dimStyle))) add(Line.from(Span.styled(" (preview truncated, $remaining more lines — ctrl+x to expand, \u2191\u2193 to scroll)", dimStyle)))
} }
if (diffExpanded) { if (diffExpanded) {
add(Line.from(Span.styled(" (full diff — ctrl+x to collapse)", dimStyle))) val showingTo = minOf(scrollOffset + PREVIEW_MAX_LINES, totalLines)
add(Line.from(Span.styled(" lines ${scrollOffset + 1}\u2013$showingTo of $totalLines\u2191\u2193 scroll, ctrl+x collapse", dimStyle)))
} }
} else { } else {
add(Line.from(Span.styled(" no preview available", dimStyle))) add(Line.from(Span.styled(" no preview available", dimStyle)))
@@ -72,6 +72,8 @@ object KeyResolver {
KeyEvent.Enter -> Action.SubmitApprovalDecision KeyEvent.Enter -> Action.SubmitApprovalDecision
KeyEvent.Escape -> Action.CancelInput KeyEvent.Escape -> Action.CancelInput
KeyEvent.Cancel -> Action.CancelInput KeyEvent.Cancel -> Action.CancelInput
KeyEvent.NavUp -> Action.NavigateUp
KeyEvent.NavDown -> Action.NavigateDown
KeyEvent.ReturnToSessionList -> Action.ReturnToSessionList KeyEvent.ReturnToSessionList -> Action.ReturnToSessionList
KeyEvent.ToggleEventStrip -> Action.ToggleEventStrip KeyEvent.ToggleEventStrip -> Action.ToggleEventStrip
KeyEvent.ToggleDiffExpand -> Action.ToggleDiffExpand KeyEvent.ToggleDiffExpand -> Action.ToggleDiffExpand
@@ -118,7 +118,10 @@ object RootReducer {
else withBgReset.copy(eventStripVisible = !withBgReset.eventStripVisible) else withBgReset.copy(eventStripVisible = !withBgReset.eventStripVisible)
} }
is Action.ToggleDiffExpand -> withBgReset.copy(diffExpanded = !withBgReset.diffExpanded) is Action.ToggleDiffExpand -> withBgReset.copy(
diffExpanded = !withBgReset.diffExpanded,
diffScrollOffset = 0,
)
// Input history append on SubmitInput in IN_SESSION + ROUTER (ARCH-HISTORY-1), // Input history append on SubmitInput in IN_SESSION + ROUTER (ARCH-HISTORY-1),
// or enter a selected session from the list in IDLE. // or enter a selected session from the list in IDLE.
@@ -188,9 +191,11 @@ object RootReducer {
} }
} }
// Input history navigation on NavigateUp in IN_SESSION + ROUTER (ARCH-HISTORY-2/3). // Diff scrolling in APPROVAL mode when expanded.
is Action.NavigateUp -> { is Action.NavigateUp -> {
if (withBgReset.displayState == DisplayState.IN_SESSION && prevInputMode == InputMode.ROUTER) { if (withBgReset.displayState == DisplayState.APPROVAL && withBgReset.diffExpanded) {
withBgReset.copy(diffScrollOffset = maxOf(0, withBgReset.diffScrollOffset - 1))
} else if (withBgReset.displayState == DisplayState.IN_SESSION && prevInputMode == InputMode.ROUTER) {
val selectedId = withBgReset.sessions.selectedId val selectedId = withBgReset.sessions.selectedId
if (selectedId != null) { if (selectedId != null) {
val history = withBgReset.inputHistory[selectedId] ?: emptyList() val history = withBgReset.inputHistory[selectedId] ?: emptyList()
@@ -227,9 +232,11 @@ object RootReducer {
} }
} }
// Input history navigation on NavigateDown in IN_SESSION + ROUTER (ARCH-HISTORY-2/3). // Diff scrolling in APPROVAL mode when expanded.
is Action.NavigateDown -> { is Action.NavigateDown -> {
if (withBgReset.displayState == DisplayState.IN_SESSION && prevInputMode == InputMode.ROUTER) { if (withBgReset.displayState == DisplayState.APPROVAL && withBgReset.diffExpanded) {
withBgReset.copy(diffScrollOffset = withBgReset.diffScrollOffset + 1)
} else if (withBgReset.displayState == DisplayState.IN_SESSION && prevInputMode == InputMode.ROUTER) {
val selectedId = withBgReset.sessions.selectedId val selectedId = withBgReset.sessions.selectedId
if (selectedId != null) { if (selectedId != null) {
val history = withBgReset.inputHistory[selectedId] ?: emptyList() val history = withBgReset.inputHistory[selectedId] ?: emptyList()
@@ -48,6 +48,7 @@ data class TuiState(
val routerConnected: Boolean = false, val routerConnected: Boolean = false,
val routerMessages: Map<String, List<String>> = emptyMap(), val routerMessages: Map<String, List<String>> = emptyMap(),
val diffExpanded: Boolean = false, val diffExpanded: Boolean = false,
val diffScrollOffset: Int = 0,
val chatMode: ChatMode = ChatMode.CHAT, val chatMode: ChatMode = ChatMode.CHAT,
/** /**
* True from connect until SnapshotComplete observed. Reset to true on Disconnected. * True from connect until SnapshotComplete observed. Reset to true on Disconnected.
@@ -92,6 +92,7 @@ import kotlinx.coroutines.withContext
import kotlinx.coroutines.withTimeout import kotlinx.coroutines.withTimeout
import kotlinx.datetime.Clock import kotlinx.datetime.Clock
import kotlinx.serialization.json.Json import kotlinx.serialization.json.Json
import kotlinx.serialization.json.JsonObject
import org.slf4j.LoggerFactory import org.slf4j.LoggerFactory
import java.util.* import java.util.*
import java.util.concurrent.* import java.util.concurrent.*
@@ -99,6 +100,7 @@ import java.util.concurrent.atomic.*
import kotlin.coroutines.cancellation.CancellationException import kotlin.coroutines.cancellation.CancellationException
private const val MAX_TOOL_ROUNDS = 10 private const val MAX_TOOL_ROUNDS = 10
private const val STAGE_COMPLETE_TOOL = "stage_complete"
@SuppressWarnings( @SuppressWarnings(
"ForbiddenComment", "ForbiddenComment",
@@ -258,6 +260,12 @@ abstract class SessionOrchestrator(
inferenceResult.response.finishReason is FinishReason.ToolCall && inferenceResult.response.finishReason is FinishReason.ToolCall &&
toolRounds < MAX_TOOL_ROUNDS toolRounds < MAX_TOOL_ROUNDS
) { ) {
// stage_complete is a meta-tool: the LLM calls it to signal the stage's goal is met.
// Skip tool dispatch — the loop exits and normal validation/transition follows
// (emitToolArtifacts + verifyProduces run on the success path after the loop).
if (inferenceResult.response.toolCalls.any { it.function.name == STAGE_COMPLETE_TOOL }) {
break
}
val toolEntries = dispatchToolCalls(sessionId, stageId, val toolEntries = dispatchToolCalls(sessionId, stageId,
inferenceResult.response.toolCalls, stageConfig) inferenceResult.response.toolCalls, stageConfig)
val fatalEntry = toolEntries.firstOrNull { it.content.startsWith("FATAL:") } val fatalEntry = toolEntries.firstOrNull { it.content.startsWith("FATAL:") }
@@ -670,7 +678,14 @@ abstract class SessionOrchestrator(
parameters = tool.parametersSchema, parameters = tool.parametersSchema,
), ),
) )
}, } + ToolDefinition(
function = ToolFunction(
name = STAGE_COMPLETE_TOOL,
description = "Call this tool when the stage's goal is fully met and no further tool calls are needed. " +
"The orchestrator will proceed to the next stage.",
parameters = JsonObject(emptyMap()),
),
),
) )
// TODO: capture rendered prompt; currently contextpack.toString() // TODO: capture rendered prompt; currently contextpack.toString()