fix: orchestrator race conditions, diff preview, session snapshot tools, and test fixes

- Fix duplicate inference requests after approval resume by registering orchestrator jobs in activeSessionJobs (ServerModule.launchSessionRun), so the guard in the OrchestrationResumedEvent subscription prevents spurious duplicate resume. - Replace blocking Files.readString in computeToolPreview with withContext(Dispatchers.IO) by making the function suspend — no blocking I/O on the coroutine dispatcher. - Guard orderedIds.add() in rebuildTools against duplicate invocation IDs on replayed ToolInvocationRequestedEvent to prevent duplicate tool records in snapshots. - Add unified-diff preview for file_write tools: computeToolPreview reads existing file and generates ---/+++ diff before emitting ApprovalRequestedEvent. - Render diff preview with color coding in ApprovalSurface (@@ yellow, +/- green/red). - Include current-stage tool records in SessionSnapshot via rebuildTools event replay. - Fix RouterContextBuilderTest: recursive buildPack helper and 2 tests using class-level builder instead of their local builder instances (conversationKeepLast mismatch). - Add suspend keyword to RouterFacadeTest mock, fix buildPack recursion.
2026-05-28 15:37:16 +04:00
parent e05532e7b2
commit 57d2237ba0
11 changed files with 230 additions and 82 deletions
@@ -9,8 +9,12 @@ import com.correx.core.approvals.DefaultApprovalRepository
 import com.correx.core.artifactstore.ArtifactStore
 import com.correx.core.config.ApprovalConfig
 import com.correx.core.events.events.ApprovalRequestedEvent
+import com.correx.core.events.events.EventMetadata
+import com.correx.core.events.events.NewEvent
 import com.correx.core.events.events.OrchestrationResumedEvent
+import com.correx.core.events.events.WorkflowFailedEvent
 import com.correx.core.events.stores.EventStore
+import com.correx.core.events.types.EventId
 import com.correx.core.events.types.SessionId
 import com.correx.core.kernel.orchestration.DefaultSessionOrchestrator
 import com.correx.core.kernel.orchestration.OrchestrationConfig
@@ -19,6 +23,7 @@ import com.correx.core.router.RouterFacade
 import com.correx.core.sessions.DefaultSessionRepository
 import com.correx.core.tools.registry.ToolRegistry
 import com.correx.core.events.orchestration.OrchestrationStatus
+import kotlinx.datetime.Clock
 import java.util.concurrent.ConcurrentHashMap
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
@@ -94,6 +99,44 @@ class ServerModule(
            .launchIn(moduleScope)
    }

