refactor(kilo-jetbrains): move settings workspace lifecycle to base

packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/settings/base/BaseSettingsUi.kt

@@ -18,7 +18,7 @@ import java.util.function.Predicate
 import javax.swing.JComponent
 
 internal abstract class BaseSettingsUi<C : BaseContentPanel, D, P, R>(
-    private val cs: CoroutineScope,
+    protected val scope: CoroutineScope,
     initial: D,
     private val loginBanner: Boolean = true,
 ) : SettingsPanel() {
@@ -109,7 +109,7 @@ internal abstract class BaseSettingsUi<C : BaseContentPanel, D, P, R>(
         disposed = true
         jobs.forEach { it.cancel() }
         jobs.clear()
-        cs.cancel()
+        scope.cancel()
     }
 
     @RequiresEdt
@@ -120,7 +120,9 @@ internal abstract class BaseSettingsUi<C : BaseContentPanel, D, P, R>(
         syncContent()
     }
 
+    @RequiresEdt
     protected fun acceptBase(base: D) {
+        checkEdt()
         val target = pending
         if (target == null) {
             val prev = baseline
@@ -147,19 +149,43 @@ internal abstract class BaseSettingsUi<C : BaseContentPanel, D, P, R>(
         check(ApplicationManager.getApplication().isDispatchThread) { "Settings UI updates must run on EDT" }
     }
 
+    @RequiresEdt
     protected abstract fun change(from: D, to: D): P?
+
+    @RequiresEdt
     protected abstract fun save(change: P, done: (R?) -> Unit)
+
+    @RequiresEdt
     protected abstract fun base(result: R): D
+
+    @RequiresEdt
     protected abstract fun syncContent()
+
+    @RequiresEdt
     protected abstract fun pendingText(): String
+
+    @RequiresEdt
     protected abstract fun failedText(): String
 
+    @RequiresEdt
     protected open fun saved(base: D, draft: D): Boolean = base == draft
+
+    @RequiresEdt
     protected open fun onSaveFailedAfterDispose(change: P) = KiloNotifications.error(failedText())
+
+    @RequiresEdt
     protected open fun logSaveStarted(change: P) = Unit
+
+    @RequiresEdt
     protected open fun logSaveCompleted(change: P) = Unit
+
+    @RequiresEdt
     protected open fun logSaveFailed(change: P) = Unit
+
+    @RequiresEdt
     protected open fun logSaveFailedAfterDispose(change: P) = Unit
+
+    @RequiresEdt
     protected open fun logSaveCompletedAfterDispose(change: P) = Unit
 
     private fun openProfile(src: JComponent) {

packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/settings/base/BaseWorkspaceSettingsUi.kt

@@ -0,0 +1,127 @@
+package ai.kilocode.client.settings.base
+
+import ai.kilocode.client.app.KiloAppService
+import ai.kilocode.client.app.KiloWorkspaceService
+import ai.kilocode.rpc.dto.KiloAppStateDto
+import ai.kilocode.rpc.dto.KiloAppStatusDto
+import ai.kilocode.rpc.dto.ModelStateDto
+import com.intellij.openapi.application.EDT
+import com.intellij.openapi.application.ModalityState
+import com.intellij.openapi.application.asContextElement
+import com.intellij.openapi.components.service
+import com.intellij.util.concurrency.annotations.RequiresBackgroundThread
+import com.intellij.util.concurrency.annotations.RequiresEdt
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.withContext
+
+internal abstract class BaseWorkspaceSettingsUi<C : BaseContentPanel, D, P, R, W>(
+    cs: CoroutineScope,
+    initial: D,
+    private val app: KiloAppService = service(),
+    private val workspaces: KiloWorkspaceService = service(),
+    private val hint: String? = null,
+    loginBanner: Boolean = true,
+) : BaseSettingsUi<C, D, P, R>(cs, initial, loginBanner) {
+    protected var appState: KiloAppStateDto = app.state.value
+        private set
+    protected var modelState: ModelStateDto = app.models.value
+        private set
+    protected var projectDirectory: String? = null
+        private set
+    protected val hasProjectDirectory get() = projectDirectory != null || hint != null
+    protected var workspaceLoading = false
+        private set
+    protected var workspaceLoaded = false
+        private set
+
+    @RequiresEdt
+    protected fun startSettings(content: C) {
+        setSettingsContent(content)
+        syncContent()
+        start()
+    }
+
+    private fun start() {
+        jobs += scope.launch {
+            app.state.collect { state -> withContext(edt) { updateApp(state) } }
+        }
+        jobs += scope.launch {
+            app.models.collect { state -> withContext(edt) { updateModels(state) } }
+        }
+        jobs += scope.launch { app.connect() }
+        val path = hint ?: return
+        jobs += scope.launch {
+            val dir = workspaces.resolveProjectDirectory(path)
+            withContext(edt) {
+                projectDirectory = dir
+                workspaceLoaded = false
+                syncContent()
+                load()
+            }
+        }
+    }
+
+    @RequiresEdt
+    private fun updateApp(state: KiloAppStateDto) {
+        appState = state
+        if (state.status != KiloAppStatusDto.READY) {
+            workspaceLoading = false
+            unavailable(state)
+            syncContent()
+            return
+        }
+        acceptBase(draft(state))
+        syncContent()
+        load()
+    }
+
+    @RequiresEdt
+    private fun updateModels(state: ModelStateDto) {
+        modelState = state
+        models(state)
+        syncContent()
+    }
+
+    @RequiresEdt
+    private fun load() {
+        val root = projectDirectory ?: return
+        if (appState.status != KiloAppStatusDto.READY || workspaceLoading || workspaceLoaded) return
+        workspaceLoading = true
+        clearWorkspaceError()
+        syncContent()
+        jobs += scope.launch {
+            val state = loadWorkspace(root)
+            withContext(edt) {
+                applyWorkspace(state)
+                workspaceLoaded = true
+                workspaceLoading = false
+                acceptBase(draft(appState))
+                syncContent()
+            }
+        }
+    }
+
+    @RequiresEdt
+    protected abstract fun draft(state: KiloAppStateDto): D
+
+    @RequiresBackgroundThread
+    protected abstract suspend fun loadWorkspace(root: String): W
+
+    @RequiresEdt
+    protected abstract fun applyWorkspace(result: W)
+
+    @RequiresEdt
+    protected open fun unavailable(state: KiloAppStateDto) = Unit
+
+    @RequiresEdt
+    protected open fun models(state: ModelStateDto) = Unit
+
+    @RequiresEdt
+    protected open fun clearWorkspaceError() = Unit
+
+    private companion object {
+        val edt = Dispatchers.EDT + ModalityState.any().asContextElement()
+    }
+}

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

@@ -6,9 +6,8 @@ import ai.kilocode.client.plugin.KiloBundle
 import ai.kilocode.client.session.ui.ReasoningPicker
 import ai.kilocode.client.session.ui.model.ModelPicker
 import ai.kilocode.client.session.ui.model.ModelText
-import ai.kilocode.client.settings.profile.edt
 import ai.kilocode.client.settings.base.BaseContentPanel
-import ai.kilocode.client.settings.base.BaseSettingsUi
+import ai.kilocode.client.settings.base.BaseWorkspaceSettingsUi
 import ai.kilocode.client.settings.base.SettingsBannerKind
 import ai.kilocode.client.settings.base.SettingsRow
 import ai.kilocode.client.settings.base.SettingsRows
@@ -21,20 +20,25 @@ import ai.kilocode.rpc.dto.ConfigPatchDto
 import ai.kilocode.rpc.dto.KiloAppStateDto
 import ai.kilocode.rpc.dto.KiloAppStatusDto
 import ai.kilocode.rpc.dto.LoadErrorDto
+import ai.kilocode.rpc.dto.ModelStateDto
 import ai.kilocode.rpc.dto.ModelsWorkspaceDto
 import ai.kilocode.rpc.dto.ProvidersDto
 import com.intellij.openapi.components.service
 import com.intellij.util.concurrency.annotations.RequiresEdt
 import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.launch
-import kotlinx.coroutines.withContext
 
 internal class ModelsSettingsUi(
-    private val cs: CoroutineScope,
+    cs: CoroutineScope,
     private val app: KiloAppService = service(),
     private val workspaces: KiloWorkspaceService = service(),
-    private val directory: String? = null,
-) : BaseSettingsUi<ModelsSettingsContent, ModelsDraft, ConfigPatchDto, KiloAppStateDto>(cs, ModelsDraft()) {
+    directory: String? = null,
+) : BaseWorkspaceSettingsUi<ModelsSettingsContent, ModelsDraft, ConfigPatchDto, KiloAppStateDto, ModelsWorkspaceDto>(
+    cs,
+    ModelsDraft(),
+    app,
+    workspaces,
+    directory,
+) {
 
     companion object {
         private val LOG = KiloLog.create(ModelsSettingsUi::class.java)
@@ -49,17 +53,11 @@ internal class ModelsSettingsUi(
 
     private var providers: ProvidersDto? = null
     private var agents: List<AgentDto> = emptyList()
-    private var appState: KiloAppStateDto = app.state.value
-    private var dir: String? = null
-    private var loading = false
-    private var loaded = false
     private var errors: List<LoadErrorDto> = emptyList()
     private var allItems: List<ModelPicker.Item> = emptyList()
 
     init {
-        setSettingsContent(ModelsSettingsContent(app, { updateDraft(it) }, ::selectSubagent))
-        syncContent()
-        start()
+        startSettings(ModelsSettingsContent(app, { updateDraft(it) }, ::selectSubagent))
     }
 
     override fun change(from: ModelsDraft, to: ModelsDraft): ConfigPatchDto? = patch(from, to).takeIf {
@@ -70,10 +68,9 @@ internal class ModelsSettingsUi(
         app.updateConfigAsync(change, done)
     }
 
-    override fun base(result: KiloAppStateDto): ModelsDraft {
-        appState = result
-        return modelsDraft(result.config, agents)
-    }
+    override fun base(result: KiloAppStateDto): ModelsDraft = modelsDraft(result.config, agents)
+
+    override fun draft(state: KiloAppStateDto): ModelsDraft = modelsDraft(state.config, agents)
 
     override fun saved(base: ModelsDraft, draft: ModelsDraft): Boolean = savedMatches(base, draft)
 
@@ -91,82 +88,34 @@ internal class ModelsSettingsUi(
 
     override fun logSaveCompletedAfterDispose(change: ConfigPatchDto) = LOG.info("model settings save: completed after dispose ${summary(change)}")
 
-    @RequiresEdt
-    fun updateApp(state: KiloAppStateDto) {
-        appState = state
-        if (state.status != KiloAppStatusDto.READY) {
-            loading = false
-            if (!loaded && providers == null) {
-                agents = emptyList()
-                errors = emptyList()
-            }
-            syncContent()
-            return
-        }
-        val base = modelsDraft(state.config, agents)
-        acceptBase(base)
-        syncContent()
-        loadModels()
-    }
-
-    @RequiresEdt
-    fun updateModelsWorkspace(state: ModelsWorkspaceDto) {
-        providers = state.providers
-        agents = state.agents?.agents ?: emptyList()
-        errors = state.errors
-        loaded = true
-        loading = false
-        val base = modelsDraft(appState.config, agents)
-        acceptBase(base)
-        syncContent()
-    }
-
-    @RequiresEdt
-    fun updateModels(state: ai.kilocode.rpc.dto.ModelStateDto) {
-        syncContent()
-    }
-
-    private fun start() {
-        jobs += cs.launch {
-            app.state.collect { state -> withContext(edt) { updateApp(state) } }
-        }
-        jobs += cs.launch {
-            app.models.collect { state -> withContext(edt) { updateModels(state) } }
-        }
-        jobs += cs.launch { app.connect() }
-        val hint = directory ?: return
-        jobs += cs.launch {
-            val resolved = workspaces.resolveProjectDirectory(hint)
-            withContext(edt) {
-                dir = resolved
-                loaded = false
-                loadModels()
-                syncContent()
-            }
+    override fun unavailable(state: KiloAppStateDto) {
+        if (!workspaceLoaded && providers == null) {
+            agents = emptyList()
+            errors = emptyList()
         }
     }
 
-    @RequiresEdt
-    private fun loadModels() {
-        val root = dir ?: return
-        if (appState.status != KiloAppStatusDto.READY || loading || loaded) return
-        loading = true
+    override fun models(state: ModelStateDto) = Unit
+
+    override fun clearWorkspaceError() {
         errors = emptyList()
-        syncContent()
-        jobs += cs.launch {
-            val state = workspaces.models(root)
-            withContext(edt) { updateModelsWorkspace(state) }
-        }
+    }
+
+    override suspend fun loadWorkspace(root: String): ModelsWorkspaceDto = workspaces.models(root)
+
+    override fun applyWorkspace(result: ModelsWorkspaceDto) {
+        providers = result.providers
+        agents = result.agents?.agents ?: emptyList()
+        errors = result.errors
     }
 
     @RequiresEdt
     override fun syncContent() {
         allItems = items(false)
         val smallItems = items(true)
-        val hasDir = dir != null || directory != null
         val state = modelsStatus(
-            ready = appState.status == KiloAppStatusDto.READY && hasDir,
-            loading = loading || (appState.status == KiloAppStatusDto.READY && !loaded && hasDir),
+            ready = appState.status == KiloAppStatusDto.READY && hasProjectDirectory,
+            loading = workspaceLoading || (appState.status == KiloAppStatusDto.READY && !workspaceLoaded && hasProjectDirectory),
             providers = providers,
             items = allItems.size,
             errors = errors,

packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/settings/base/BaseWorkspaceSettingsUiTest.kt

@@ -0,0 +1,191 @@
+package ai.kilocode.client.settings.base
+
+import ai.kilocode.client.app.KiloAppService
+import ai.kilocode.client.app.KiloWorkspaceService
+import ai.kilocode.client.testing.FakeAppRpcApi
+import ai.kilocode.client.testing.FakeWorkspaceRpcApi
+import ai.kilocode.rpc.dto.ConfigDto
+import ai.kilocode.rpc.dto.KiloAppStateDto
+import ai.kilocode.rpc.dto.KiloAppStatusDto
+import ai.kilocode.rpc.dto.ModelSelectionDto
+import ai.kilocode.rpc.dto.ModelStateDto
+import com.intellij.openapi.application.ApplicationManager
+import com.intellij.testFramework.fixtures.BasePlatformTestCase
+import com.intellij.util.ui.UIUtil
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.SupervisorJob
+import kotlinx.coroutines.cancel
+import kotlinx.coroutines.delay
+import kotlinx.coroutines.runBlocking
+
+class BaseWorkspaceSettingsUiTest : BasePlatformTestCase() {
+    private lateinit var appScope: CoroutineScope
+    private lateinit var uiScope: CoroutineScope
+    private lateinit var rpc: FakeAppRpcApi
+    private lateinit var workspaceRpc: FakeWorkspaceRpcApi
+    private lateinit var app: KiloAppService
+    private lateinit var workspaces: KiloWorkspaceService
+    private var panel: FakePanel? = null
+
+    override fun setUp() {
+        super.setUp()
+        appScope = CoroutineScope(SupervisorJob())
+        uiScope = CoroutineScope(SupervisorJob())
+        rpc = FakeAppRpcApi()
+        workspaceRpc = FakeWorkspaceRpcApi()
+        app = KiloAppService(appScope, rpc)
+        workspaces = KiloWorkspaceService(appScope, workspaceRpc)
+    }
+
+    override fun tearDown() {
+        try {
+            val view = panel
+            if (view != null) edt { view.dispose() }
+            panel = null
+            uiScope.cancel()
+            appScope.cancel()
+        } finally {
+            super.tearDown()
+        }
+    }
+
+    fun `test startup accepts ready app state and loads resolved workspace`() {
+        rpc.state.value = state("new")
+        workspaceRpc.directory = "/resolved"
+        val view = create("/hint")
+
+        flushUntil { edt { view.value() == "new" && view.roots == listOf("/resolved") } }
+
+        edt {
+            assertEquals("new", view.value())
+            assertEquals("/resolved", view.dir())
+            assertTrue(view.loaded())
+            assertFalse(view.loading())
+            assertFalse(view.loadOnEdt)
+        }
+    }
+
+    fun `test non ready app state calls unavailable hook`() {
+        rpc.state.value = state("ready")
+        val view = create("/test")
+        flushUntil { edt { view.value() == "ready" } }
+        val before = edt { view.unavailable }
+
+        rpc.state.value = KiloAppStateDto(KiloAppStatusDto.DISCONNECTED)
+        flushUntil { edt { view.unavailable > before } }
+
+        edt { assertFalse(view.loading()) }
+    }
+
+    fun `test model state updates are delivered on edt`() {
+        rpc.models = ModelStateDto(favorite = listOf(ModelSelectionDto("kilo", "new")))
+        rpc.state.value = state("ready")
+        val view = create("/test")
+
+        flushUntil { edt { view.favoriteCount == 1 } }
+
+        edt {
+            assertEquals(1, view.favoriteCount)
+            assertTrue(view.modelsOnEdt)
+        }
+    }
+
+    private fun create(hint: String): FakePanel {
+        val view = edt { FakePanel(uiScope, app, workspaces, hint) }
+        panel = view
+        return view
+    }
+
+    private fun state(model: String) = KiloAppStateDto(
+        KiloAppStatusDto.READY,
+        config = ConfigDto(model = model),
+    )
+
+    private fun <T> edt(block: () -> T): T {
+        var result: T? = null
+        ApplicationManager.getApplication().invokeAndWait { result = block() }
+        @Suppress("UNCHECKED_CAST")
+        return result as T
+    }
+
+    private fun flushUntil(done: () -> Boolean) = runBlocking {
+        repeat(20) {
+            delay(100)
+            edt { UIUtil.dispatchAllInvocationEvents() }
+            if (done()) return@runBlocking
+        }
+        edt { UIUtil.dispatchAllInvocationEvents() }
+        assertTrue(done())
+    }
+
+    private data class Draft(val value: String)
+    private data class Change(val value: String)
+
+    private class FakeContent : BaseContentPanel()
+
+    private class FakePanel(
+        cs: CoroutineScope,
+        app: KiloAppService,
+        workspaces: KiloWorkspaceService,
+        hint: String,
+    ) : BaseWorkspaceSettingsUi<FakeContent, Draft, Change, Draft, String>(
+        cs,
+        Draft("old"),
+        app,
+        workspaces,
+        hint,
+    ) {
+        val roots = mutableListOf<String>()
+        var unavailable = 0
+            private set
+        var favoriteCount = 0
+            private set
+        var loadOnEdt = true
+            private set
+        var modelsOnEdt = false
+            private set
+
+        init {
+            startSettings(FakeContent())
+        }
+
+        fun value(): String = draft.value
+
+        fun dir(): String? = projectDirectory
+
+        fun loading(): Boolean = workspaceLoading
+
+        fun loaded(): Boolean = workspaceLoaded
+
+        override fun change(from: Draft, to: Draft): Change? = if (from == to) null else Change(to.value)
+
+        override fun save(change: Change, done: (Draft?) -> Unit) = done(Draft(change.value))
+
+        override fun base(result: Draft): Draft = result
+
+        override fun draft(state: KiloAppStateDto): Draft = Draft(state.config?.model ?: "none")
+
+        override suspend fun loadWorkspace(root: String): String {
+            loadOnEdt = ApplicationManager.getApplication().isDispatchThread
+            roots += root
+            return root
+        }
+
+        override fun applyWorkspace(result: String) = Unit
+
+        override fun unavailable(state: KiloAppStateDto) {
+            unavailable++
+        }
+
+        override fun models(state: ModelStateDto) {
+            favoriteCount = state.favorite.size
+            modelsOnEdt = ApplicationManager.getApplication().isDispatchThread
+        }
+
+        override fun syncContent() = Unit
+
+        override fun pendingText(): String = "Saving"
+
+        override fun failedText(): String = "Failed"
+    }
+}