opensource/kilocode
3
0
Fork 0
mirror of https://github.com/Kilo-Org/kilocode.git synced 2026-09-21 05:52:35 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(jetbrains): bound app load recovery

Browse Source
This commit is contained in:
kirillk
2026-05-25 13:17:07 -04:00
parent c81397044c
commit 0851f4024a
9 changed files with 211 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/kilo-jetbrains/backend/src/main/kotlin/ai/kilocode/backend/app/KiloBackendAppService.kt
+26 -9
View File
@@ -34,7 +34,9 @@ import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.launch
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.withContext
import kotlinx.coroutines.withTimeout
import kotlinx.coroutines.sync.withLock
import kotlinx.coroutines.TimeoutCancellationException
import kotlinx.serialization.json.JsonNull
import kotlinx.serialization.json.JsonObject
import kotlinx.serialization.json.JsonPrimitive
@@ -68,6 +70,7 @@ class KiloBackendAppService private constructor(
private val cs: CoroutineScope,
private val server: CliServer,
private val log: KiloLog,
private val loadTimeoutMs: Long,
) : Disposable {
/** IntelliJ service injection entry point. */
@@ -75,24 +78,27 @@ class KiloBackendAppService private constructor(
cs,
KiloBackendCliManager(),
KiloLog.create(KiloBackendAppService::class.java),
APP_LOAD_TIMEOUT_MS,
)
companion object {
private const val MAX_RETRIES = 3
private const val RETRY_DELAY_MS = 1000L
private const val APP_LOAD_TIMEOUT_MS = 30_000L
/** Test factory — no IntelliJ deps needed. */
internal fun create(
cs: CoroutineScope,
server: CliServer,
log: KiloLog,
) = KiloBackendAppService(cs, server, log)
loadTimeoutMs: Long = APP_LOAD_TIMEOUT_MS,
) = KiloBackendAppService(cs, server, log, loadTimeoutMs)
}
private val mutex = Mutex()
private val connection = KiloConnectionService(cs, server, onReconnect = {
cs.launch { reconnect() }
}, log = log)
}, appLoadTimeoutMs = loadTimeoutMs, log = log)
private var watcher: Job? = null
private var eventWatcher: Job? = null
@@ -250,7 +256,8 @@ class KiloBackendAppService private constructor(
var warns: List<ConfigWarning> = emptyList()
try {
coroutineScope {
withTimeout(loadTimeoutMs) {
coroutineScope {
launch {
val result = fetchProfile()
val status = when {
@@ -291,11 +298,11 @@ class KiloBackendAppService private constructor(
throw LoadFailure(err)
}
}
launch {
warns = fetchWarnings()
}
}
warns = fetchWarnings()
ensureActive()
profile = prof
config = cfg
@@ -313,6 +320,16 @@ class KiloBackendAppService private constructor(
)
log.info("Application started — config, profile, notifications loaded")
startWatchingGlobalSseEvents()
} catch (e: TimeoutCancellationException) {
val err = LoadError(
resource = "app",
detail = "Timed out loading app data after ${loadTimeoutMs}ms",
)
log.warn("Application start timed out after ${loadTimeoutMs}ms")
setAppError(
message = "Failed to load required data",
errors = errors.toList() + err,
)
} catch (e: CancellationException) {
throw e
} catch (e: Exception) {
@@ -336,7 +353,7 @@ class KiloBackendAppService private constructor(
* as failures.
*/
private suspend fun fetchProfile(): FetchResult<KiloProfile200Response?> {
val client = connection.api
val client = connection.appLoadApi
?: return FetchResult.ok(null)
return try {
val response = client.kiloProfile()
@@ -364,7 +381,7 @@ class KiloBackendAppService private constructor(
}
private suspend fun fetchConfig(): FetchResult<Config> {
val client = connection.api
val client = connection.appLoadApi
?: return FetchResult.fail("config", detail = "Not connected")
return try {
FetchResult.ok(client.globalConfigGet())
@@ -376,7 +393,7 @@ class KiloBackendAppService private constructor(
}
private suspend fun fetchNotifications(): FetchResult<List<KiloNotifications200ResponseInner>> {
val client = connection.api
val client = connection.appLoadApi
?: return FetchResult.fail("notifications", detail = "Not connected")
return try {
FetchResult.ok(client.kiloNotifications())
@@ -388,7 +405,7 @@ class KiloBackendAppService private constructor(
}
private suspend fun fetchWarnings(): List<ConfigWarning> {
val client = connection.api ?: return emptyList()
val client = connection.appLoadApi ?: return emptyList()
return try {
client.configWarnings().map(::warning)
} catch (e: Exception) {
packages/kilo-jetbrains/backend/src/main/kotlin/ai/kilocode/backend/app/KiloBackendConnectionService.kt
+25 -1
View File
@@ -44,6 +44,7 @@ data class SseEvent(val type: String, val data: String)
*
* Uses two separate OkHttp clients mirroring the VS Code architecture:
* - [apiClient]: no call/read timeout — used for the generated API client and SSE
* - app-load client: bounded timeout — used for startup REST calls
* - [healthClient]: 3 s timeout — used only for `/global/health` polling
*
* The generated [DefaultApi] is configured with [apiClient] and exposed via [api]
@@ -59,9 +60,23 @@ class KiloConnectionService(
private val cs: CoroutineScope,
private val server: CliServer,
private val onReconnect: () -> Unit,
private val log: KiloLog = KiloLog.create(KiloConnectionService::class.java),
private val log: KiloLog,
private val appLoadTimeoutMs: Long,
) {
constructor(
cs: CoroutineScope,
server: CliServer,
onReconnect: () -> Unit,
) : this(cs, server, onReconnect, KiloLog.create(KiloConnectionService::class.java), 30_000L)
constructor(
cs: CoroutineScope,
server: CliServer,
onReconnect: () -> Unit,
log: KiloLog,
) : this(cs, server, onReconnect, log, 30_000L)
companion object {
private const val HEARTBEAT_TIMEOUT_MS = 15_000L
private const val HEALTH_POLL_INTERVAL_MS = 10_000L
@@ -81,6 +96,9 @@ class KiloConnectionService(
/** OkHttp client used for API calls — no call/read timeout. Null when disconnected. */
var apiClient: OkHttpClient? = null
private set
var appLoadApi: DefaultApi? = null
private set
private var appLoadClient: OkHttpClient? = null
private var healthClient: OkHttpClient? = null
/** Port the CLI server is listening on. Zero when disconnected. */
var port = 0
@@ -175,12 +193,15 @@ class KiloConnectionService(
// Create dual OkHttp clients (bundled — no IntelliJ platform deps)
val ac = KiloBackendHttpClients.api(password)
val lc = KiloBackendHttpClients.appLoad(password, appLoadTimeoutMs)
val hc = KiloBackendHttpClients.health(password)
apiClient = ac
appLoadClient = lc
healthClient = hc
// Configure generated API client with the no-timeout api client
api = DefaultApi(basePath = "http://127.0.0.1:$port", client = ac)
appLoadApi = DefaultApi(basePath = "http://127.0.0.1:$port", client = lc)
startSse()
startHeartbeatWatcher()
@@ -327,8 +348,11 @@ class KiloConnectionService(
private fun close() {
api = null
appLoadApi = null
apiClient?.let { KiloBackendHttpClients.shutdown(it) }
apiClient = null
appLoadClient?.let { KiloBackendHttpClients.shutdown(it) }
appLoadClient = null
healthClient?.let { KiloBackendHttpClients.shutdown(it) }
healthClient = null
}
packages/kilo-jetbrains/backend/src/main/kotlin/ai/kilocode/backend/cli/KiloBackendHttpClients.kt
+14 -1
View File
@@ -7,10 +7,11 @@ import java.util.Base64
import java.util.concurrent.TimeUnit
/**
* Factory for the two OkHttp clients used by the plugin.
* Factory for the OkHttp clients used by the plugin.
*
* Mirrors the VS Code architecture:
* - [api] client has no call/read timeout (streaming ops like prompt/SSE can run long)
* - [appLoad] client has a bounded timeout for startup REST calls
* - [health] client has a short 3 s timeout and a small dedicated connection pool
*
* Both clients bundle Basic Auth via an interceptor and are fully independent
@@ -30,6 +31,18 @@ object KiloBackendHttpClients {
.readTimeout(0, TimeUnit.MILLISECONDS)
.build()
/** App-load client — bounded timeout for required startup REST calls. */
fun appLoad(password: String, timeoutMs: Long): OkHttpClient {
val timeout = timeoutMs.coerceAtLeast(1L)
return OkHttpClient.Builder()
.addInterceptor(auth(password))
.connectTimeout(CONNECT_TIMEOUT_MS.coerceAtMost(timeout), TimeUnit.MILLISECONDS)
.callTimeout(timeout, TimeUnit.MILLISECONDS)
.readTimeout(timeout, TimeUnit.MILLISECONDS)
.connectionPool(ConnectionPool(2, 30, TimeUnit.SECONDS))
.build()
}
/** Health client — short timeout, dedicated connection pool. */
fun health(password: String): OkHttpClient =
OkHttpClient.Builder()
packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/app/KiloBackendAppServiceTest.kt
+100 -2
View File
@@ -37,8 +37,8 @@ class KiloBackendAppServiceTest {
mock.close()
}
private fun create(): KiloBackendAppService =
KiloBackendAppService.create(scope, FakeCliServer(mock), log)
private fun create(loadTimeoutMs: Long = 30_000L): KiloBackendAppService =
KiloBackendAppService.create(scope, FakeCliServer(mock), log, loadTimeoutMs)
@Test
fun `full lifecycle reaches Ready`() = runBlocking {
@@ -404,6 +404,104 @@ class KiloBackendAppServiceTest {
}
}
@Test
fun `hung app load transitions from Loading to Error`() = runBlocking {
val gate = CountDownLatch(1)
mock.responseGate = gate
val svc = create(loadTimeoutMs = 300L)
try {
svc.connect()
withTimeout(10_000) {
svc.appState.first { it is KiloAppState.Loading }
}
val err = withTimeout(10_000) {
svc.appState.first { it is KiloAppState.Error }
} as KiloAppState.Error
assertEquals("Failed to load required data", err.message)
assertTrue(err.errors.any { it.detail?.contains("timeout", ignoreCase = true) == true })
} finally {
gate.countDown()
}
}
@Test
fun `hung warnings do not prevent Ready`() = runBlocking {
val gate = CountDownLatch(1)
mock.warningsGate = gate
val svc = create(loadTimeoutMs = 300L)
try {
svc.connect()
val ready = withTimeout(10_000) {
svc.appState.first { it is KiloAppState.Ready }
} as KiloAppState.Ready
assertTrue(ready.data.warnings.isEmpty())
assertTrue(svc.warnings.isEmpty())
} finally {
gate.countDown()
}
}
@Test
fun `restart during Loading cancels stale load and reaches Ready`() = runBlocking {
val gate = CountDownLatch(1)
mock.responseGate = gate
val svc = create(loadTimeoutMs = 500L)
try {
svc.connect()
withTimeout(10_000) {
svc.appState.first { it is KiloAppState.Loading }
}
gate.countDown()
svc.restart()
withTimeout(10_000) {
svc.appState.first { it is KiloAppState.Ready }
}
assertIs<KiloAppState.Ready>(svc.appState.value)
assertFalse(log.messages.any { it.contains("Application start timed out") })
} finally {
gate.countDown()
}
}
@Test
fun `reinstall during Loading cancels stale load and reaches Ready`() = runBlocking {
val gate = CountDownLatch(1)
mock.responseGate = gate
val svc = create(loadTimeoutMs = 500L)
try {
svc.connect()
withTimeout(10_000) {
svc.appState.first { it is KiloAppState.Loading }
}
gate.countDown()
svc.reinstall()
withTimeout(10_000) {
svc.appState.first { it is KiloAppState.Ready }
}
assertIs<KiloAppState.Ready>(svc.appState.value)
assertFalse(log.messages.any { it.contains("Application start timed out") })
} finally {
gate.countDown()
}
}
@Test
fun `SSE config updated event refreshes config`() = runBlocking {
mock.config = """{"model":"initial"}"""
packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/testing/MockCliServer.kt
+4
View File
@@ -94,6 +94,9 @@ class MockCliServer : AutoCloseable {
/** Optional gate for REST responses; SSE stays unblocked so the app can enter Loading. */
@Volatile var responseGate: CountDownLatch? = null
/** Optional gate for config warnings only. */
@Volatile var warningsGate: CountDownLatch? = null
/** Request counts by bare path (e.g. "/session" or "/global/config"). Thread-safe. */
private val counts = ConcurrentHashMap<String, AtomicInteger>()
@@ -223,6 +226,7 @@ class MockCliServer : AutoCloseable {
val delay = responseDelay
if (delay > 0) Thread.sleep(delay)
if (bare != "/global/event") responseGate?.await()
if (bare.startsWith("/config/warnings")) warningsGate?.await()
when {
path == "/global/health" -> respond(output, 200, health)
packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/actions/ReinstallKiloAction.kt
+1 -3
View File
@@ -1,7 +1,6 @@
package ai.kilocode.client.actions
import ai.kilocode.client.app.KiloAppService
import ai.kilocode.rpc.dto.KiloAppStatusDto
import com.intellij.openapi.actionSystem.AnAction
import com.intellij.openapi.actionSystem.AnActionEvent
import com.intellij.openapi.components.service
@@ -13,7 +12,6 @@ class ReinstallKiloAction : AnAction(), DumbAware {
}
override fun update(e: AnActionEvent) {
val status = service<KiloAppService>().state.value.status
e.presentation.isEnabled = status != KiloAppStatusDto.CONNECTING && status != KiloAppStatusDto.LOADING
e.presentation.isEnabled = true
}
}
packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/actions/RestartKiloAction.kt
+1 -3
View File
@@ -1,7 +1,6 @@
package ai.kilocode.client.actions
import ai.kilocode.client.app.KiloAppService
import ai.kilocode.rpc.dto.KiloAppStatusDto
import com.intellij.openapi.actionSystem.AnAction
import com.intellij.openapi.actionSystem.AnActionEvent
import com.intellij.openapi.components.service
@@ -13,7 +12,6 @@ class RestartKiloAction : AnAction(), DumbAware {
}
override fun update(e: AnActionEvent) {
val status = service<KiloAppService>().state.value.status
e.presentation.isEnabled = status != KiloAppStatusDto.CONNECTING && status != KiloAppStatusDto.LOADING
e.presentation.isEnabled = true
}
}
packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/actions/KiloRecoveryActionsTest.kt
+34
View File
@@ -0,0 +1,34 @@
package ai.kilocode.client.actions
import com.intellij.openapi.actionSystem.AnAction
import com.intellij.openapi.actionSystem.AnActionEvent
import com.intellij.openapi.actionSystem.Presentation
import com.intellij.openapi.actionSystem.ex.ActionUtil
import com.intellij.testFramework.fixtures.BasePlatformTestCase
@Suppress("UnstableApiUsage")
class KiloRecoveryActionsTest : BasePlatformTestCase() {
fun `test restart action stays enabled for all app states`() {
val action = RestartKiloAction()
val event = event(action)
ActionUtil.updateAction(action, event)
assertTrue("Restart should force-enable recovery action", event.presentation.isEnabled)
}
fun `test reinstall action stays enabled for all app states`() {
val action = ReinstallKiloAction()
val event = event(action)
ActionUtil.updateAction(action, event)
assertTrue("Reinstall should force-enable recovery action", event.presentation.isEnabled)
}
private fun event(action: AnAction): AnActionEvent {
val presentation = Presentation().apply { copyFrom(action.templatePresentation) }
presentation.isEnabled = false
return AnActionEvent.createFromDataContext("", presentation) { null }
}
}
packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/testing/FakeAppRpcApi.kt
+6
View File
@@ -35,6 +35,10 @@ class FakeAppRpcApi : KiloAppRpcApi {
private set
var retries = 0
private set
var restarts = 0
private set
var reinstalls = 0
private set
override suspend fun connect() {
assertNotEdt("connect")
@@ -58,10 +62,12 @@ class FakeAppRpcApi : KiloAppRpcApi {
override suspend fun restart() {
assertNotEdt("restart")
restarts += 1
}
override suspend fun reinstall() {
assertNotEdt("reinstall")
reinstalls += 1
}
override suspend fun modelState(): ModelStateDto {