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.
This commit is contained in:
2026-05-28 15:37:16 +04:00
parent e05532e7b2
commit 57d2237ba0
11 changed files with 230 additions and 82 deletions
@@ -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()
}