From 750bb778998a405719d3132d17867d2868e4defd Mon Sep 17 00:00:00 2001 From: kirillk Date: Tue, 2 Jun 2026 20:57:50 -0400 Subject: [PATCH] fix(jetbrains): confine session subscriptions to EDT --- .../jetbrains-controller-subscriptions.md | 5 +++ .../session/controller/SessionController.kt | 43 +++++++++++++------ .../session/controller/ChatLoggingFlowTest.kt | 7 +++ .../session/controller/PromptLifecycleTest.kt | 15 +++++++ specs/jetbrains-session-ui-findings-todo.md | 26 ++++++----- 5 files changed, 74 insertions(+), 22 deletions(-) create mode 100644 .changeset/jetbrains-controller-subscriptions.md diff --git a/.changeset/jetbrains-controller-subscriptions.md b/.changeset/jetbrains-controller-subscriptions.md new file mode 100644 index 00000000000..d820e2ddb22 --- /dev/null +++ b/.changeset/jetbrains-controller-subscriptions.md @@ -0,0 +1,5 @@ +--- +"kilo-code": patch +--- + +Improve JetBrains session stability by keeping controller subscription state on the UI thread. diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/controller/SessionController.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/controller/SessionController.kt index 0075a889582..609cfdd45d6 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/controller/SessionController.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/controller/SessionController.kt @@ -50,6 +50,7 @@ import com.intellij.openapi.application.ApplicationManager import ai.kilocode.log.ChatLogSummary import ai.kilocode.log.KiloLog import com.intellij.openapi.util.Disposer +import com.intellij.util.concurrency.annotations.RequiresEdt import kotlinx.coroutines.CoroutineScope import kotlinx.coroutines.Job import kotlinx.coroutines.cancel @@ -223,7 +224,10 @@ class SessionController( val meta = if (LOG.isDebugEnabled) ChatLogSummary.dir(directory) else "kind=session" LOG.info("${ChatLogSummary.sid(session.id)} kind=session $meta created=true") capture("Task Created", sessionProps(session.id) + mapOf("source" to "jetbrains")) - subscribeEvents() + runEdt { + if (disposed) return@runEdt + subscribeEvents() + } session.id } sessions.prompt(id, directory, dto) @@ -458,7 +462,9 @@ class SessionController( } } + @RequiresEdt private fun drainAutoApprove(skip: Set = emptySet()) { + assertEdt() val id = sid ?: return val ids = (childIds + id).toSet() drainJob?.cancel() @@ -749,13 +755,12 @@ class SessionController( if (!model.showSession) setControllerViewState(SessionControllerEvent.ViewChanged.ShowProgress) } + @RequiresEdt private fun subscribeEvents() { + assertEdt() val id = sid ?: return LOG.debug { "${ChatLogSummary.sid(id)} kind=subscription subscribe=true" } - eventJob?.cancel() - childJobs.values.forEach { it.cancel() } - childJobs.clear() - childIds.clear() + cancelSubscriptions() eventJob = cs.launch { try { sessions.events(id, directory).collect { event -> @@ -772,7 +777,9 @@ class SessionController( } } + @RequiresEdt private fun subscribeChild(child: String) { + assertEdt() if (childJobs.containsKey(child)) return LOG.debug { "${ChatLogSummary.sid(sid ?: "pending")} kind=child-subscription child=$child subscribe=true" } val job = cs.launch { @@ -789,12 +796,24 @@ class SessionController( childJobs[child] = job } + @RequiresEdt private fun trackChild(child: String) { + assertEdt() if (!childIds.add(child)) return subscribeChild(child) cs.launch { recoverChildPermissions(child) } } + @RequiresEdt + private fun cancelSubscriptions() { + assertEdt() + eventJob?.cancel() + eventJob = null + childJobs.values.forEach { it.cancel() } + childJobs.clear() + childIds.clear() + } + private suspend fun recoverChildPermissions(child: String) { try { val permissions = sessions.pendingPermissions(directory).filter { it.sessionID == child } @@ -1696,13 +1715,13 @@ class SessionController( } override fun dispose() { - disposed = true - connectionDelay.dispose() - eventJob?.cancel() - drainJob?.cancel() - childJobs.values.forEach { it.cancel() } - childJobs.clear() - childIds.clear() + runEdt { + disposed = true + connectionDelay.dispose() + cancelSubscriptions() + drainJob?.cancel() + drainJob = null + } cs.cancel() } diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/controller/ChatLoggingFlowTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/controller/ChatLoggingFlowTest.kt index 312913c29c1..1f5f444dd4f 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/controller/ChatLoggingFlowTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/controller/ChatLoggingFlowTest.kt @@ -10,10 +10,16 @@ import ai.kilocode.rpc.dto.QuestionInfoDto import ai.kilocode.rpc.dto.QuestionRequestDto import ai.kilocode.rpc.dto.SessionStatusDto import kotlinx.coroutines.flow.flow +import java.util.concurrent.CopyOnWriteArrayList class ChatLoggingFlowTest : SessionControllerTestBase() { fun `test prompt creates session and subscribes before dispatch`() { + val calls = CopyOnWriteArrayList() + rpc.eventFlow = { id, _ -> + calls.add(id) + rpc.events + } projectRpc.state.value = workspaceReady() appRpc.state.value = KiloAppStateDto(KiloAppStatusDto.READY) val m = controller() @@ -23,6 +29,7 @@ class ChatLoggingFlowTest : SessionControllerTestBase() { flush() assertEquals(1, rpc.creates) + assertEquals(listOf("ses_test"), calls.toList()) assertEquals(1, rpc.prompts.size) assertEquals("ses_test", rpc.prompts[0].first) } diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/controller/PromptLifecycleTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/controller/PromptLifecycleTest.kt index 04aeadf6630..07dcfb934f3 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/controller/PromptLifecycleTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/controller/PromptLifecycleTest.kt @@ -19,6 +19,7 @@ import ai.kilocode.rpc.dto.QuestionOptionDto import ai.kilocode.rpc.dto.QuestionReplyDto import ai.kilocode.rpc.dto.QuestionRequestDto import ai.kilocode.rpc.dto.ToolRefDto +import java.util.concurrent.CopyOnWriteArrayList class PromptLifecycleTest : SessionControllerTestBase() { @@ -496,6 +497,20 @@ class PromptLifecycleTest : SessionControllerTestBase() { assertEquals("ses_child", perm.sessionId) } + fun `test repeated task part subscribes to child once`() { + val calls = CopyOnWriteArrayList() + rpc.eventFlow = { id, _ -> + calls.add(id) + rpc.events + } + prompted() + + emit(taskPart("ses_child"), flush = false) + emit(taskPart("ses_child")) + + assertEquals(1, calls.count { it == "ses_child" }) + } + fun `test child PermissionAsked moves root model to AwaitingPermission`() { val (m, _, _) = prompted() diff --git a/specs/jetbrains-session-ui-findings-todo.md b/specs/jetbrains-session-ui-findings-todo.md index c72e95cc9b1..7b97510fd91 100644 --- a/specs/jetbrains-session-ui-findings-todo.md +++ b/specs/jetbrains-session-ui-findings-todo.md @@ -4,31 +4,37 @@ Use this as a planning backlog for the JetBrains session UI performance, memory, Last status check: 2026-06-02. +## Current Summary + +- Implemented high-priority work: streaming markdown no longer rebuilds the full rendered tree on normal deltas, and hidden cached session UIs are disposed after a configurable timeout while `SessionUpdateQueue` uses a shared coroutine ticker instead of per-session scheduler threads. +- Remaining high-priority work: none. `SessionController` subscription-state mutation is now EDT-confined while RPC event collection remains on background coroutines. +- Remaining non-high-priority work: repaint/revalidate cleanup, lazy collapsed bodies, question editor disposal, EDT annotations/assertions, style callback disposal guards, `SessionUpdateQueue` Swing listener EDT confinement, question body retention, semantic timeline colors, and additional retained Swing regression tests. + ## High Priority - [x] Fix streaming markdown full-tree rebuilds - Severity: High - - Status: Addressed for the high-priority full-tree rebuild issue. Remaining repaint/revalidate cleanup is tracked separately under medium priority. + - Status: Implemented for the high-priority full-tree rebuild issue. Remaining repaint/revalidate cleanup is tracked separately under medium priority. - Files: `packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/SessionMessageListPanel.kt`, `packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/views/TextView.kt`, `packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/md/MdViewHybrid.kt` - Issue: `ContentDelta` updates fetch full content and route through `TextView.update()`/`md.set()`, causing markdown to reparse and recreate all child components, including code editors, on each streaming chunk. - - Plan direction: Restore incremental delta handling where safe, avoid full `set()` for append-only text, and retain/reuse markdown/code block components where possible. - - Current evidence: `ContentDelta` carries `created`, `SessionMessageListPanel` routes normal deltas through `appendDelta()`, and `MdViewHybrid` retains compatible HTML/code block views instead of clearing all rendered blocks. + - Implemented: `ContentDelta` carries `created`, `SessionMessageListPanel` routes normal deltas through `appendDelta()`, `TextView.appendDelta()` calls `md.append(delta)`, and `MdViewHybrid` retains compatible HTML/code block views instead of clearing all rendered blocks. + - Tests/release note: `SessionMessageListPanelTest` covers preserving the `TextView` and markdown component during deltas; `.changeset/retained-jetbrains-markdown.md` documents the user-facing fix. - [x] Prevent hidden cached sessions from accumulating resources - Severity: High - - Status: Addressed with timed hidden-session disposal and shared coroutine-based queue ticking. + - Status: Implemented with timed hidden-session disposal and shared coroutine-based queue ticking. - Files: `packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/SessionSidePanelManager.kt`, `packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/controller/SessionUpdateQueue.kt`, `packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/controller/SessionController.kt` - Issue: Inactive cached session UIs keep controller subscriptions, scheduler threads, models, Swing trees, and pending event queues alive. Non-metadata events can accumulate while hidden. - - Plan direction: Bound hidden transcript event queues with unconditional timed disposal, and avoid one scheduler thread per session UI. - - Current evidence: `kilo.session.inactive.dispose` was removed, `kilo.session.inactive.disposeTimeoutMs` now defaults to 180000 ms, hidden cached `SessionUi` instances are disposed by an EDT timer after the timeout unless shown again, and `SessionUpdateQueue` now uses the controller coroutine scope for its ticker instead of creating a per-queue scheduler thread. + - Implemented: `kilo.session.inactive.dispose` was removed, `kilo.session.inactive.disposeTimeoutMs` now defaults to 180000 ms, hidden cached `SessionUi` instances are disposed by an EDT timer after the timeout unless shown again, and `SessionUpdateQueue` now uses the controller coroutine scope for its ticker instead of creating a per-queue scheduler thread. + - Tests/release note: `SessionSidePanelManagerTest` covers hidden timeout disposal, reopen-after-timeout recreation, busy hidden disposal, and permission hidden disposal; `.changeset/hidden-session-timers.md` documents the user-facing fix. -- [ ] Confine controller subscription mutation to EDT +- [x] Confine controller subscription mutation to EDT - Severity: High - - Status: Not addressed; still needs work. + - Status: Implemented with EDT-only subscription bookkeeping and background RPC event collection. - Files: `packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/controller/SessionController.kt` - Issue: `prompt()` calls `subscribeEvents()` from a background coroutine after session creation; `subscribeEvents()` mutates `eventJob`, `childJobs`, and `childIds` without EDT confinement. - - Plan direction: Run subscription-state mutation on EDT or separate thread-safe subscription state from EDT-only controller state. - - Current evidence: `prompt()` still calls `subscribeEvents()` from inside `cs.launch`, and `subscribeEvents()` still mutates `eventJob`, `childJobs`, and `childIds` without an EDT assertion or EDT handoff. + - Implemented: `subscribeEvents()`, `subscribeChild()`, `trackChild()`, `drainAutoApprove()`, and subscription cleanup are `@RequiresEdt`/asserted EDT-only. New-session prompt creation hands subscription setup back to EDT before prompt dispatch, and `dispose()` clears subscription state on EDT before cancelling the coroutine scope. + - Tests/release note: `ChatLoggingFlowTest` covers new-session event subscription, `PromptLifecycleTest` covers duplicate child task parts subscribing once, and `.changeset/jetbrains-controller-subscriptions.md` documents the stability fix. ## Medium Priority