Kilo-Org · markijbema · Jun 17, 2026 · Jun 17, 2026 · Jun 17, 2026 · Jun 17, 2026
diff --git a/...s/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/SessionUiTestBase.kt b/...s/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/SessionUiTestBase.kt
@@ -12,6 +12,7 @@ import ai.kilocode.client.session.controller.SessionController
 import ai.kilocode.client.testing.FakeAppRpcApi
 import ai.kilocode.client.testing.FakeSessionRpcApi
 import ai.kilocode.client.testing.FakeWorkspaceRpcApi
+import ai.kilocode.client.testing.TestCoroutines
 import ai.kilocode.client.session.SessionRef
 import ai.kilocode.client.session.scroll.SessionScroll
 import ai.kilocode.rpc.dto.ChatEventDto
@@ -29,8 +30,6 @@ import com.intellij.testFramework.fixtures.BasePlatformTestCase
 import com.intellij.openapi.util.Disposer
 import com.intellij.util.ui.UIUtil
 import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.SupervisorJob
-import kotlinx.coroutines.cancel
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.runBlocking
 import java.awt.Container
@@ -42,6 +41,7 @@ import javax.swing.JScrollBar
 
 @Suppress("UnstableApiUsage")
 abstract class SessionUiTestBase : BasePlatformTestCase() {
+    private lateinit var coroutines: TestCoroutines
     protected lateinit var scope: CoroutineScope
     protected lateinit var sessions: KiloSessionService
     protected lateinit var app: KiloAppService
@@ -53,7 +53,8 @@ abstract class SessionUiTestBase : BasePlatformTestCase() {
 
     override fun setUp() {
         super.setUp()
-        scope = CoroutineScope(SupervisorJob())
+        coroutines = TestCoroutines()
+        scope = coroutines.scope
 
         rpc = FakeSessionRpcApi()
         appRpc = FakeAppRpcApi().also {
@@ -75,7 +76,7 @@ abstract class SessionUiTestBase : BasePlatformTestCase() {
     override fun tearDown() {
         try {
             Disposer.dispose(ui)
-            scope.cancel()
+            coroutines.close { UIUtil.dispatchAllInvocationEvents() }
         } finally {
             super.tearDown()
         }
@@ -116,11 +117,8 @@ abstract class SessionUiTestBase : BasePlatformTestCase() {
         (scrollView() as? Container)?.doLayout()
     }
 
-    protected fun settle() = runBlocking {
-        repeat(5) {
-            delay(100)
-            UIUtil.dispatchAllInvocationEvents()
-        }
+    protected fun settle() {
+        coroutines.drain { UIUtil.dispatchAllInvocationEvents() }
     }
 
     protected fun settleShort(ms: Long) = runBlocking {
@@ -141,15 +139,15 @@ abstract class SessionUiTestBase : BasePlatformTestCase() {
             emit(ChatEventDto.MessageUpdated("ses_test", message(id)), flush = false)
             emit(ChatEventDto.PartUpdated("ses_test", part("part_$i", id, "text", text(i))), flush = false)
         }
-        settleShort(100)
+        settle()
         forceFlush()
         drainScroll()
     }
 
     protected fun emit(event: ChatEventDto, flush: Boolean = true) {
         runBlocking { rpc.events.emit(event) }
         if (flush) {
-            settleShort(20)
+            settle()
             forceFlush()
         }
     }

diff --git a/...ontend/src/test/kotlin/ai/kilocode/client/session/controller/SessionControllerTestBase.kt b/...ontend/src/test/kotlin/ai/kilocode/client/session/controller/SessionControllerTestBase.kt
@@ -8,6 +8,7 @@ import ai.kilocode.client.session.model.SessionState
 import ai.kilocode.client.testing.FakeAppRpcApi
 import ai.kilocode.client.testing.FakeWorkspaceRpcApi
 import ai.kilocode.client.testing.FakeSessionRpcApi
+import ai.kilocode.client.testing.TestCoroutines
 import ai.kilocode.client.testing.TestUiTimers
 import ai.kilocode.client.app.KiloWorkspaceService
 import ai.kilocode.client.app.Workspace
@@ -36,8 +37,6 @@ import com.intellij.testFramework.fixtures.BasePlatformTestCase
 import com.intellij.util.ui.UIUtil
 import java.awt.event.HierarchyEvent
 import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.SupervisorJob
-import kotlinx.coroutines.cancel
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.runBlocking
 
@@ -98,6 +97,7 @@ abstract class SessionControllerTestBase : BasePlatformTestCase() {
     protected lateinit var workspace: Workspace
     protected lateinit var timers: TestUiTimers
 
+    private lateinit var coroutines: TestCoroutines
     protected lateinit var scope: CoroutineScope
     protected lateinit var parent: Disposable
 
@@ -108,7 +108,8 @@ abstract class SessionControllerTestBase : BasePlatformTestCase() {
         projectRpc = FakeWorkspaceRpcApi()
         timers = TestUiTimers()
 
-        scope = CoroutineScope(SupervisorJob())
+        coroutines = TestCoroutines()
+        scope = coroutines.scope
         parent = Disposer.newDisposable("test")
 
         sessions = KiloSessionService(project, scope, rpc)
@@ -120,7 +121,7 @@ abstract class SessionControllerTestBase : BasePlatformTestCase() {
     override fun tearDown() {
         try {
             Disposer.dispose(parent)
-            scope.cancel()
+            coroutines.close { edt { UIUtil.dispatchAllInvocationEvents() } }
         } finally {
             super.tearDown()
         }
@@ -226,23 +227,14 @@ abstract class SessionControllerTestBase : BasePlatformTestCase() {
 
     // ------ EDT + coroutine helpers ------
 
-    /** Let coroutines settle without forcing buffered controller delivery. */
-    protected fun settle() = runBlocking {
-        repeat(5) {
-            delay(100)
-            edt { UIUtil.dispatchAllInvocationEvents() }
-        }
+    /** Drain background coroutine and EDT work without forcing buffered controller delivery. */
+    protected fun settle() {
+        drain(false)
     }
 
-    /** Let coroutines settle, force buffered controller delivery, then drain EDT. */
-    protected fun flush() = runBlocking {
-        repeat(5) {
-            delay(100)
-            edt {
-                controllers.forEach { it.flushEvents() }
-                UIUtil.dispatchAllInvocationEvents()
-            }
-        }
+    /** Drain background work, force buffered controller delivery, then drain EDT. */
+    protected fun flush() {
+        drain(true)
     }
 
     protected fun pause(ms: Long) = runBlocking {
@@ -258,6 +250,15 @@ abstract class SessionControllerTestBase : BasePlatformTestCase() {
         }
     }
 
+    private fun drain(force: Boolean) {
+        coroutines.drain {
+            edt {
+                if (force) controllers.forEach { it.flushEvents() }
+                UIUtil.dispatchAllInvocationEvents()
+            }
+        }
+    }
+
     protected fun edt(block: () -> Unit) {
         ApplicationManager.getApplication().invokeAndWait(block)
     }

diff --git a/.../frontend/src/test/kotlin/ai/kilocode/client/session/controller/SessionUpdateQueueTest.kt b/.../frontend/src/test/kotlin/ai/kilocode/client/session/controller/SessionUpdateQueueTest.kt
@@ -404,11 +404,10 @@ class SessionUpdateQueueTest : SessionControllerTestBase() {
         val sb = StringBuilder()
         for (token in tokens) {
             sb.append(token)
-            emit(ChatEventDto.PartUpdated("ses_test", part("prt1", "ses_test", "msg1", "text", text = sb.toString())))
-            emit(ChatEventDto.PartDelta("ses_test", "msg1", "prt1", "text", token))
+            emit(ChatEventDto.PartUpdated("ses_test", part("prt1", "ses_test", "msg1", "text", text = sb.toString())), flush = false)
+            emit(ChatEventDto.PartDelta("ses_test", "msg1", "prt1", "text", token), flush = false)
+            flush()
         }
-        settle()
-        flush()
 
         val text = (m.model.message("msg1")!!.parts["prt1"] as ai.kilocode.client.session.model.Text).content.toString()
         assertEquals(sb.toString(), text)

diff --git a/...s/frontend/src/test/kotlin/ai/kilocode/client/session/ui/header/SessionHeaderPanelTest.kt b/...s/frontend/src/test/kotlin/ai/kilocode/client/session/ui/header/SessionHeaderPanelTest.kt
@@ -433,8 +433,9 @@ class SessionHeaderPanelTest : SessionControllerTestBase() {
         val c = promptedHeader()
         val panel = SessionHeaderPanel(c, parent)
         repeat(12) { idx ->
-            emit(ChatEventDto.PartUpdated("ses_test", tool("tool_more_$idx", "bash", "running", "More $idx")))
+            emit(ChatEventDto.PartUpdated("ses_test", tool("tool_more_$idx", "bash", "running", "More $idx")), flush = false)
         }
+        flush()
         panel.timelineViewport().setSize(panel.timelineBarWidth() * 4, panel.timelineViewportPreferredSize().height)
         panel.timelineViewport().doLayout()
         panel.timelineViewport().viewPosition = Point(0, 0)
@@ -471,8 +472,9 @@ class SessionHeaderPanelTest : SessionControllerTestBase() {
         val c = promptedHeader()
         val panel = SessionHeaderPanel(c, parent)
         repeat(12) { idx ->
-            emit(ChatEventDto.PartUpdated("ses_test", tool("tool_touch_$idx", "bash", "running", "Touch $idx")))
+            emit(ChatEventDto.PartUpdated("ses_test", tool("tool_touch_$idx", "bash", "running", "Touch $idx")), flush = false)
         }
+        flush()
         panel.timelineViewport().setSize(panel.timelineBarWidth() * 4, panel.timelineViewportPreferredSize().height)
         panel.timelineViewport().doLayout()
         panel.timelineViewport().viewPosition = Point(0, 0)
@@ -505,8 +507,9 @@ class SessionHeaderPanelTest : SessionControllerTestBase() {
         val c = promptedHeader()
         val panel = SessionHeaderPanel(c, parent)
         repeat(12) { idx ->
-            emit(ChatEventDto.PartUpdated("ses_test", tool("tool_wheel_$idx", "bash", "running", "Wheel $idx")))
+            emit(ChatEventDto.PartUpdated("ses_test", tool("tool_wheel_$idx", "bash", "running", "Wheel $idx")), flush = false)
         }
+        flush()
         panel.timelineViewport().setSize(panel.timelineBarWidth() * 4, panel.timelineViewportPreferredSize().height)
         panel.timelineViewport().doLayout()
         panel.timelineViewport().viewPosition = Point(0, 0)
@@ -539,8 +542,9 @@ class SessionHeaderPanelTest : SessionControllerTestBase() {
         val c = promptedHeader()
         val panel = SessionHeaderPanel(c, parent)
         repeat(12) { idx ->
-            emit(ChatEventDto.PartUpdated("ses_test", tool("tool_more_$idx", "bash", "running", "More $idx")))
+            emit(ChatEventDto.PartUpdated("ses_test", tool("tool_more_$idx", "bash", "running", "More $idx")), flush = false)
         }
+        flush()
         panel.timelineViewport().setSize(panel.timelineBarWidth() * 4, panel.timelineViewportPreferredSize().height)
         panel.timelineViewport().doLayout()
         panel.timelineViewport().viewPosition = Point(0, 0)
@@ -580,16 +584,17 @@ class SessionHeaderPanelTest : SessionControllerTestBase() {
         edt { c.prompt("go") }
         flush()
 
-        emit(ChatEventDto.SessionUpdated("ses_test", session("ses_test", title = "Generated title")))
-        emit(ChatEventDto.MessageUpdated("ses_test", assistant()))
-        emit(ChatEventDto.PartUpdated("ses_test", reasoning(done = false, text = "Thinking")))
-        emit(ChatEventDto.PartUpdated("ses_test", tool("tool_1", "bash", "running", "Run tests", input = mapOf("cmd" to "test", "files" to "src"))))
-        emit(ChatEventDto.PartUpdated("ses_test", tool("tool_2", "edit", "error", "Edit file", input = mapOf("cmd" to "test", "files" to "src"))))
-        emit(ChatEventDto.PartUpdated("ses_test", stepFinish()))
+        emit(ChatEventDto.SessionUpdated("ses_test", session("ses_test", title = "Generated title")), flush = false)
+        emit(ChatEventDto.MessageUpdated("ses_test", assistant()), flush = false)
+        emit(ChatEventDto.PartUpdated("ses_test", reasoning(done = false, text = "Thinking")), flush = false)
+        emit(ChatEventDto.PartUpdated("ses_test", tool("tool_1", "bash", "running", "Run tests", input = mapOf("cmd" to "test", "files" to "src"))), flush = false)
+        emit(ChatEventDto.PartUpdated("ses_test", tool("tool_2", "edit", "error", "Edit file", input = mapOf("cmd" to "test", "files" to "src"))), flush = false)
+        emit(ChatEventDto.PartUpdated("ses_test", stepFinish()), flush = false)
         emit(ChatEventDto.TodoUpdated("ses_test", listOf(
             TodoDto("Write tests", "completed", "high"),
             TodoDto("Ship it", "pending", "medium"),
-        )))
+        )), flush = false)
+        flush()
         return c
     }
 

diff --git a/...ages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/testing/TestCoroutines.kt b/...ages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/testing/TestCoroutines.kt
@@ -0,0 +1,39 @@
+package ai.kilocode.client.testing
+
+import java.util.concurrent.Executors
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.SupervisorJob
+import kotlinx.coroutines.asCoroutineDispatcher
+import kotlinx.coroutines.launch
+
+class TestCoroutines {
+    private val dispatcher = Executors.newSingleThreadExecutor().asCoroutineDispatcher()
+    private val job = SupervisorJob()
+
+    val scope = CoroutineScope(job + dispatcher)
+
+    fun drain(pump: () -> Unit) {
+        repeat(5) {
+            await(scope.launch {}, pump)
+            pump()
+        }
+    }
+
+    fun close(pump: () -> Unit) {
+        job.cancel()
+        try {
+            await(job, pump)
+        } finally {
+            dispatcher.close()
+        }
+    }
+
+    private fun await(job: kotlinx.coroutines.Job, pump: () -> Unit) {
+        val end = System.nanoTime() + 5_000_000_000L
+        while (!job.isCompleted) {
+            check(System.nanoTime() < end) { "Timed out draining test coroutines" }
+            pump()
+            Thread.yield()
+        }
+    }
+}