diff --git a/.changeset/jetbrains-restartless-unload.md b/.changeset/jetbrains-restartless-unload.md new file mode 100644 index 00000000000..62e2f7941a7 --- /dev/null +++ b/.changeset/jetbrains-restartless-unload.md @@ -0,0 +1,5 @@ +--- +"@kilocode/kilo-jetbrains": patch +--- + +Stop Kilo backend processes and clear JetBrains UI resources during restartless plugin unloads. 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 bd34fa476b8..7c3ab88b702 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 @@ -110,6 +110,7 @@ class KiloBackendAppService private constructor( private var watcher: Job? = null private var eventWatcher: Job? = null private var loader: Job? = null + private var closed = false private val loadLock = Any() private val _appState = MutableStateFlow(KiloAppState.Disconnected) @@ -159,6 +160,12 @@ class KiloBackendAppService private constructor( } } + suspend fun shutdownForUnload() { + mutex.withLock { + shutdown() + } + } + suspend fun retry() { mutex.withLock { when (val current = _appState.value) { @@ -806,6 +813,12 @@ class KiloBackendAppService private constructor( } override fun dispose() { + shutdown() + } + + private fun shutdown() { + if (closed) return + closed = true watcher?.cancel() watcher = null clear() diff --git a/packages/kilo-jetbrains/backend/src/main/kotlin/ai/kilocode/backend/cli/KiloBackendCliManager.kt b/packages/kilo-jetbrains/backend/src/main/kotlin/ai/kilocode/backend/cli/KiloBackendCliManager.kt index a22af3eee79..5ad06704ee7 100644 --- a/packages/kilo-jetbrains/backend/src/main/kotlin/ai/kilocode/backend/cli/KiloBackendCliManager.kt +++ b/packages/kilo-jetbrains/backend/src/main/kotlin/ai/kilocode/backend/cli/KiloBackendCliManager.kt @@ -40,7 +40,10 @@ class KiloBackendCliManager( @Volatile private var process: Process? = null + @Volatile + private var closing: Process? = null private var hook: Thread? = null + private var stderr: Thread? = null @Volatile override var forceExtract = false @@ -59,8 +62,7 @@ class KiloBackendCliManager( process?.let { proc -> log.info("Cleaning up orphaned CLI process (pid=${proc.pid()})") process = null - uninstall() - kill(proc, "startup failure cleanup") + cleanup(proc, "startup failure cleanup") } CliServer.State.Error( message = e.message ?: "Unknown error", @@ -73,13 +75,13 @@ class KiloBackendCliManager( if (process != proc) return process = null uninstall() + stderr = null } override fun stop() { val proc = process ?: return process = null - uninstall() - kill(proc, "stop()") + cleanup(proc, "stop()") } private fun extractCli(): File { @@ -158,14 +160,19 @@ class KiloBackendCliManager( val stderr = StringBuilder() - Thread({ - BufferedReader(InputStreamReader(proc.errorStream)).use { reader -> - reader.lineSequence().forEach { line -> - log.warn("CLI stderr: $line") - synchronized(stderr) { stderr.appendLine(line) } + val err = Thread({ + runCatching { + BufferedReader(InputStreamReader(proc.errorStream)).use { reader -> + reader.lineSequence().forEach { line -> + log.warn("CLI stderr: $line") + synchronized(stderr) { stderr.appendLine(line) } + } } + }.onFailure { err -> + if (proc.isAlive && closing !== proc) log.warn("CLI stderr reader failed", err) } }, "kilo-cli-stderr").apply { isDaemon = true; start() } + this@KiloBackendCliManager.stderr = err BufferedReader(InputStreamReader(proc.inputStream)).use { reader -> for (line in reader.lineSequence()) { @@ -185,6 +192,7 @@ class KiloBackendCliManager( val details = synchronized(stderr) { stderr.toString().trim() } process = null uninstall() + this@KiloBackendCliManager.stderr = null log.warn("CLI process exited with code $code before announcing a port: $details") CliServer.State.Error( message = "CLI process exited with code $code before announcing a port", @@ -195,8 +203,23 @@ class KiloBackendCliManager( override fun dispose() { val proc = process ?: return process = null - uninstall() - kill(proc, "Disposing") + cleanup(proc, "Disposing") + } + + private fun cleanup(proc: Process, source: String) { + closing = proc + try { + uninstall() + close(proc) + kill(proc, source) + val thread = stderr + stderr = null + if (thread != null && thread != Thread.currentThread()) { + thread.join(TimeUnit.SECONDS.toMillis(1)) + } + } finally { + closing = null + } } private fun install(proc: Process) { @@ -237,6 +260,12 @@ class KiloBackendCliManager( private fun children(proc: Process): List = proc.toHandle().descendants().toList().asReversed() + private fun close(proc: Process) { + runCatching { proc.errorStream.close() }.onFailure { log.info("CLI stderr stream close skipped: ${it.message}") } + runCatching { proc.inputStream.close() }.onFailure { log.info("CLI stdout stream close skipped: ${it.message}") } + runCatching { proc.outputStream.close() }.onFailure { log.info("CLI stdin stream close skipped: ${it.message}") } + } + private fun platform(): String { val os = when { SystemInfo.isMac -> "darwin" diff --git a/packages/kilo-jetbrains/backend/src/main/kotlin/ai/kilocode/backend/plugin/KiloBackendDynamicPluginListener.kt b/packages/kilo-jetbrains/backend/src/main/kotlin/ai/kilocode/backend/plugin/KiloBackendDynamicPluginListener.kt new file mode 100644 index 00000000000..ffcf17f387c --- /dev/null +++ b/packages/kilo-jetbrains/backend/src/main/kotlin/ai/kilocode/backend/plugin/KiloBackendDynamicPluginListener.kt @@ -0,0 +1,21 @@ +package ai.kilocode.backend.plugin + +import ai.kilocode.KiloPlugin +import ai.kilocode.backend.app.KiloBackendAppService +import ai.kilocode.log.KiloLog +import com.intellij.ide.plugins.DynamicPluginListener +import com.intellij.ide.plugins.IdeaPluginDescriptor +import com.intellij.openapi.components.service +import kotlinx.coroutines.runBlocking + +class KiloBackendDynamicPluginListener : DynamicPluginListener { + private val log = KiloLog.create(KiloBackendDynamicPluginListener::class.java) + + override fun beforePluginUnload(pluginDescriptor: IdeaPluginDescriptor, isUpdate: Boolean) { + if (pluginDescriptor.pluginId != KiloPlugin.id) return + log.info("Shutting down Kilo backend for plugin unload (isUpdate=$isUpdate)") + runBlocking { + service().shutdownForUnload() + } + } +} diff --git a/packages/kilo-jetbrains/backend/src/main/resources/kilo.jetbrains.backend.xml b/packages/kilo-jetbrains/backend/src/main/resources/kilo.jetbrains.backend.xml index b5e04928d9f..83a1cc6fab2 100644 --- a/packages/kilo-jetbrains/backend/src/main/resources/kilo.jetbrains.backend.xml +++ b/packages/kilo-jetbrains/backend/src/main/resources/kilo.jetbrains.backend.xml @@ -12,4 +12,9 @@ + + + + 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 c0101b0e7cd..d19eb944e65 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 @@ -54,6 +54,28 @@ class KiloBackendAppServiceTest { assertNotNull(ready.data.notifications) } + @Test + fun `shutdown for unload clears runtime and disposes server once`() = runBlocking { + val server = FakeCliServer(mock) + val svc = KiloBackendAppService.create(scope, server, log) + svc.connect() + + withTimeout(10_000) { + svc.appState.first { it is KiloAppState.Ready } + } + + svc.shutdownForUnload() + svc.shutdownForUnload() + svc.dispose() + + assertEquals(KiloAppState.Disconnected, svc.appState.value) + assertNull(svc.profile) + assertNull(svc.config) + assertTrue(svc.notifications.isEmpty()) + assertTrue(svc.warnings.isEmpty()) + assertEquals(1, server.disposeCount) + } + @Test fun `config is loaded`() = runBlocking { mock.config = """{"model":"claude-4","username":"testuser"}""" diff --git a/packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/testing/FakeCliServer.kt b/packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/testing/FakeCliServer.kt index 885b327be8d..aa29477ff96 100644 --- a/packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/testing/FakeCliServer.kt +++ b/packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/testing/FakeCliServer.kt @@ -13,6 +13,10 @@ import ai.kilocode.backend.cli.CliServer class FakeCliServer(private val mock: MockCliServer) : CliServer { override var forceExtract = false + var stopCount = 0 + private set + var disposeCount = 0 + private set override fun process(): Process? = null @@ -23,11 +27,13 @@ class FakeCliServer(private val mock: MockCliServer) : CliServer { /** Shutdown the server socket but keep the mock alive for restart. */ override fun stop() { + stopCount++ mock.shutdown() } /** Final cleanup. */ override fun dispose() { + disposeCount++ mock.close() } } diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/KiloToolWindowFactory.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/KiloToolWindowFactory.kt index ecbe1f14f7d..b0b640d3ea8 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/KiloToolWindowFactory.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/KiloToolWindowFactory.kt @@ -6,6 +6,7 @@ import ai.kilocode.client.session.SessionSidePanelManager import ai.kilocode.client.telemetry.Telemetry import ai.kilocode.log.KiloLog import com.intellij.openapi.actionSystem.ActionManager +import com.intellij.openapi.components.Service import com.intellij.openapi.components.service import com.intellij.openapi.project.DumbAware import com.intellij.openapi.project.Project @@ -14,7 +15,6 @@ import com.intellij.openapi.wm.ToolWindowFactory import com.intellij.ui.content.ContentFactory import kotlinx.coroutines.CoroutineScope import kotlinx.coroutines.Dispatchers -import kotlinx.coroutines.SupervisorJob import kotlinx.coroutines.launch import kotlinx.coroutines.withContext @@ -27,16 +27,22 @@ import kotlinx.coroutines.withContext * completes. */ class KiloToolWindowFactory : ToolWindowFactory, DumbAware { - - companion object { - private val LOG = KiloLog.create(KiloToolWindowFactory::class.java) + override fun createToolWindowContent(project: Project, toolWindow: ToolWindow) { + project.service().create(toolWindow) } +} - override fun createToolWindowContent(project: Project, toolWindow: ToolWindow) { +private val LOG = KiloLog.create(KiloToolWindowFactory::class.java) + +@Service(Service.Level.PROJECT) +internal class KiloToolWindowSetupService( + private val project: Project, + private val cs: CoroutineScope, +) { + fun create(toolWindow: ToolWindow) { val start = System.currentTimeMillis() try { val workspaces = service() - val cs = CoroutineScope(SupervisorJob()) val hint = project.basePath ?: "" cs.launch { diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/plugin/KiloFrontendDynamicPluginListener.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/plugin/KiloFrontendDynamicPluginListener.kt new file mode 100644 index 00000000000..d721fc7cc2f --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/plugin/KiloFrontendDynamicPluginListener.kt @@ -0,0 +1,51 @@ +package ai.kilocode.client.plugin + +import ai.kilocode.KiloPlugin +import ai.kilocode.client.session.ui.attachment.unregisterAttachmentEditorKind +import ai.kilocode.client.vfs.KiloEditorKindRegistry +import ai.kilocode.client.vfs.KiloVirtualFileSystem +import ai.kilocode.log.KiloLog +import com.intellij.ide.plugins.DynamicPluginListener +import com.intellij.ide.plugins.IdeaPluginDescriptor +import com.intellij.openapi.components.service +import com.intellij.openapi.fileEditor.FileEditorManager +import com.intellij.openapi.project.ProjectManager +import com.intellij.openapi.wm.ToolWindowManager +import javax.swing.SwingUtilities + +class KiloFrontendDynamicPluginListener : DynamicPluginListener { + override fun beforePluginUnload(pluginDescriptor: IdeaPluginDescriptor, isUpdate: Boolean) { + if (pluginDescriptor.pluginId != KiloPlugin.id) return + KiloFrontendUnloadCleanup.cleanup(isUpdate) + } +} + +object KiloFrontendUnloadCleanup { + private val log = KiloLog.create(KiloFrontendUnloadCleanup::class.java) + + fun cleanup(isUpdate: Boolean) { + log.info("Cleaning up Kilo frontend for plugin unload (isUpdate=$isUpdate)") + runEdt { + ProjectManager.getInstance().openProjects.forEach { project -> + if (project.isDisposed) return@forEach + ToolWindowManager.getInstance(project).getToolWindow("Kilo Code") + ?.contentManager + ?.removeAllContents(true) + val editors = FileEditorManager.getInstance(project).openFiles + .filter { it.fileSystem === KiloVirtualFileSystem.getInstance() } + editors.forEach { file -> FileEditorManager.getInstance(project).closeFile(file) } + } + } + unregisterAttachmentEditorKind() + service().clear() + KiloVirtualFileSystem.getInstance().clear() + } + + private fun runEdt(block: () -> Unit) { + if (SwingUtilities.isEventDispatchThread()) { + block() + return + } + SwingUtilities.invokeAndWait(block) + } +} diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/attachment/AttachmentEditorKind.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/attachment/AttachmentEditorKind.kt index 38da71dc4f2..a10409fc32a 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/attachment/AttachmentEditorKind.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/attachment/AttachmentEditorKind.kt @@ -249,6 +249,10 @@ fun ensureAttachmentEditorKind() { service().register(AttachmentEditorKind) } +internal fun unregisterAttachmentEditorKind() { + service().unregister(AttachmentEditorKind.ID) +} + internal fun attachmentParams( sessionId: String, messageId: String, diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/vfs/KiloEditorKindRegistry.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/vfs/KiloEditorKindRegistry.kt index 207e79da489..236019d7f35 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/vfs/KiloEditorKindRegistry.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/vfs/KiloEditorKindRegistry.kt @@ -18,5 +18,9 @@ class KiloEditorKindRegistry { service().unregister(id) } + fun clear() { + kinds.keys.forEach { id -> unregister(id) } + } + fun get(id: String): KiloEditorKind? = kinds[id] } diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/vfs/KiloVirtualFileKindRegistry.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/vfs/KiloVirtualFileKindRegistry.kt index 89af66be9e0..08ad685a9ed 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/vfs/KiloVirtualFileKindRegistry.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/vfs/KiloVirtualFileKindRegistry.kt @@ -15,5 +15,9 @@ class KiloVirtualFileKindRegistry { kinds.remove(id) } + fun clear() { + kinds.clear() + } + fun get(id: String): KiloVirtualFileKind? = kinds[id] } diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/vfs/KiloVirtualFileSystem.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/vfs/KiloVirtualFileSystem.kt index c87aa9e1911..cac74dd041f 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/vfs/KiloVirtualFileSystem.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/vfs/KiloVirtualFileSystem.kt @@ -26,6 +26,10 @@ class KiloVirtualFileSystem : VirtualFileSystem(), NonPhysicalFileSystem { files.remove(path.canonical()) } + fun clear() { + files.clear() + } + override fun findFileByPath(path: String): VirtualFile? { val parsed = decode(path) ?: return null return findOrCreateFile(parsed) diff --git a/packages/kilo-jetbrains/frontend/src/main/resources/kilo.jetbrains.frontend.xml b/packages/kilo-jetbrains/frontend/src/main/resources/kilo.jetbrains.frontend.xml index 2c8207d401e..b7515a2cd2d 100644 --- a/packages/kilo-jetbrains/frontend/src/main/resources/kilo.jetbrains.frontend.xml +++ b/packages/kilo-jetbrains/frontend/src/main/resources/kilo.jetbrains.frontend.xml @@ -63,6 +63,11 @@ overrides="false"/> + + + + diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/ui/attachment/AttachmentEditorKindTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/ui/attachment/AttachmentEditorKindTest.kt index b18e32a35d1..0bdb0dca4d8 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/ui/attachment/AttachmentEditorKindTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/ui/attachment/AttachmentEditorKindTest.kt @@ -6,7 +6,9 @@ import ai.kilocode.client.session.model.FileAttachment import ai.kilocode.client.testing.FakeAppRpcApi import ai.kilocode.client.testing.FakeSessionRpcApi import ai.kilocode.client.vfs.KiloPath +import ai.kilocode.client.vfs.KiloEditorKindRegistry import ai.kilocode.client.vfs.KiloVirtualFile +import ai.kilocode.client.vfs.KiloVirtualFileKindRegistry import ai.kilocode.client.vfs.KiloVirtualFileSystem import ai.kilocode.rpc.dto.KiloAppStateDto import ai.kilocode.rpc.dto.KiloAppStatusDto @@ -105,6 +107,30 @@ class AttachmentEditorKindTest : BasePlatformTestCase() { assertNull(VirtualFileManager.getInstance().findFileByUrl(file.url)) } + fun testAttachmentEditorKindAndVirtualFilesCanBeCleared() { + ensureAttachmentEditorKind() + val fs = KiloVirtualFileSystem.getInstance() + val path = KiloPath(AttachmentEditorKind.ID, mapOf( + "directory" to "/repo", + "sessionId" to "ses1", + "messageId" to "msg1", + "partId" to "part1", + "filename" to "note.txt", + )) + val file = fs.findOrCreateFile(path) + + assertNotNull(file) + assertNotNull(service().get(AttachmentEditorKind.ID)) + assertNotNull(service().get(AttachmentEditorKind.ID)) + + unregisterAttachmentEditorKind() + fs.clear() + + assertNull(service().get(AttachmentEditorKind.ID)) + assertNull(service().get(AttachmentEditorKind.ID)) + assertNull(fs.findOrCreateFile(path)) + } + @Suppress("UnstableApiUsage") fun testFetchUsesAttachmentKeyBeforeDuplicatePartId() { val cs = CoroutineScope(SupervisorJob())