From 57d2237ba048aebdafeaae5645c9a081623812de Mon Sep 17 00:00:00 2001 From: kami Date: Thu, 28 May 2026 15:37:16 +0400 Subject: [PATCH] fix: orchestrator race conditions, diff preview, session snapshot tools, and test fixes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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. --- .../com/correx/apps/server/ServerModule.kt | 43 +++++++++++++ .../apps/server/bridge/SessionEventBridge.kt | 62 +++++++++++++++++++ .../com/correx/apps/server/protocol/Dtos.kt | 8 +++ .../apps/server/protocol/ServerMessage.kt | 1 + .../apps/server/routes/SessionRoutes.kt | 32 +--------- .../apps/server/ws/GlobalStreamHandler.kt | 30 ++------- .../apps/tui/components/ApprovalSurface.kt | 33 ++++++++-- .../apps/tui/reducer/SessionsReducer.kt | 16 +++++ .../orchestration/SessionOrchestrator.kt | 54 +++++++++++++++- .../test/kotlin/RouterContextBuilderTest.kt | 31 +++++----- .../src/test/kotlin/RouterFacadeTest.kt | 2 +- 11 files changed, 230 insertions(+), 82 deletions(-) diff --git a/apps/server/src/main/kotlin/com/correx/apps/server/ServerModule.kt b/apps/server/src/main/kotlin/com/correx/apps/server/ServerModule.kt index ba6e2936..cdef5095 100644 --- a/apps/server/src/main/kotlin/com/correx/apps/server/ServerModule.kt +++ b/apps/server/src/main/kotlin/com/correx/apps/server/ServerModule.kt @@ -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 { diff --git a/apps/server/src/main/kotlin/com/correx/apps/server/bridge/SessionEventBridge.kt b/apps/server/src/main/kotlin/com/correx/apps/server/bridge/SessionEventBridge.kt index 67476699..758452d7 100644 --- a/apps/server/src/main/kotlin/com/correx/apps/server/bridge/SessionEventBridge.kt +++ b/apps/server/src/main/kotlin/com/correx/apps/server/bridge/SessionEventBridge.kt @@ -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): List { + val orderedIds = mutableListOf() + val byInvocation = mutableMapOf() + + 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 -> diff --git a/apps/server/src/main/kotlin/com/correx/apps/server/protocol/Dtos.kt b/apps/server/src/main/kotlin/com/correx/apps/server/protocol/Dtos.kt index 1bf76b6e..4eb69a3d 100644 --- a/apps/server/src/main/kotlin/com/correx/apps/server/protocol/Dtos.kt +++ b/apps/server/src/main/kotlin/com/correx/apps/server/protocol/Dtos.kt @@ -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, diff --git a/apps/server/src/main/kotlin/com/correx/apps/server/protocol/ServerMessage.kt b/apps/server/src/main/kotlin/com/correx/apps/server/protocol/ServerMessage.kt index 7ae50111..53a60b6f 100644 --- a/apps/server/src/main/kotlin/com/correx/apps/server/protocol/ServerMessage.kt +++ b/apps/server/src/main/kotlin/com/correx/apps/server/protocol/ServerMessage.kt @@ -88,6 +88,7 @@ sealed interface ServerMessage { val workflowId: String, val state: SessionStateDto, val pendingApprovals: List, + val tools: List = emptyList(), val recentEvents: List = emptyList(), val lastSequence: Long, val lastSessionSequence: Long, diff --git a/apps/server/src/main/kotlin/com/correx/apps/server/routes/SessionRoutes.kt b/apps/server/src/main/kotlin/com/correx/apps/server/routes/SessionRoutes.kt index 17af7e3c..ddcc1b86 100644 --- a/apps/server/src/main/kotlin/com/correx/apps/server/routes/SessionRoutes.kt +++ b/apps/server/src/main/kotlin/com/correx/apps/server/routes/SessionRoutes.kt @@ -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)) } } diff --git a/apps/server/src/main/kotlin/com/correx/apps/server/ws/GlobalStreamHandler.kt b/apps/server/src/main/kotlin/com/correx/apps/server/ws/GlobalStreamHandler.kt index 50dbddfb..45fc24d1 100644 --- a/apps/server/src/main/kotlin/com/correx/apps/server/ws/GlobalStreamHandler.kt +++ b/apps/server/src/main/kotlin/com/correx/apps/server/ws/GlobalStreamHandler.kt @@ -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) } } diff --git a/apps/tui/src/main/kotlin/com/correx/apps/tui/components/ApprovalSurface.kt b/apps/tui/src/main/kotlin/com/correx/apps/tui/components/ApprovalSurface.kt index 1096e974..9deb46d1 100644 --- a/apps/tui/src/main/kotlin/com/correx/apps/tui/components/ApprovalSurface.kt +++ b/apps/tui/src/main/kotlin/com/correx/apps/tui/components/ApprovalSurface.kt @@ -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 diff --git a/apps/tui/src/main/kotlin/com/correx/apps/tui/reducer/SessionsReducer.kt b/apps/tui/src/main/kotlin/com/correx/apps/tui/reducer/SessionsReducer.kt index 606e4974..4dcc4ac0 100644 --- a/apps/tui/src/main/kotlin/com/correx/apps/tui/reducer/SessionsReducer.kt +++ b/apps/tui/src/main/kotlin/com/correx/apps/tui/reducer/SessionsReducer.kt @@ -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( 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 9b5dbc75..fb572e4c 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 @@ -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() @@ -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? { + 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 diff --git a/testing/deterministic/src/test/kotlin/RouterContextBuilderTest.kt b/testing/deterministic/src/test/kotlin/RouterContextBuilderTest.kt index 41ea0ce6..3817cab8 100644 --- a/testing/deterministic/src/test/kotlin/RouterContextBuilderTest.kt +++ b/testing/deterministic/src/test/kotlin/RouterContextBuilderTest.kt @@ -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]!! diff --git a/testing/deterministic/src/test/kotlin/RouterFacadeTest.kt b/testing/deterministic/src/test/kotlin/RouterFacadeTest.kt index 161c156a..4432ade1 100644 --- a/testing/deterministic/src/test/kotlin/RouterFacadeTest.kt +++ b/testing/deterministic/src/test/kotlin/RouterFacadeTest.kt @@ -247,7 +247,7 @@ class RouterFacadeTest { fun `budget is passed through to context builder`(): Unit = runBlocking { val capturedBudget = mutableListOf() 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() }