fix(kilo-jetbrains): stabilize model settings saves

.changeset/quiet-garden-model-settings.md

@@ -2,4 +2,4 @@
 "@kilocode/kilo-jetbrains": patch
 ---
 
-Fix model settings layout and prevent login banner flicker while saving.
+Fix model settings layout and keep pending model selections visible while saving.

packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/settings/models/ModelsSettingsUi.kt

@@ -87,13 +87,13 @@ internal class ModelsSettingsUi(
     @RequiresEdt
     fun modified(): Boolean {
         checkEdt()
-        return draft != baseline
+        return draft != (pending ?: baseline)
     }
 
     @RequiresEdt
     fun resetDraft() {
         checkEdt()
-        draft = baseline
+        draft = pending ?: baseline
         saveError = null
         if (!saving) clearProgress()
         sync()
@@ -125,10 +125,11 @@ internal class ModelsSettingsUi(
                 }
                 if (state != null) {
                     LOG.info("model settings save: completed ${summary(patch)}")
+                    val edit = draft
                     appState = state
                     val base = modelsDraft(state.config, agents)
                     baseline = if (savedMatches(base, next)) base else next
-                    draft = next
+                    draft = if (edit == next) baseline else edit
                     pending = null
                     saving = false
                     saveError = null
@@ -136,8 +137,9 @@ internal class ModelsSettingsUi(
                     sync()
                     return@invokeLater
                 }
+                val edit = draft
                 baseline = prev
-                draft = next
+                draft = if (edit == next) next else edit
                 pending = null
                 saving = false
                 LOG.warn("model settings save: failed ${summary(patch)}")
@@ -243,6 +245,7 @@ internal class ModelsSettingsUi(
             saving = saving,
         )
         val ready = state == ModelsStatus.READY || state == ModelsStatus.MODES_FAILED
+        val editable = !saving && (ready || state == ModelsStatus.LOADING)
         val bannerVisible = modelsLoginBannerVisible(
             ready = appState.status == KiloAppStatusDto.READY,
             authenticated = appState.profile != null,
@@ -262,9 +265,9 @@ internal class ModelsSettingsUi(
         defaults.setItems(allItems, draft.model)
         small.setItems(smallItems, draft.small)
         subagent.setItems(allItems, draft.subagent)
-        listOf(defaults, small, subagent).forEach { it.isEnabled = ready }
-        layout = syncVariant(ready) || layout
-        layout = syncModes(ready) || layout
+        listOf(defaults, small, subagent).forEach { it.isEnabled = editable }
+        layout = syncVariant(editable) || layout
+        layout = syncModes(editable) || layout
         if (layout) {
             revalidate()
             repaint()
@@ -388,14 +391,14 @@ internal class ModelsSettingsUi(
     private fun acceptBase(base: ModelsDraft) {
         val save = pending
         if (save == null) {
+            val prev = baseline
+            val edit = draft
             baseline = base
-            if (!saving) draft = base
+            if (edit == prev) draft = base
             return
         }
         if (!savedMatches(base, save)) return
         baseline = base
-        draft = save
-        pending = null
     }
 
     private fun checkEdt() {

packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/settings/models/ModelsSettingsUiTest.kt

@@ -103,6 +103,141 @@ class ModelsSettingsUiTest : BasePlatformTestCase() {
         flushUntil { !text(panel).contains("Failed to save model settings") }
     }
 
+    fun `test reset during pending save keeps applied selection visible`() {
+        val panel = requireUi()
+        rpc.configUpdateGate = CompletableDeferred()
+
+        edt {
+            select(panel, "new")
+            panel.applyDraft()
+            panel.resetDraft()
+            assertTrue(text(panel.progress).contains("Saving model settings"))
+            assertFalse(panel.modified())
+            assertSelected(panel, "kilo/new")
+        }
+
+        rpc.configUpdateGate?.complete(Unit)
+        flushUntil { rpc.configPatches.isNotEmpty() }
+        edt { assertSelected(panel, "kilo/new") }
+    }
+
+    fun `test pickers are disabled during pending save`() {
+        val panel = requireUi()
+        rpc.configUpdateGate = CompletableDeferred()
+
+        edt {
+            select(panel, "new")
+            panel.applyDraft()
+            assertTrue(pickers(panel).isNotEmpty())
+            assertTrue(pickers(panel).all { !it.isEnabled })
+        }
+
+        rpc.configUpdateGate?.complete(Unit)
+        flushUntil { rpc.configPatches.isNotEmpty() }
+    }
+
+    fun `test matching app state before save callback keeps pending target`() {
+        val panel = requireUi()
+        rpc.configUpdateGate = CompletableDeferred()
+
+        edt {
+            select(panel, "new")
+            panel.applyDraft()
+        }
+        rpc.state.value = state("kilo/new")
+        flushUntil { edt { !panel.modified() && pickers(panel).all { !it.isEnabled } } }
+        edt {
+            panel.resetDraft()
+            assertSelected(panel, "kilo/new")
+            assertTrue(text(panel.progress).contains("Saving model settings"))
+        }
+
+        rpc.configUpdateGate?.complete(Unit)
+        flushUntil { rpc.configPatches.isNotEmpty() }
+        edt { assertSelected(panel, "kilo/new") }
+    }
+
+    fun `test matching app state after save callback preserves dirty draft`() {
+        val panel = requireUi()
+        rpc.configUpdateGate = CompletableDeferred()
+
+        edt {
+            select(panel, "new")
+            panel.applyDraft()
+            select(panel, "old")
+        }
+        rpc.configUpdateGate?.complete(Unit)
+        flushUntil { rpc.configPatches.isNotEmpty() }
+        rpc.state.value = state("kilo/new")
+        flushUntil { edt { panel.modified() } }
+
+        edt {
+            assertSelected(panel, "kilo/old")
+            assertTrue(panel.modified())
+        }
+    }
+
+    fun `test stale app state during pending save is ignored`() {
+        val panel = requireUi()
+        rpc.configUpdateGate = CompletableDeferred()
+
+        edt {
+            select(panel, "new")
+            panel.applyDraft()
+        }
+        rpc.state.value = state("kilo/old")
+        flushUntil { edt { text(panel.progress).contains("Saving model settings") } }
+
+        edt {
+            assertSelected(panel, "kilo/new")
+            assertFalse(panel.modified())
+        }
+
+        rpc.configUpdateGate?.complete(Unit)
+        flushUntil { rpc.configPatches.isNotEmpty() }
+    }
+
+    fun `test pickers stay enabled while models load`() {
+        edt {
+            requireUi().dispose()
+            ui = null
+        }
+        uiScope = CoroutineScope(SupervisorJob())
+        workspaceRpc.modelsGate = CompletableDeferred()
+        workspaceRpc.models = ModelsWorkspaceDto(providers = providers())
+
+        edt { ui = ModelsSettingsUi(uiScope, app, workspaces, directory = "/test") }
+        val panel = requireUi()
+
+        flushUntil { text(panel.progress).contains("Loading models") }
+        edt {
+            assertTrue(pickers(panel).isNotEmpty())
+            assertTrue(pickers(panel).all { it.isEnabled })
+        }
+
+        workspaceRpc.modelsGate?.complete(Unit)
+        flushUntil { text(panel).contains("Old") }
+    }
+
+    fun `test edit during pending save is preserved after completion`() {
+        val panel = requireUi()
+        rpc.configUpdateGate = CompletableDeferred()
+
+        edt {
+            select(panel, "new")
+            panel.applyDraft()
+            select(panel, "old")
+            assertSelected(panel, "kilo/old")
+        }
+
+        rpc.configUpdateGate?.complete(Unit)
+        flushUntil { rpc.configPatches.isNotEmpty() }
+        edt {
+            assertSelected(panel, "kilo/old")
+            assertTrue(panel.modified())
+        }
+    }
+
     fun `test failed save after dispose shows notification`() {
         val notes = mutableListOf<Notification>()
         ApplicationManager.getApplication().messageBus.connect(testRootDisposable).subscribe(
@@ -190,11 +325,23 @@ class ModelsSettingsUiTest : BasePlatformTestCase() {
         defaults = emptyMap(),
     )
 
+    private fun state(model: String): KiloAppStateDto = KiloAppStateDto(
+        KiloAppStatusDto.READY,
+        config = ConfigDto(model = model),
+        profile = ProfileDto(email = "alice@test.com"),
+    )
+
     private fun select(panel: ModelsSettingsUi, id: String) {
-        val picker = components(panel).filterIsInstance<ModelPicker>().first()
+        val picker = pickers(panel).first()
         picker.onSelect(ModelPicker.Item(id, id.replaceFirstChar { it.titlecase() }, "kilo", "Kilo"))
     }
 
+    private fun assertSelected(panel: ModelsSettingsUi, key: String) {
+        assertEquals(key, pickers(panel).first().selectionKeyForTest())
+    }
+
+    private fun pickers(panel: ModelsSettingsUi): List<ModelPicker> = components(panel).filterIsInstance<ModelPicker>()
+
     private fun requireUi(): ModelsSettingsUi = requireNotNull(ui)
 
     private fun <T> edt(block: () -> T): T {

packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/testing/FakeWorkspaceRpcApi.kt

@@ -5,6 +5,7 @@ import ai.kilocode.rpc.dto.KiloWorkspaceStateDto
 import ai.kilocode.rpc.dto.KiloWorkspaceStatusDto
 import ai.kilocode.rpc.dto.ModelsWorkspaceDto
 import ai.kilocode.rpc.dto.WorkspaceFileDto
+import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
 
@@ -23,6 +24,7 @@ class FakeWorkspaceRpcApi : KiloWorkspaceRpcApi {
     var reloads = 0
         private set
     var models = ModelsWorkspaceDto()
+    var modelsGate: CompletableDeferred<Unit>? = null
     var fileMatches = emptyList<WorkspaceFileDto>()
     var openResult = true
     val fileCalls = mutableListOf<Pair<String, String>>()
@@ -45,6 +47,7 @@ class FakeWorkspaceRpcApi : KiloWorkspaceRpcApi {
 
     override suspend fun models(directory: String): ModelsWorkspaceDto {
         assertNotEdt("models")
+        modelsGate?.await()
         return models
     }