fix(jetbrains): improve worktree base branch selector

2026-09-21 05:52:35 +08:00 · 2026-08-11 19:48:52 -04:00
parent d17cd80cc3
commit f5ff70ea3b
6 changed files with 236 additions and 43 deletions
@@ -0,0 +1,5 @@
+---
+"@kilocode/kilo-jetbrains": patch
+---
+
+Improve the New Worktree base branch selector with fuzzy matching, default fallback, and validation for unknown branches.
@@ -1,5 +1,6 @@
 package ai.kilocode.client.agentManager.worktree

+import ai.kilocode.client.KiloNotifications
 import ai.kilocode.client.app.KiloAppService
 import ai.kilocode.client.app.KiloWorkspaceService
 import ai.kilocode.client.plugin.KiloBundle
@@ -9,6 +10,7 @@ import ai.kilocode.client.session.ui.model.ModelPicker
 import ai.kilocode.client.session.ui.model.modelItems
 import ai.kilocode.client.session.ui.prompt.KiloPromptCompletionProvider
 import ai.kilocode.client.session.ui.prompt.MentionAction
+import ai.kilocode.client.session.ui.prompt.PromptFuzzyRanker
 import ai.kilocode.client.session.ui.prompt.PromptPanel
 import ai.kilocode.client.session.ui.prompt.SlashAction
 import ai.kilocode.client.ui.UiStyle
@@ -20,6 +22,7 @@ import com.intellij.openapi.components.service
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.ui.ComboBox
 import com.intellij.openapi.ui.DialogWrapper
+import com.intellij.ui.DocumentAdapter
 import com.intellij.ui.components.JBTextField
 import com.intellij.util.ui.FormBuilder
 import com.intellij.util.ui.JBUI
@@ -30,8 +33,14 @@ import kotlinx.coroutines.cancel
 import kotlinx.coroutines.launch
 import java.awt.Component
 import java.awt.GridBagConstraints
+import java.awt.event.FocusAdapter
+import java.awt.event.FocusEvent
+import javax.swing.ComboBoxModel
 import javax.swing.DefaultComboBoxModel
 import javax.swing.JComponent
+import javax.swing.JTextField
+import javax.swing.event.DocumentEvent
+import javax.swing.plaf.basic.BasicComboPopup

 private const val NAME_COLUMNS = 100

@@ -82,10 +91,13 @@ internal class NewWorktreeDialog(
        showEnhance = false,
    )
    private val branch = JBTextField(suggestedName)