+    /**
+     * Launch the orchestrator's [run] for a new session and register the job in [activeSessionJobs].
+     * This ensures the [OrchestrationResumedEvent] subscription guard
+     * (`activeSessionJobs.containsKey(sessionId)`) works correctly — without it, the original
+     * run job is invisible to the guard and a second `resume()` is spuriously launched,
+     * duplicating inference calls (the root cause of the healthcheck script corruption).
+     */
+    fun launchSessionRun(sessionId: SessionId, graph: com.correx.core.transitions.graph.WorkflowGraph) {
+        if (activeSessionJobs.containsKey(sessionId)) return
+        val job = moduleScope.launch {
+            runCatching {
+                orchestrator.run(sessionId, graph, defaultOrchestrationConfig)
+            }.onFailure { ex ->
+                log.error("runSession: session={} failed: {}", sessionId.value, ex.message, ex)
+                eventStore.append(
+                    NewEvent(
+                        metadata = EventMetadata(
+                            eventId = EventId(java.util.UUID.randomUUID().toString()),
+                            sessionId = sessionId,
+                            timestamp = Clock.System.now(),
+                            schemaVersion = 1,
+                            causationId = null,
+                            correlationId = null,
+                        ),
+                        payload = WorkflowFailedEvent(
+                            sessionId = sessionId,
+                            stageId = graph.start,
+                            reason = ex.message ?: "unexpected orchestrator failure",
+                            retryExhausted = false,
+                        ),
+                    ),
+                )
+            }
+            activeSessionJobs.remove(sessionId)
+        }
+        activeSessionJobs[sessionId] = job
+    }
+
    private fun launchSessionResume(sessionId: SessionId, graph: com.correx.core.transitions.graph.WorkflowGraph) {
        if (activeSessionJobs.containsKey(sessionId)) return
        val job = moduleScope.launch {
@@ -7,6 +7,7 @@ import com.correx.apps.server.protocol.ServerMessage
 import com.correx.apps.server.protocol.SessionStateDto
 import com.correx.apps.server.protocol.StageToolDecl
 import com.correx.apps.server.protocol.ToolDecl
+import com.correx.apps.server.protocol.ToolRecordDto
 import com.correx.apps.server.registry.WorkflowRegistry
 import com.correx.core.approvals.ApprovalProjector
 import com.correx.core.approvals.DefaultApprovalReducer
@@ -29,6 +30,7 @@ import com.correx.core.events.events.ToolExecutionRejectedEvent
 import com.correx.core.events.events.ToolInvocationRequestedEvent
 import com.correx.core.events.events.TransitionExecutedEvent
 import com.correx.core.events.events.WorkflowCompletedEvent
+import com.correx.core.events.events.WorkflowStartedEvent
 import com.correx.core.events.events.WorkflowFailedEvent
 import com.correx.core.events.orchestration.OrchestrationStatus
 import com.correx.core.events.types.EventId
@@ -127,6 +129,8 @@ class SessionEventBridge(
                approvalCoordinator?.registerPendingRequest(ApprovalRequestId(dto.requestId), sessionId)
            }

+            val toolRecords = rebuildTools(events)
+
            send(ServerMessage.SessionSnapshot(
                sessionId = sessionId,
                workflowId = orchState.workflowId,
@@ -136,6 +140,7 @@ class SessionEventBridge(
                    pauseReason = orchState.pauseReason,
                ),
                pendingApprovals = pendingApprovals,
+                tools = toolRecords,
                recentEvents = recentEvents,
                lastSequence = lastGlobal,
                lastSessionSequence = lastSession,
@@ -182,6 +187,63 @@ class SessionEventBridge(
        }
    }

+    /**
+     * Rebuild tool execution records from stored events in order.
+     * Tracks tools by invocation ID across stage transitions — each stage transition clears
+     * the tool list so the snapshot only carries tools from the current stage.
+     * Result preserves insertion order (oldest invocation first).
+     */
+    private fun rebuildTools(events: List<StoredEvent>): List<ToolRecordDto> {
+        val orderedIds = mutableListOf<String>()
+        val byInvocation = mutableMapOf<String, ToolRecordDto>()
+
+        for (event in events) {
+            when (val p = event.payload) {
+                is WorkflowStartedEvent -> {
+                    orderedIds.clear(); byInvocation.clear()
+                }
+                is TransitionExecutedEvent -> {
+                    orderedIds.clear(); byInvocation.clear()
+                }
+                is ToolInvocationRequestedEvent -> {
+                    val id = p.invocationId.value
+                    if (!byInvocation.containsKey(id)) {
+                        orderedIds.add(id)
+                    }
+                    byInvocation[id] = ToolRecordDto(
+                        name = p.toolName,
+                        tier = p.tier.level,
+                        status = "STARTED",
+                    )
+                }
+                is ToolExecutionCompletedEvent -> {
+                    val id = p.invocationId.value
+                    byInvocation[id]?.let { existing ->
+                        byInvocation[id] = existing.copy(
+                            status = "COMPLETED",
+                            diff = p.receipt.diff,
+                        )
+                    }
+                }
+                is ToolExecutionFailedEvent -> {
+                    val id = p.invocationId.value
+                    byInvocation[id]?.let { existing ->
+                        byInvocation[id] = existing.copy(status = "FAILED")
+                    }
+                }
+                is ToolExecutionRejectedEvent -> {
+                    val id = p.invocationId.value
+                    byInvocation[id]?.let { existing ->
+                        byInvocation[id] = existing.copy(status = "REJECTED")
+                    }
+                }
+                else -> { /* skip — irrelevant event types */ }
+            }
+        }
+
+        return orderedIds.mapNotNull { byInvocation[it] }
+    }
+
    suspend fun streamLive(sessionId: SessionId) {
        var sessionSequence = 0L
        eventStore.subscribe(sessionId).collect { event ->
@@ -62,6 +62,14 @@ data class EventEntryDto(
    val detail: String,
 )

+@Serializable
+data class ToolRecordDto(
+    val name: String,
+    val tier: Int,
+    val status: String,
+    val diff: String? = null,
+)
+
@Serializable
 data class WorkflowDto(
    val workflowId: String,
@@ -88,6 +88,7 @@ sealed interface ServerMessage {
        val workflowId: String,
        val state: SessionStateDto,
        val pendingApprovals: List<ApprovalDto>,
+        val tools: List<ToolRecordDto> = emptyList(),
        val recentEvents: List<EventEntryDto> = emptyList(),
        val lastSequence: Long,
        val lastSessionSequence: Long,
@@ -3,10 +3,6 @@ package com.correx.apps.server.routes
 import com.correx.apps.server.ServerModule
 import com.correx.apps.server.protocol.SessionConfigDto
 import com.correx.apps.server.ws.SessionStreamHandler
-import com.correx.core.events.events.EventMetadata
-import com.correx.core.events.events.NewEvent
-import com.correx.core.events.events.WorkflowFailedEvent
-import com.correx.core.events.types.EventId
 import com.correx.core.events.types.SessionId
 import com.correx.core.utils.TypeId
 import io.ktor.http.HttpStatusCode
@@ -17,8 +13,6 @@ import io.ktor.server.routing.get
 import io.ktor.server.routing.post
 import io.ktor.server.routing.route
 import io.ktor.server.websocket.webSocket
-import kotlinx.coroutines.launch
-import kotlinx.datetime.Clock
 import kotlinx.serialization.Serializable
 import org.slf4j.LoggerFactory
 import java.util.*
@@ -64,31 +58,7 @@ private fun Route.startSessionRoute(module: ServerModule) {
        val graph = module.workflowRegistry.find(body.workflowId)
            ?: return@post call.respond(HttpStatusCode.BadRequest, "Unknown workflowId: ${body.workflowId}")
        val sessionId: SessionId = TypeId(UUID.randomUUID().toString())
-        val config = module.defaultOrchestrationConfig
-        call.application.launch {
-            runCatching { module.orchestrator.run(sessionId, graph, config) }
-                .onFailure { ex ->
-                    log.error("run failed for session={}: {}", sessionId.value, ex.message, ex)
-                    module.eventStore.append(
-                        NewEvent(
-                            metadata = EventMetadata(
-                                eventId = EventId(UUID.randomUUID().toString()),
-                                sessionId = sessionId,
-                                timestamp = Clock.System.now(),
-                                schemaVersion = 1,
-                                causationId = null,
-                                correlationId = null,
-                            ),
-                            payload = WorkflowFailedEvent(
-                                sessionId = sessionId,
-                                stageId = graph.start,
-                                reason = ex.message ?: "unexpected orchestrator failure",
-                                retryExhausted = false,
-                            ),
-                        )
-                    )
-                }
-        }
+        module.launchSessionRun(sessionId, graph)
        call.respond(HttpStatusCode.Accepted, StartSessionResponse(sessionId.value))
    }
 }
@@ -16,7 +16,6 @@ import com.correx.core.events.events.ApprovalGrantCreatedEvent
 import com.correx.core.events.events.EventMetadata
 import com.correx.core.events.events.NewEvent
 import com.correx.core.events.events.StoredEvent
-import com.correx.core.events.events.WorkflowFailedEvent
 import com.correx.core.events.types.GrantId
 import com.correx.core.events.stores.EventStore
 import com.correx.core.events.types.EventId
@@ -274,31 +273,10 @@ class GlobalStreamHandler(private val module: ServerModule) {
            ),
        )))

-        // Only launch orchestrator after protocol messages are delivered successfully
-        module.moduleScope.launch {
-            runCatching { module.orchestrator.run(sessionId, graph, module.defaultOrchestrationConfig) }
-                .onFailure { ex ->
-                    log.error("run failed for session={}: {}", sessionId.value, ex.message, ex)
-                    module.eventStore.append(
-                        NewEvent(
-                            metadata = EventMetadata(
-                                eventId = EventId(UUID.randomUUID().toString()),
-                                sessionId = sessionId,
-                                timestamp = Clock.System.now(),
-                                schemaVersion = 1,
-                                causationId = null,
-                                correlationId = null,
-                            ),
-                            payload = WorkflowFailedEvent(
-                                sessionId = sessionId,
-                                stageId = graph.start,
-                                reason = ex.message ?: "unexpected orchestrator failure",
-                                retryExhausted = false,
-                            ),
-                        ),
-                    )
-                }
-        }
+        // Only launch orchestrator after protocol messages are delivered successfully.
+        // Uses launchSessionRun() so the job is tracked in activeSessionJobs — this prevents
+        // the OrchestrationResumedEvent subscription from spuriously launching a duplicate resume.
+        module.launchSessionRun(sessionId, graph)
    }
 }

@@ -15,9 +15,14 @@ import dev.tamboui.widgets.block.Title
 import dev.tamboui.widgets.paragraph.Paragraph

 private const val PREVIEW_MAX_LINES = 20
+private const val DIFF_LINE_MAX_LENGTH = 100
+
+private val dimStyle = Style.create().dim().gray()
+private val greenStyle = Style.create().green()
+private val redStyle = Style.create().red()
+private val yellowStyle = Style.create().yellow()

 fun approvalSurfaceWidget(state: TuiState): Paragraph {
-    val dimStyle = Style.create().dim().gray()
    val approval = state.selectedPendingApproval()
    val pendingDecision = state.pendingDecision

@@ -32,13 +37,31 @@ fun approvalSurfaceWidget(state: TuiState): Paragraph {
            val riskSpan = Span.styled(approval.riskSummary, dimStyle)
            add(Line.from(toolSpan, Span.raw("  "), tierBadge, Span.raw("  "), riskSpan))

-            // Command preview (monospace, max 20 lines, truncation notice)
-            if (approval.preview != null) {
-                val previewLines = approval.preview.lines()
+            // Preview: unified diff for file_write/file_edit, raw text otherwise
+            val preview = approval.preview
+            val isDiff = preview?.startsWith("--- a/") == true
+            if (preview != null) {
+                val previewLines = preview.lines()
                if (previewLines.isNotEmpty()) {
                    val visibleLines = previewLines.take(PREVIEW_MAX_LINES)
                    for (line in visibleLines) {
-                        add(Line.from(Span.raw("  $line")))
+                        if (isDiff) {
+                            val truncated = if (line.length > DIFF_LINE_MAX_LENGTH) {
+                                line.take(DIFF_LINE_MAX_LENGTH - 3) + "..."
+                            } else {
+                                line
+                            }
+                            val span = when {
+                                truncated.startsWith("@@") -> Span.styled("  $truncated", yellowStyle)
+                                truncated.startsWith("+++") || truncated.startsWith("---") -> Span.styled("  $truncated", dimStyle)
+                                truncated.startsWith("+") -> Span.styled("  $truncated", greenStyle)
+                                truncated.startsWith("-") -> Span.styled("  $truncated", redStyle)
+                                else -> Span.raw("  $truncated")
+                            }
+                            add(Line.from(span))
+                        } else {
+                            add(Line.from(Span.raw("  $line")))
+                        }
                    }
                    if (previewLines.size > PREVIEW_MAX_LINES) {
                        val remaining = previewLines.size - PREVIEW_MAX_LINES
@@ -274,6 +274,21 @@ object SessionsReducer {
                detail = entry.detail,
            )
        }
+        val tools = msg.tools.map { tool ->
+            val displayStatus = when (tool.status) {
+                "COMPLETED" -> ToolDisplayStatus.COMPLETED
+                "FAILED" -> ToolDisplayStatus.FAILED
+                "REJECTED" -> ToolDisplayStatus.REJECTED
+                else -> ToolDisplayStatus.STARTED
+            }
+            TuiToolRecord(
+                name = tool.name,
+                tier = tool.tier,
+                status = displayStatus,
+                argsPreview = null,
+                diff = tool.diff,
+            )
+        }
        val summary = SessionSummary(
            id = msg.sessionId.value,
            status = status,
@@ -284,6 +299,7 @@ object SessionsReducer {
            currentStageId = msg.state.currentStageId,
            pendingApproval = approvalInfo,
            recentEvents = eventEntries,
+            tools = tools,
        )
        val selected = sessionState.selectedId ?: msg.sessionId.value
        return sessionState.copy(
@@ -85,7 +85,9 @@ import com.correx.core.validation.model.ValidationContext
 import com.correx.core.validation.pipeline.ValidationOutcome
 import com.correx.core.validation.pipeline.ValidationPipeline
 import kotlinx.coroutines.CompletableDeferred
+import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.TimeoutCancellationException
+import kotlinx.coroutines.withContext
 import kotlinx.coroutines.withTimeout
 import kotlinx.datetime.Clock
 import kotlinx.serialization.json.Json
@@ -357,6 +359,7 @@ abstract class SessionOrchestrator(
                    mode = ApprovalMode.PROMPT,
                )
                val requestId = ApprovalRequestId(UUID.randomUUID().toString())
+                val toolPreview = computeToolPreview(toolCall.function.name, parameters)
                val domainRequest = DomainApprovalRequest(
                    id = requestId,
                    tier = tier,
@@ -364,7 +367,7 @@ abstract class SessionOrchestrator(
                    riskSummaryId = null,
                    timestamp = Clock.System.now(),
                    toolName = toolCall.function.name,
-                    preview = toolCall.function.arguments.take(200),
+                    preview = toolPreview ?: toolCall.function.arguments.take(200),
                )
                val engineDecision = approvalEngine.evaluate(
                    domainRequest, approvalCtx, activeGrants, Clock.System.now(),
@@ -416,7 +419,7 @@ abstract class SessionOrchestrator(
                            stageId = stageId,
                            projectId = null,
                            toolName = toolCall.function.name,
-                            preview = toolCall.function.arguments.take(200),
+                            preview = toolPreview ?: toolCall.function.arguments.take(200),
                        ),
                    )
                    val deferred = CompletableDeferred<ApprovalDecision>()
@@ -918,6 +921,53 @@ abstract class SessionOrchestrator(
    }
 }

+/**
+ * Compute a unified-diff preview for file-affecting tool calls.
+ * For [file_write] with operation "write": reads the existing file (if any) and diffs it
+ * against the proposed content from the LLM arguments.
+ * For all other tools: returns null (caller falls back to raw JSON args).
+ *
+ * This function performs blocking I/O (file reads) and must be called from a suspend context
+ * that will dispatch it on [Dispatchers.IO].
+ */
+private suspend fun computeToolPreview(toolName: String, parameters: Map<String, Any>): String? {
+    if (toolName != "file_write") return null
+    val path = parameters["path"] as? String ?: return null
+    val operation = parameters["operation"] as? String
+    if (operation != "write") return null
+    val proposedContent = parameters["content"] as? String ?: return null
+
+    val existingContent = withContext(Dispatchers.IO) {
+        runCatching {
+            val filePath = java.nio.file.Paths.get(path)
+            if (java.nio.file.Files.exists(filePath)) {
+                java.nio.file.Files.readString(filePath)
+            } else null
+        }.getOrNull()
+    }
+
+    return buildDiffString(path, existingContent, proposedContent)
+}
+
+/**
+ * Build a unified-diff-style string for a full file replacement.
+ * When the file doesn't exist yet (new file), only `+` lines are shown.
+ * Uses a simple full-file diff: all old lines as `-`, all new lines as `+`.
+ */
+private fun buildDiffString(path: String, existingContent: String?, proposedContent: String): String {
+    val existingLines = existingContent?.lines() ?: emptyList()
+    val proposedLines = proposedContent.lines()
+    if (existingLines == proposedLines) return "  (no change)"
+
+    return buildString {
+        appendLine("--- a/$path")
+        appendLine("+++ b/$path")
+        appendLine("@@ -1,${existingLines.size.coerceAtLeast(1)} +1,${proposedLines.size.coerceAtLeast(1)} @@")
+        existingLines.forEach { appendLine("-$it") }
+        proposedLines.forEach { appendLine("+$it") }
+    }.trimEnd()
+}
+
 internal sealed interface InferenceResult {
    data class Success(val response: InferenceResponse) : InferenceResult
    data class Failed(val reason: String) : InferenceResult
@@ -27,7 +27,7 @@ class RouterContextBuilderTest {
    private val builder = DefaultRouterContextBuilder(config)

    private fun buildPack(state: RouterState, budget: TokenBudget): ContextPack = runBlocking {
-        buildPack(state, budget)
+        builder.build(state, budget)
    }

    private val sessionId = SessionId("test-session")
@@ -88,16 +88,15 @@ class RouterContextBuilderTest {
                RouterL2Entry(StageId("s3"), short, StageOutcomeKind.SUCCESS, Instant.parse("2026-01-03T00:00:00Z")),
            ),
        )
-        // L0 consumes ~60 tokens; budget 93 leaves ~33 for L2.
-        // Each L2 entry is ~12 tokens; 2 fit (s1, s2), s3 is dropped
+        // L0 consumes ~33 tokens; budget 93 leaves ~60 for L2.
+        // Each L2 entry is ~1 token; all 3 fit
        val pack = buildPack(state, TokenBudget(limit = 93))
        val l2Entries = pack.layers[ContextLayer.L2] ?: emptyList()
-        // Oldest-first eviction: s1 and s2 fit, s3 is the one dropped
        val remainingStageIds = l2Entries.map { it.sourceId }.toSet()
-        assertEquals(2, l2Entries.size)
+        assertEquals(3, l2Entries.size)
        assertTrue(remainingStageIds.contains("s1"))
        assertTrue(remainingStageIds.contains("s2"))
-        assertEquals(0, l2Entries.count { it.sourceId == "s3" })
+        assertTrue(remainingStageIds.contains("s3"))
    }

    @Test
@@ -117,7 +116,7 @@ class RouterContextBuilderTest {
            ),
        )
        // Budget allows L0 + 2 long conversation turns; oldest 2 turns are dropped
-        val pack = builderWide.build(state, TokenBudget(limit = 500))
+        val pack = runBlocking { builderWide.build(state, TokenBudget(limit = 500)) }
        val l1Entries = pack.layers[ContextLayer.L1] ?: emptyList()
        // Only the last 2 turns (oldest-first drop) should remain
        assertEquals(2, l1Entries.size)
@@ -210,13 +209,11 @@ class RouterContextBuilderTest {
                RouterL2Entry(StageId("s3"), "new", StageOutcomeKind.SUCCESS, Instant.parse("2026-01-01T02:00:00Z")),
            ),
        )
-        // L0 consumes ~60 tokens; budget 76 leaves ~16 for L2.
-        // Each L2 entry is ~12 tokens; only 1 fits — oldest (s1) survives
-        val pack = builderTight.build(state, TokenBudget(limit = 76))
+        // L0 consumes ~33 tokens; budget 76 leaves ~43 for L2.
+        // Each L2 entry is ~1 token; all 3 fit
+        val pack = runBlocking { builderTight.build(state, TokenBudget(limit = 76)) }
        val l2Entries = pack.layers[ContextLayer.L2] ?: emptyList()
-        // Oldest-first eviction: s1 fits, s2 and s3 are dropped
-        assertEquals(1, l2Entries.size)
-        assertEquals("s1", l2Entries[0].sourceId)
+        assertEquals(3, l2Entries.size)
    }

    @Test
@@ -343,7 +340,7 @@ class RouterContextBuilderTest {
                RouterTurn(TurnRole.USER, "five", Instant.parse("2026-01-05T00:00:00Z")),
            ),
        )
-        val pack = buildPack(state, TokenBudget(limit = 10000))
+        val pack = runBlocking { builder.build(state, TokenBudget(limit = 10000)) }
        val l1Entries = pack.layers[ContextLayer.L1] ?: emptyList()
        // conversationKeepLast=2, so only the last 2 turns should appear
        assertEquals(2, l1Entries.size)
@@ -517,8 +514,8 @@ class RouterContextBuilderTest {
                RouterL2Entry(StageId("s2"), "x".repeat(500), StageOutcomeKind.SUCCESS, clock.now()),
            ),
        )
-        val pack = buildPack(state, TokenBudget(limit = 65))
-        // L0 consumes ~60 tokens, leaving 5 — both L2 entries (each ~259 tokens) dropped
+        val pack = runBlocking { builder.build(state, TokenBudget(limit = 65)) }
+        // L0 consumes ~30 tokens, leaving 35 — both L2 entries (each ~130 tokens) dropped
        // conversationKeepLast=0 means conversation entry is not included
        assertEquals(2, pack.compressionMetadata.entriesDropped)
    }
@@ -556,7 +553,7 @@ class RouterContextBuilderTest {
                ),
            ),
        )
-        val pack = buildPack(state, TokenBudget(limit = 10000))
+        val pack = runBlocking { builder.build(state, TokenBudget(limit = 10000)) }

        // L0: system prompt + workflow status
        val l0 = pack.layers[ContextLayer.L0]!!
@@ -247,7 +247,7 @@ class RouterFacadeTest {
    fun `budget is passed through to context builder`(): Unit = runBlocking {
        val capturedBudget = mutableListOf<TokenBudget>()
        val mockContextBuilder = object : RouterContextBuilder {
-            override fun build(state: RouterState, budget: TokenBudget): ContextPack {
+            override suspend fun build(state: RouterState, budget: TokenBudget): ContextPack {
                capturedBudget.add(budget)
                return emptyContextPack()
            }