Merge pull request #10890 from Kilo-Org/thoracic-plume

test(jetbrains): stabilize backend SSE tests

packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/app/KiloBackendAppServiceTest.kt

@@ -285,10 +285,12 @@ class KiloBackendAppServiceTest {
         svc.connect()
 
         withTimeout(10_000) {
-            svc.appState.first { it is KiloAppState.Ready }
+            svc.appState.first { state ->
+                state is KiloAppState.Ready && state.data.warnings.any { it.path == ".kilo/kilo.json" }
+            }
         }
 
-        assertTrue(log.messages.any {
+        assertTrue(log.awaitMessage {
             it.contains("App warnings:") && it.contains(".kilo/kilo.json: Invalid JSON")
         })
     }
@@ -516,13 +518,14 @@ class KiloBackendAppServiceTest {
 
         // Change the config response and push an SSE event
         mock.config = """{"model":"updated"}"""
+        val before = mock.requestCount("/global/config")
         mock.awaitSseConnection()
         mock.pushEvent("global.config.updated", """{"type":"global.config.updated"}""")
 
-        // Wait for config to be refreshed
+        assertTrue(mock.awaitRequestCount("/global/config", before + 1))
         withTimeout(5_000) {
-            while (svc.config?.model != "updated") {
-                delay(100)
+            svc.appState.first { state ->
+                state is KiloAppState.Ready && state.data.config.model == "updated"
             }
         }
 
@@ -542,12 +545,14 @@ class KiloBackendAppServiceTest {
         assertEquals(1, (svc.appState.value as KiloAppState.Ready).data.warnings.size)
 
         mock.warnings = "[]"
+        val before = mock.requestCount("/config/warnings")
         mock.awaitSseConnection()
         mock.pushEvent("global.config.updated", """{"type":"global.config.updated"}""")
 
+        assertTrue(mock.awaitRequestCount("/config/warnings", before + 1))
         withTimeout(5_000) {
-            while ((svc.appState.value as? KiloAppState.Ready)?.data?.warnings?.isNotEmpty() == true) {
-                delay(100)
+            svc.appState.first { state ->
+                state is KiloAppState.Ready && state.data.warnings.isEmpty()
             }
         }
 

packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/app/KiloConnectionServiceTest.kt

@@ -8,11 +8,11 @@ import ai.kilocode.backend.testing.MockCliServer
 import ai.kilocode.backend.testing.TestLog
 import ai.kilocode.log.KiloLog
 import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.CoroutineStart
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.SupervisorJob
 import kotlinx.coroutines.async
 import kotlinx.coroutines.cancel
-import kotlinx.coroutines.delay
 import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.flow.take
 import kotlinx.coroutines.flow.toList
@@ -89,13 +89,11 @@ class KiloConnectionServiceTest {
         }
 
         // Use first{} on the flow to capture the event — avoids race with SharedFlow subscription
-        val deferred = scope.launch {
+        val deferred = scope.launch(start = CoroutineStart.UNDISPATCHED) {
             val event = svc.events.first { it.type == "global.config.updated" }
             assertEquals("global.config.updated", event.type)
         }
 
-        // Small delay to ensure the collector subscription is active
-        delay(200)
         mock.pushEvent("global.config.updated", """{"type":"global.config.updated"}""")
 
         withTimeout(5_000) {
@@ -114,13 +112,12 @@ class KiloConnectionServiceTest {
         }
 
         val count = 100
-        val received = async {
+        val received = async(start = CoroutineStart.UNDISPATCHED) {
             withTimeout(5_000) {
                 svc.events.take(count).toList()
             }
         }
 
-        delay(200)
         repeat(count) { idx ->
             mock.pushEvent("test.event", idx.toString())
         }
@@ -144,18 +141,23 @@ class KiloConnectionServiceTest {
         @Suppress("UNCHECKED_CAST")
         val ref = sourceField.get(svc) as AtomicReference<EventSource?>
         ref.set(source)
-        val received = scope.async {
+        val received = scope.async(start = CoroutineStart.UNDISPATCHED) {
             withTimeout(5_000) {
                 svc.events.take(2).toList()
             }
         }
 
-        delay(200)
+        val entered = CountDownLatch(1)
         val first = Thread { listener.onEvent(source, null, "first.event", "first") }
-        val second = Thread { listener.onEvent(source, null, "second.event", "second") }
+        val second = Thread {
+            entered.countDown()
+            listener.onEvent(source, null, "second.event", "second")
+        }
         first.start()
         assertTrue(blocked.started.await(1, TimeUnit.SECONDS))
         second.start()
+        assertTrue(entered.await(1, TimeUnit.SECONDS))
+        blocked.release.countDown()
         first.join()
         second.join()
 
@@ -247,9 +249,11 @@ class KiloConnectionServiceTest {
         }
 
         // Restart should tear down and re-open
+        val seen = mock.sseConnectionCount
         svc.restart()
 
         // Wait for reconnection
+        assertTrue(mock.awaitSseConnections(seen + 1))
         withTimeout(10_000) {
             svc.state.first { it is ConnectionState.Connected }
         }
@@ -281,13 +285,14 @@ class KiloConnectionServiceTest {
 
     private class BlockingLog : KiloLog {
         val started = CountDownLatch(1)
+        val release = CountDownLatch(1)
         override var isDebugEnabled: Boolean = true
 
         override fun debug(block: () -> String) {
             val msg = block()
             if (!msg.contains("evt=first.event")) return
             started.countDown()
-            Thread.sleep(250)
+            release.await(1, TimeUnit.SECONDS)
         }
 
         override fun info(msg: String) {}

packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/testing/MockCliServer.kt

@@ -99,10 +99,40 @@ class MockCliServer : AutoCloseable {
 
     /** Request counts by bare path (e.g. "/session" or "/global/config"). Thread-safe. */
     private val counts = ConcurrentHashMap<String, AtomicInteger>()
+    private val requests = Object()
+    private val streams = Object()
+    private val sse = AtomicInteger(0)
+
+    val sseConnectionCount: Int
+        get() = sse.get()
 
     /** Return the number of requests received for [path] (bare, no query). */
     fun requestCount(path: String): Int = counts[path]?.get() ?: 0
 
+    fun awaitRequestCount(path: String, target: Int, timeout: Long = 5_000): Boolean {
+        val end = System.currentTimeMillis() + timeout
+        synchronized(requests) {
+            while (requestCount(path) < target) {
+                val wait = end - System.currentTimeMillis()
+                if (wait <= 0) return false
+                requests.wait(wait)
+            }
+            return true
+        }
+    }
+
+    fun awaitSseConnections(target: Int, timeout: Long = 5_000): Boolean {
+        val end = System.currentTimeMillis() + timeout
+        synchronized(streams) {
+            while (sse.get() < target) {
+                val wait = end - System.currentTimeMillis()
+                if (wait <= 0) return false
+                streams.wait(wait)
+            }
+            return true
+        }
+    }
+
     @Volatile var lastExperimentalSessionPath: String? = null
 
     /** Reset all request counters. */
@@ -127,7 +157,8 @@ class MockCliServer : AutoCloseable {
         // Clean up any previous instance
         shutdownServer()
 
-        sseLatch = CountDownLatch(1)
+        val latch = CountDownLatch(1)
+        sseLatch = latch
         sseConnected = CountDownLatch(1)
         sseWriter = null
 
@@ -218,9 +249,11 @@ class MockCliServer : AutoCloseable {
 
             val output = BufferedWriter(OutputStreamWriter(socket.getOutputStream()))
             val bare = path.substringBefore("?")
+            val latch = sseLatch
 
             // Track request counts
             counts.computeIfAbsent(bare) { AtomicInteger(0) }.incrementAndGet()
+            synchronized(requests) { requests.notifyAll() }
 
             // Optional delay for race condition testing
             val delay = responseDelay
@@ -255,7 +288,7 @@ class MockCliServer : AutoCloseable {
                     lastOrganizationSetBody = body
                     respond(output, organizationSetStatus, "true")
                 }
-                path == "/global/event" -> handleSse(output)
+                path == "/global/event" -> handleSse(output, latch)
                 path == "/path" -> respond(output, 200, this.path)
                 bare == "/provider" -> respond(output, providersStatus, providers)
                 bare == "/agent" -> respond(output, agentsStatus, agents)
@@ -319,7 +352,7 @@ class MockCliServer : AutoCloseable {
         writer.flush()
     }
 
-    private fun handleSse(writer: BufferedWriter) {
+    private fun handleSse(writer: BufferedWriter, latch: CountDownLatch) {
         writer.write("HTTP/1.1 200 OK\r\n")
         writer.write("Content-Type: text/event-stream\r\n")
         writer.write("Cache-Control: no-cache\r\n")
@@ -327,8 +360,10 @@ class MockCliServer : AutoCloseable {
         writer.write("\r\n")
         writer.flush()
         sseWriter = writer
+        sse.incrementAndGet()
+        synchronized(streams) { streams.notifyAll() }
         sseConnected.countDown()
         // Block until SSE is closed or server shuts down
-        sseLatch.await()
+        latch.await()
     }
 }

packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/testing/TestLog.kt

@@ -7,31 +7,51 @@ import ai.kilocode.log.KiloLog
  */
 class TestLog : KiloLog {
     private val items = mutableListOf<String>()
+    private val lock = Object()
     val messages: List<String>
-        get() = synchronized(items) { items.toList() }
+        get() = synchronized(lock) { items.toList() }
     override var isDebugEnabled: Boolean = true
 
+    fun awaitMessage(timeout: Long = 5_000, predicate: (String) -> Boolean): Boolean {
+        val end = System.currentTimeMillis() + timeout
+        synchronized(lock) {
+            while (items.none(predicate)) {
+                val wait = end - System.currentTimeMillis()
+                if (wait <= 0) return false
+                lock.wait(wait)
+            }
+            return true
+        }
+    }
+
     override fun debug(block: () -> String) {
         if (!isDebugEnabled) return
         val msg = block()
-        synchronized(items) { items.add("DEBUG: $msg") }
+        add("DEBUG: $msg")
         println("[test] DEBUG: $msg")
     }
 
     override fun info(msg: String) {
-        synchronized(items) { items.add("INFO: $msg") }
+        add("INFO: $msg")
         println("[test] INFO: $msg")
     }
 
     override fun warn(msg: String, t: Throwable?) {
-        synchronized(items) { items.add("WARN: $msg") }
+        add("WARN: $msg")
         println("[test] WARN: $msg")
         t?.printStackTrace()
     }
 
     override fun error(msg: String, t: Throwable?) {
-        synchronized(items) { items.add("ERROR: $msg") }
+        add("ERROR: $msg")
         System.err.println("[test] ERROR: $msg")
         t?.printStackTrace()
     }
+
+    private fun add(msg: String) {
+        synchronized(lock) {
+            items.add(msg)
+            lock.notifyAll()
+        }
+    }
 }