-    private val base = ComboBox(baseModel(branches, defaultBase)).apply {
+    private val bases = baseBranches(branches, defaultBase)
+    private val baseSet = bases.toSet()
+    private val base = ComboBox(baseModel(bases)).apply {
        isEditable = true
        selectedItem = defaultBase
    }
+    private var syncing = false

    /** The agent (mode) for the new session; model selections persist against it. */
    private var agent: String? = null
@@ -102,6 +114,7 @@ internal class NewWorktreeDialog(
    private var center: JComponent? = null

    init {
+        wireBase()
        title = KiloBundle.message("worktree.configure.title")
        init()
        setOKButtonText(KiloBundle.message("worktree.dialog.create"))
@@ -199,10 +212,96 @@ internal class NewWorktreeDialog(
        )
    }

+    private fun wireBase() {
+        val field = baseField() ?: return
+        field.document.addDocumentListener(object : DocumentAdapter() {
+            override fun textChanged(e: DocumentEvent) {
+                if (!syncing) syncBase(field.text, popup = true)
+            }
+        })
+        field.addFocusListener(object : FocusAdapter() {
+            override fun focusLost(e: FocusEvent) {
+                restoreBase()
+            }
+        })
+    }
+
+    private fun restoreBase() {
+        if (baseText().isNotEmpty() || defaultBase.isBlank()) return
+        setBase(defaultBase)
+    }
+
+    private fun syncBase(text: String, popup: Boolean) {
+        val value = text.trim()
+        if (value.isEmpty()) return
+        if (popup && base.isShowing && !base.isPopupVisible) {
+            base.isPopupVisible = true
+        }
+        val idx = matchBase(value) ?: return
+        val list = popupList() ?: return
+        if (list.selectedIndex != idx) list.selectedIndex = idx
+        list.ensureIndexIsVisible(idx)
+    }
+
+    private fun matchBase(text: String): Int? {
+        val rank = PromptFuzzyRanker(text)
+        return bases.withIndex().mapNotNull { item ->
+            rank.score(item.value, emptyList())?.let { score -> item.index to score }
+        }.maxByOrNull { it.second }?.first
+    }
+
+    private fun popupList() = (base.accessibleContext?.getAccessibleChild(0) as? BasicComboPopup)?.list
+
+    private fun baseField() = base.editor.editorComponent as? JTextField
+
+    private fun baseText() = baseField()?.text?.trim()
+        ?: base.editor.item?.toString()?.trim().orEmpty()
+
+    private fun setBase(value: String) {
+        syncing = true
+        try {
+            base.selectedItem = value
+            baseField()?.text = value
+        } finally {
+            syncing = false
+        }
+    }
+
+    private fun resolvedBase(): String? {
+        val value = baseText()
+        if (value.isEmpty()) {
+            val fallback = defaultBase.trim()
+            if (fallback.isNotEmpty()) setBase(fallback)
+            return fallback.takeIf { it.isNotEmpty() }
+        }
+        if (value in baseSet) return value
+        val idx = matchBase(value) ?: return value
+        val target = bases[idx]
+        setBase(target)
+        return target
+    }
+
+    private fun validBase(value: String?): Boolean {
+        if (value == null || value in baseSet) return true
+        KiloNotifications.error(
+            project,
+            KiloBundle.message("worktree.configure.base.invalid.title"),
+            KiloBundle.message("worktree.configure.base.invalid.content", value),
+        )
+        baseField()?.apply {
+            requestFocusInWindow()
+            selectAll()
+        }
+        syncBase(value, popup = true)
+        return false
+    }
+
    private fun submitCreate(text: String = prompt.text()) {
        val explicit = branch.text.trim()
        val resolved = explicit.ifEmpty { name.text.trim() }.ifEmpty { suggestedName }
-        onCreate(resolved, base.editor.item?.toString()?.trim()?.takeIf { it.isNotEmpty() }, pending(text))
+        val target = resolvedBase()
+        if (!validBase(target)) return
+        onCreate(resolved, target, pending(text))
        close(OK_EXIT_CODE)
    }

@@ -245,11 +344,15 @@ internal class NewWorktreeDialog(
        spec.available,
    )

-    private fun baseModel(branches: List<String>, default: String): DefaultComboBoxModel<String> {
+    private fun baseBranches(branches: List<String>, default: String): List<String> {
        val ordered = LinkedHashSet<String>()
        if (default.isNotBlank()) ordered.add(default)
        ordered.addAll(branches)
-        return DefaultComboBoxModel(ordered.toTypedArray())
+        return ordered.toList()
+    }
+
+    private fun baseModel(branches: List<String>): ComboBoxModel<String> {
+        return DefaultComboBoxModel(branches.toTypedArray())
    }

    private fun variantTitle(value: String): String = value.replaceFirstChar { it.titlecase() }
@@ -20,10 +20,7 @@ import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.fileTypes.FileTypeManager
 import com.intellij.openapi.project.DumbAware
 import com.intellij.openapi.progress.runBlockingCancellable
-import com.intellij.psi.codeStyle.MinusculeMatcher
-import com.intellij.psi.codeStyle.NameUtil
 import com.intellij.util.textCompletion.TextCompletionProvider
-import com.intellij.util.text.matching.MatchingMode
 import java.util.Collections
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.launch
@@ -166,7 +163,7 @@ class KiloPromptCompletionProvider(
    private fun slash(prefix: String, result: CompletionResultSet) {
        result.restartCompletionOnAnyPrefixChange()
        val out = result.withPrefixMatcher(PlainPrefixMatcher.ALWAYS_TRUE)
-        val rank = Ranker(prefix)
+        val rank = PromptFuzzyRanker(prefix)
        val names = clientTokens()
        val clients = actions.mapNotNull { action ->
            rank.score(action.name, action.hints)?.let { Hit(client(action), it) }
@@ -187,7 +184,7 @@ class KiloPromptCompletionProvider(
        result.restartCompletionOnAnyPrefixChange()
        val out = result.withPrefixMatcher(PlainPrefixMatcher.ALWAYS_TRUE)
        val search = search(prefix)
-        val rank = Ranker(prefix)
+        val rank = PromptFuzzyRanker(prefix)
        val known = mentions.filter { action -> rank.matches(action.name, action.hints) && action.available(search) }
        known.forEach { action -> out.addElement(prioritize(resource(action))) }
        if (search.indexing) {
@@ -224,38 +221,6 @@ class KiloPromptCompletionProvider(
        }
        .withAutoCompletionPolicy(AutoCompletionPolicy.NEVER_AUTOCOMPLETE)

-    private class Ranker(prefix: String) {
-        private val start = matcher(prefix)
-        private val middle = if (prefix.any { separator(it) }) null else matcher("*$prefix")
-
-        fun matches(name: String, hints: List<String>): Boolean = score(name, hints) != null
-
-        fun score(name: String, hints: List<String>): Int? = (listOf(name) + hints).maxOfOrNull { value ->
-            score(value) ?: Int.MIN_VALUE
-        }?.takeIf { it != Int.MIN_VALUE }
-
-        private fun score(value: String): Int? {
-            val exact = start.match(value)
-            if (exact != null) return START + start.matchingDegree(value, true, exact)
-            val fallback = middle ?: return null
-            val fuzzy = fallback.match(value) ?: return null
-            return fallback.matchingDegree(value, false, fuzzy)
-        }
-
-        private companion object {
-            const val START = 10_000
-
-            fun matcher(prefix: String): MinusculeMatcher = NameUtil.buildMatcher(prefix)
-                .withMatchingMode(MatchingMode.IGNORE_CASE)
-                .build()
-
-            fun separator(c: Char): Boolean = when (c) {
-                '_', '-', ':', '+', '.' -> true
-                else -> c.isWhitespace()
-            }
-        }
-    }
-
    private fun commandName(text: String): String? {
        val raw = text.trimStart()
        if (!raw.startsWith('/')) return null
@@ -0,0 +1,37 @@
+package ai.kilocode.client.session.ui.prompt
+
+import com.intellij.psi.codeStyle.MinusculeMatcher
+import com.intellij.psi.codeStyle.NameUtil
+import com.intellij.util.text.matching.MatchingMode
+
+internal class PromptFuzzyRanker(prefix: String) {
+    private val start = matcher(prefix)
+    private val middle = if (prefix.any { separator(it) }) null else matcher("*$prefix")
+
+    fun matches(name: String, hints: List<String>): Boolean = score(name, hints) != null
+
+    fun score(name: String, hints: List<String>): Int? = (listOf(name) + hints).maxOfOrNull { value ->
+        score(value) ?: Int.MIN_VALUE
+    }?.takeIf { it != Int.MIN_VALUE }
+
+    private fun score(value: String): Int? {
+        val exact = start.match(value)
+        if (exact != null) return START + start.matchingDegree(value, true, exact)
+        val fallback = middle ?: return null
+        val fuzzy = fallback.match(value) ?: return null
+        return fallback.matchingDegree(value, false, fuzzy)
+    }
+
+    private companion object {
+        const val START = 10_000
+
+        fun matcher(prefix: String): MinusculeMatcher = NameUtil.buildMatcher(prefix)
+            .withMatchingMode(MatchingMode.IGNORE_CASE)
+            .build()
+
+        fun separator(c: Char): Boolean = when (c) {
+            '_', '-', ':', '+', '.' -> true
+            else -> c.isWhitespace()
+        }
+    }
+}
@@ -368,6 +368,8 @@ worktree.configure.title=New Worktree
 worktree.configure.branch=Branch name:
 worktree.configure.base=Base branch:
 worktree.configure.branch.required=Branch name is required
+worktree.configure.base.invalid.title=Base branch not found
+worktree.configure.base.invalid.content=Select an existing base branch before creating a worktree: {0}
 worktree.dialog.name.placeholder=Worktree name (optional)
 worktree.dialog.prompt.placeholder=Describe what you want to start working on ({0} to create)
 worktree.dialog.create=Create Worktree
@@ -17,6 +17,7 @@ import ai.kilocode.rpc.dto.ModelSelectionDto
 import ai.kilocode.rpc.dto.ModelsWorkspaceDto
 import ai.kilocode.rpc.dto.ProviderDto
 import ai.kilocode.rpc.dto.ProvidersDto
+import com.intellij.openapi.ui.ComboBox
 import com.intellij.openapi.util.Disposer
 import com.intellij.testFramework.fixtures.BasePlatformTestCase
 import com.intellij.ui.components.JBPanel
@@ -28,6 +29,9 @@ import kotlinx.coroutines.delay
 import kotlinx.coroutines.runBlocking
 import java.awt.Component
 import java.awt.Container
+import java.awt.event.FocusEvent
+import javax.swing.JTextField
+import javax.swing.plaf.basic.BasicComboPopup

 class NewWorktreeDialogTest : BasePlatformTestCase() {
    private lateinit var scope: CoroutineScope
@@ -125,7 +129,71 @@ class NewWorktreeDialogTest : BasePlatformTestCase() {
        assertEquals("gpt-5", payload.model)
    }

-    private fun open() {
+    fun `test base branch fuzzy search selects matching popup item`() {
+        open(branches = listOf("main", "release/candidate", "feature/refactor-ui"))
+
+        edt { field().text = "relcan" }
+
+        edt { assertEquals("release/candidate", popup().list.selectedValue) }
+    }
+
+    fun `test empty base branch restores default on focus lost`() {
+        open()
+
+        edt {
+            val field = field()
+            field.text = ""
+            field.focusListeners.forEach { it.focusLost(FocusEvent(field, FocusEvent.FOCUS_LOST)) }
+
+            assertEquals("main", field.text)
+        }
+    }
+
+    fun `test creating with empty base branch falls back to default`() {
+        open()
+        flushUntil { edt { model().selectionKeyForTest() != null } }
+        edt {
+            field().text = ""
+            prompt().setText("build the thing")
+        }
+        flushUntil { edt { prompt().isSendEnabled } }
+        edt { prompt().send() }
+        flushUntil { created.isNotEmpty() }
+        dialog = null
+
+        assertEquals("main", created.single().second)
+    }
+
+    fun `test creating with fuzzy base branch uses matching branch`() {
+        open(branches = listOf("main", "release/candidate", "feature/refactor-ui"))
+        flushUntil { edt { model().selectionKeyForTest() != null } }
+        edt {
+            field().text = "relcan"
+            prompt().setText("build the thing")
+        }
+        flushUntil { edt { prompt().isSendEnabled } }
+        edt { prompt().send() }
+        flushUntil { created.isNotEmpty() }
+        dialog = null
+
+        assertEquals("release/candidate", created.single().second)
+    }
+
+    fun `test creating with unknown base branch does not create`() {
+        open(branches = listOf("main", "release/candidate"))
+        flushUntil { edt { model().selectionKeyForTest() != null } }
+        edt {
+            field().text = "zzzzzz"
+            prompt().setText("build the thing")
+        }
+        flushUntil { edt { prompt().isSendEnabled } }
+        edt { prompt().send() }
+        flush()
+
+        assertTrue(created.isEmpty())
+    }
+
+    private fun open(branches: List<String> = listOf("main")) {
        dialog = edt {
            NewWorktreeDialog(
                JBPanel<Nothing>(),
@@ -133,7 +201,7 @@ class NewWorktreeDialogTest : BasePlatformTestCase() {
                "/test",
                "agent/foo",
                "main",
-                listOf("main"),
+                branches,
                onCreate = { branch, base, prompt -> created.add(Triple(branch, base, prompt)) },
                app,
                workspaces,
@@ -167,6 +235,12 @@ class NewWorktreeDialogTest : BasePlatformTestCase() {

    private fun prompt(): PromptPanel = descendants(root()).filterIsInstance<PromptPanel>().single()

+    private fun combo(): ComboBox<*> = descendants(root()).filterIsInstance<ComboBox<*>>().single()
+
+    private fun field(): JTextField = combo().editor.editorComponent as JTextField
+
+    private fun popup(): BasicComboPopup = combo().accessibleContext.getAccessibleChild(0) as BasicComboPopup
+
    private fun root(): Component = requireNotNull(dialog).centerComponent()

    private fun descendants(root: Component): List<Component> {
@@ -181,6 +255,13 @@ class NewWorktreeDialogTest : BasePlatformTestCase() {

    private fun <T> edt(block: () -> T): T = edtWait(block)

+    private fun flush() = runBlocking {
+        repeat(20) {
+            delay(10)
+            edt { UIUtil.dispatchAllInvocationEvents() }
+        }
+    }
+
    private fun flushUntil(done: () -> Boolean) = runBlocking {
        repeat(200) {
            delay(10)