diff --git a/packages/kilo-jetbrains/backend/src/main/kotlin/ai/kilocode/backend/app/KiloBackendAppService.kt b/packages/kilo-jetbrains/backend/src/main/kotlin/ai/kilocode/backend/app/KiloBackendAppService.kt index fb0f6d8430e..81492b4799f 100644 --- a/packages/kilo-jetbrains/backend/src/main/kotlin/ai/kilocode/backend/app/KiloBackendAppService.kt +++ b/packages/kilo-jetbrains/backend/src/main/kotlin/ai/kilocode/backend/app/KiloBackendAppService.kt @@ -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 = 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 { - 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 { - 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> { - 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 { - val client = connection.api ?: return emptyList() + val client = connection.appLoadApi ?: return emptyList() return try { client.configWarnings().map(::warning) } catch (e: Exception) { diff --git a/packages/kilo-jetbrains/backend/src/main/kotlin/ai/kilocode/backend/app/KiloBackendConnectionService.kt b/packages/kilo-jetbrains/backend/src/main/kotlin/ai/kilocode/backend/app/KiloBackendConnectionService.kt index d9a891cbfce..4eccd0b4c74 100644 --- a/packages/kilo-jetbrains/backend/src/main/kotlin/ai/kilocode/backend/app/KiloBackendConnectionService.kt +++ b/packages/kilo-jetbrains/backend/src/main/kotlin/ai/kilocode/backend/app/KiloBackendConnectionService.kt @@ -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 } diff --git a/packages/kilo-jetbrains/backend/src/main/kotlin/ai/kilocode/backend/cli/KiloBackendHttpClients.kt b/packages/kilo-jetbrains/backend/src/main/kotlin/ai/kilocode/backend/cli/KiloBackendHttpClients.kt index 3e71cf2bf33..b305c8f4a27 100644 --- a/packages/kilo-jetbrains/backend/src/main/kotlin/ai/kilocode/backend/cli/KiloBackendHttpClients.kt +++ b/packages/kilo-jetbrains/backend/src/main/kotlin/ai/kilocode/backend/cli/KiloBackendHttpClients.kt @@ -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() diff --git a/packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/app/KiloBackendAppServiceTest.kt b/packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/app/KiloBackendAppServiceTest.kt index 6864fb20f0f..f63b25cb5d9 100644 --- a/packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/app/KiloBackendAppServiceTest.kt +++ b/packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/app/KiloBackendAppServiceTest.kt @@ -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(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(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"}""" diff --git a/packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/testing/MockCliServer.kt b/packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/testing/MockCliServer.kt index bdde10d818f..12ff956be7e 100644 --- a/packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/testing/MockCliServer.kt +++ b/packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/testing/MockCliServer.kt @@ -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() @@ -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) diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/actions/ReinstallKiloAction.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/actions/ReinstallKiloAction.kt index 7eab2842c67..c5bd9574cd6 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/actions/ReinstallKiloAction.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/actions/ReinstallKiloAction.kt @@ -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().state.value.status - e.presentation.isEnabled = status != KiloAppStatusDto.CONNECTING && status != KiloAppStatusDto.LOADING + e.presentation.isEnabled = true } } diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/actions/RestartKiloAction.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/actions/RestartKiloAction.kt index 335cf611f0f..1df776223e4 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/actions/RestartKiloAction.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/actions/RestartKiloAction.kt @@ -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().state.value.status - e.presentation.isEnabled = status != KiloAppStatusDto.CONNECTING && status != KiloAppStatusDto.LOADING + e.presentation.isEnabled = true } } diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/actions/KiloRecoveryActionsTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/actions/KiloRecoveryActionsTest.kt new file mode 100644 index 00000000000..f99429a1679 --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/actions/KiloRecoveryActionsTest.kt @@ -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 } + } +} diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/testing/FakeAppRpcApi.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/testing/FakeAppRpcApi.kt index a91623f165b..6ebca7de308 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/testing/FakeAppRpcApi.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/testing/FakeAppRpcApi.kt @@ -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 {