fix(jetbrains): simplify new worktree dialog

packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/agentManager/AgentManagerPanel.kt

@@ -1,7 +1,6 @@
 package ai.kilocode.client.agentManager
 
 import ai.kilocode.client.KiloNotifications
-import ai.kilocode.client.app.KiloSessionService
 import ai.kilocode.client.agentManager.worktree.NewWorktreeDialog
 import ai.kilocode.client.agentManager.worktree.WorktreeController
 import ai.kilocode.client.agentManager.worktree.WorktreeDataKeys
@@ -145,7 +144,7 @@ class AgentManagerPanel(
         project?.service<WorktreeStatusService>()?.refreshPr()
     }
 
-    /** Opens the New Worktree dialog (New + Import tabs). */
+    /** Opens the New Worktree dialog. */
     fun configure() {
         val target = project ?: return
         NewWorktreeDialog(
@@ -158,9 +157,6 @@ class AgentManagerPanel(
             onCreate = { branch, base, prompt ->
                 controller.create(branch, base, prompt = prompt)
             },
-            onImportPr = { url -> controller.importPr(url) },
-            onImportBranch = { branch -> controller.importBranch(branch) },
-            sessions = target.service<KiloSessionService>(),
         ).show()
     }
 

packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/agentManager/worktree/NewWorktreeDialog.kt

@@ -1,7 +1,6 @@
 package ai.kilocode.client.agentManager.worktree
 
 import ai.kilocode.client.app.KiloAppService
-import ai.kilocode.client.app.KiloSessionService
 import ai.kilocode.client.app.KiloWorkspaceService
 import ai.kilocode.client.plugin.KiloBundle
 import ai.kilocode.client.session.ui.ReasoningPicker
@@ -13,6 +12,7 @@ import ai.kilocode.client.session.ui.prompt.MentionAction
 import ai.kilocode.client.session.ui.prompt.PromptPanel
 import ai.kilocode.client.session.ui.prompt.SlashAction
 import ai.kilocode.client.ui.UiStyle
+import ai.kilocode.client.ui.layout.Stack
 import ai.kilocode.rpc.dto.ModelsWorkspaceDto
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.application.ModalityState
@@ -20,39 +20,30 @@ 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.components.JBLabel
-import com.intellij.ui.components.JBTabbedPane
 import com.intellij.ui.components.JBTextField
 import com.intellij.util.ui.FormBuilder
 import com.intellij.util.ui.JBUI
-import com.intellij.util.ui.UIUtil
-import com.intellij.util.ui.components.BorderLayoutPanel
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.SupervisorJob
 import kotlinx.coroutines.cancel
 import kotlinx.coroutines.launch
 import java.awt.Component
-import java.awt.event.ActionEvent
-import javax.swing.AbstractAction
+import java.awt.GridBagConstraints
 import javax.swing.DefaultComboBoxModel
-import javax.swing.JButton
 import javax.swing.JComponent
 
 private const val NAME_COLUMNS = 100
 
 /**
- * New Worktree dialog with parity to the VS Code Agent Manager dialog:
+ * New Worktree dialog with parity to the VS Code Agent Manager dialog: a worktree name (top), an
+ * initial prompt with the same mode / model / reasoning pickers as the chat prompt (center), and the
+ * branch name + base branch (bottom). Creating a worktree starts a session automatically with the
+ * prompt.
  *
- * - **New** tab: a worktree name (top), an initial prompt with the same mode / model / reasoning
- *   pickers as the chat prompt (center), and the branch name + base branch (bottom). Creating a
- *   worktree here starts a session automatically with the prompt.
- * - **Import** tab: import a worktree from a GitHub pull request URL or from an existing branch.
- *
- * The dialog performs no worktree work itself — it invokes [onCreate], [onImportPr], or
- * [onImportBranch] and closes; the panel drives the controller. Mode, model, and reasoning
- * selections are persisted the same way the chat prompt does, so the freshly-started session
- * inherits them.
+ * The dialog performs no worktree work itself — it invokes [onCreate] and closes; the panel drives
+ * the controller. Mode, model, and reasoning selections are persisted the same way the chat prompt
+ * does, so the freshly-started session inherits them.
  */
 internal class NewWorktreeDialog(
     parent: Component,
@@ -62,12 +53,8 @@ internal class NewWorktreeDialog(
     private val defaultBase: String,
     private val branches: List<String>,
     private val onCreate: (branch: String, base: String?, prompt: PendingPrompt?) -> Unit,
-    private val onImportPr: (url: String) -> Unit,
-    private val onImportBranch: (branch: String) -> Unit,
     private val app: KiloAppService = service(),
     private val workspaces: KiloWorkspaceService = service(),
-    // Project-scoped; supplied by the panel. When absent, mode selection is not persisted.
-    private val sessions: KiloSessionService? = null,
 ) : DialogWrapper(parent, false) {
 
     private val scope = CoroutineScope(SupervisorJob() + Dispatchers.Default)
@@ -87,22 +74,18 @@ internal class NewWorktreeDialog(
         project = project,
         onSend = { text, _ -> submitCreate(text) },
         onAbort = {},
-        onEnhance = ::enhance,
         completion = completion,
         cs = scope,
         rounded = false,
         showSubmit = false,
         approve = false,
+        showEnhance = false,
     )
     private val branch = JBTextField(suggestedName)
     private val base = ComboBox(baseModel(branches, defaultBase)).apply {
         isEditable = true
         selectedItem = defaultBase
     }
-    private val prUrl = JBTextField().apply {
-        emptyText.text = KiloBundle.message("worktree.import.pr.placeholder")
-    }
-    private val branchPicker = ComboBox(importBranchModel(branches))
 
     /** The agent (mode) for the new session; model selections persist against it. */
     private var agent: String? = null
@@ -117,30 +100,22 @@ internal class NewWorktreeDialog(
     private var disposed = false
 
     private var center: JComponent? = null
-    private var tabs: JBTabbedPane? = null
 
     init {
         title = KiloBundle.message("worktree.configure.title")
         init()
         setOKButtonText(KiloBundle.message("worktree.dialog.create"))
-        syncOkEnabled()
     }
 
-    override fun createCenterPanel(): JComponent = JBTabbedPane().apply {
-        addTab(KiloBundle.message("worktree.dialog.tab.new"), newTab())
-        addTab(KiloBundle.message("worktree.dialog.tab.import"), importTab())
-        addChangeListener { syncOkEnabled() }
-        tabs = this
-    }.also { center = it }
+    override fun createCenterPanel(): JComponent = content().also { center = it }
 
-    /** The built tabbed content, so tests can drive the real Swing tree before the dialog is shown. */
+    /** The built content, so tests can drive the real Swing tree before the dialog is shown. */
     internal fun centerComponent(): JComponent = center ?: error("center panel not built")
 
     override fun getPreferredFocusedComponent(): JComponent = prompt.defaultFocusedComponent
 
     override fun getDimensionServiceKey(): String = "ai.kilocode.NewWorktreeDialog"
 
-    // "Create Worktree" is the dialog's default action; it applies only to the New tab.
     override fun doOKAction() = submitCreate()
 
     override fun dispose() {
@@ -149,22 +124,25 @@ internal class NewWorktreeDialog(
         super.dispose()
     }
 
-    private fun newTab(): JComponent {
+    private fun content(): JComponent {
         wirePickers()
-        val south = FormBuilder.createFormBuilder()
-            .addLabeledComponent(KiloBundle.message("worktree.configure.branch"), branch)
-            .addLabeledComponent(KiloBundle.message("worktree.configure.base"), base)
-            .panel
-
-        val root = BorderLayoutPanel()
-        root.border = JBUI.Borders.empty(UiStyle.Gap.sm())
-        root.addToTop(name)
-        root.addToCenter(prompt)
-        root.addToBottom(south)
         loadModels()
-        return root
+        return Stack.vertical(gap = UiStyle.Gap.pad())
+            .next(name)
+            .next(prompt)
+            .next(fields())
+            .apply { border = JBUI.Borders.empty(UiStyle.Gap.sm()) }
     }
 
+    // A FormBuilder that stretches every field to the full width, so the base-branch combo matches
+    // the name field and prompt above it.
+    private fun fields(): JComponent = object : FormBuilder() {
+        override fun getFill(component: JComponent) = GridBagConstraints.HORIZONTAL
+    }
+        .addLabeledComponent(KiloBundle.message("worktree.configure.branch"), branch)
+        .addLabeledComponent(KiloBundle.message("worktree.configure.base"), base)
+        .panel
+
     private fun wirePickers() {
         prompt.mode.onSelect = { item -> selectAgent(item.id) }
         prompt.model.favorites = { app.favorites.value }
@@ -221,37 +199,7 @@ internal class NewWorktreeDialog(
         )
     }
 
-    private fun importTab(): JComponent {
-        prUrl.addActionListener { submitImportPr() }
-        val open = createButton(KiloBundle.message("worktree.import.pr.open")) { submitImportPr() }
-        val prRow = BorderLayoutPanel().apply {
-            addToCenter(prUrl)
-            addToRight(open)
-        }
-        branchPicker.addActionListener {
-            val value = branchPicker.selectedItem as? String ?: return@addActionListener
-            if (value.isBlank() || value == importBranchPlaceholder()) return@addActionListener
-            onImportBranch(value)
-            close(OK_EXIT_CODE)
-        }
-
-        val builder = FormBuilder.createFormBuilder()
-            .addComponent(sectionLabel(KiloBundle.message("worktree.import.pr.section")))
-            .addComponent(prRow)
-            .addSeparator()
-            .addComponent(sectionLabel(KiloBundle.message("worktree.import.branches.section")))
-        if (branches.isEmpty()) {
-            builder.addComponent(sectionLabel(KiloBundle.message("worktree.import.branches.empty")))
-        } else {
-            builder.addComponent(branchPicker)
-        }
-        val panel = builder.panel
-        panel.border = JBUI.Borders.empty(UiStyle.Gap.sm())
-        return panel
-    }
-
     private fun submitCreate(text: String = prompt.text()) {
-        if (tabs?.selectedIndex != 0) return
         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))
@@ -272,18 +220,9 @@ internal class NewWorktreeDialog(
         )
     }
 
-    private fun enhance(text: String, done: (Result<String>) -> Unit) {
-        val svc = sessions ?: return done(Result.failure(IllegalStateException("Session service unavailable")))
-        scope.launch {
-            val result = runCatching { svc.enhancePrompt(directory, text) }
-            ui { if (!project.isDisposed) done(result) }
-        }
-    }
-
     // The dialog is modal, so its EDT runs a nested event loop. A plain invokeLater carries the
     // caller's (non-modal) modality and would be deferred until the dialog closes, leaving the
-    // pickers empty and enhance disabled. ModalityState.any() lets these UI-only updates run while
-    // the dialog is showing.
+    // pickers empty. ModalityState.any() lets these UI-only updates run while the dialog is showing.
     private fun ui(block: () -> Unit) {
         ApplicationManager.getApplication().invokeLater({ if (!disposed) block() }, ModalityState.any())
     }
@@ -306,26 +245,6 @@ internal class NewWorktreeDialog(
         spec.available,
     )
 
-    private fun submitImportPr() {
-        val url = prUrl.text.trim()
-        if (url.isEmpty()) return
-        onImportPr(url)
-        close(OK_EXIT_CODE)
-    }
-
-    /** The New-tab create action is the default button; disable it while the Import tab is active. */
-    private fun syncOkEnabled() {
-        getOKAction().isEnabled = tabs?.selectedIndex == 0
-    }
-
-    private fun sectionLabel(text: String) = JBLabel(text).apply {
-        foreground = UIUtil.getContextHelpForeground()
-    }
-
-    private fun createButton(text: String, run: () -> Unit) = JButton(object : AbstractAction(text) {
-        override fun actionPerformed(e: ActionEvent) = run()
-    })
-
     private fun baseModel(branches: List<String>, default: String): DefaultComboBoxModel<String> {
         val ordered = LinkedHashSet<String>()
         if (default.isNotBlank()) ordered.add(default)
@@ -333,13 +252,5 @@ internal class NewWorktreeDialog(
         return DefaultComboBoxModel(ordered.toTypedArray())
     }
 
-    private fun importBranchModel(branches: List<String>): DefaultComboBoxModel<String> {
-        val ordered = mutableListOf(importBranchPlaceholder())
-        ordered.addAll(branches)
-        return DefaultComboBoxModel(ordered.toTypedArray())
-    }
-
-    private fun importBranchPlaceholder() = KiloBundle.message("worktree.import.branches.placeholder")
-
     private fun variantTitle(value: String): String = value.replaceFirstChar { it.titlecase() }
 }

packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/prompt/PromptPanel.kt

@@ -108,7 +108,7 @@ class PromptPanel(
     private val project: Project,
     private val onSend: (String, List<PromptPartDto>) -> Unit,
     private val onAbort: () -> Unit,
-    private val onEnhance: (String, (Result<String>) -> Unit) -> Unit,
+    private val onEnhance: (String, (Result<String>) -> Unit) -> Unit = { _, _ -> },
     private val onMentions: suspend (String) -> List<PromptPartDto> = { emptyList() },
     private val completion: KiloPromptCompletionProvider? = null,
     private val selection: SessionSelection? = null,
@@ -116,6 +116,7 @@ class PromptPanel(
     private val rounded: Boolean = true,
     private val showSubmit: Boolean = true,
     private val approve: Boolean = true,
+    private val showEnhance: Boolean = true,
 ) : BorderLayoutPanel(), SessionEditorStyleTarget, SendPromptContext, UiDataProvider {
 
     companion object {
@@ -315,7 +316,7 @@ class PromptPanel(
             bar.add(auto)
             bar.add(Box.createHorizontalStrut(JBUI.scale(SessionUiStyle.View.Prompt.CONTROL_GAP)))
         }
-        bar.add(enhance)
+        if (showEnhance) bar.add(enhance)
         if (showSubmit) {
             bar.add(Box.createHorizontalStrut(JBUI.scale(SessionUiStyle.View.Prompt.CONTROL_GAP)))
             bar.add(separator)

packages/kilo-jetbrains/frontend/src/main/resources/messages/KiloBundle.properties

@@ -368,18 +368,11 @@ worktree.configure.title=New Worktree
 worktree.configure.branch=Branch name:
 worktree.configure.base=Base branch:
 worktree.configure.branch.required=Branch name is required
-worktree.dialog.tab.new=New
-worktree.dialog.tab.import=Import
 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
 worktree.dialog.creating=Creating…
 worktree.import.pr.section=Pull Request
-worktree.import.pr.placeholder=Paste PR URL...
-worktree.import.pr.open=Open
-worktree.import.branches.section=Branches
-worktree.import.branches.placeholder=Select branch...
-worktree.import.branches.empty=No other branches to import.
 worktree.import.pr.failed.title=Couldn''t import pull request
 worktree.import.branch.failed.title=Couldn''t import branch "{0}"
 worktree.stats.diff.tooltip={0} additions, {1} deletions

packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/agentManager/worktree/NewWorktreeDialogTest.kt

@@ -1,7 +1,6 @@
 package ai.kilocode.client.agentManager.worktree
 
 import ai.kilocode.client.app.KiloAppService
-import ai.kilocode.client.app.KiloSessionService
 import ai.kilocode.client.app.KiloWorkspaceService
 import ai.kilocode.client.session.ui.ReasoningPicker
 import ai.kilocode.client.session.ui.mode.ModePicker
@@ -35,7 +34,6 @@ class NewWorktreeDialogTest : BasePlatformTestCase() {
     private lateinit var app: KiloAppService
     private lateinit var workspaces: KiloWorkspaceService
     private lateinit var sessionRpc: FakeSessionRpcApi
-    private lateinit var sessions: KiloSessionService
     private var dialog: NewWorktreeDialog? = null
     private val created = mutableListOf<Triple<String, String?, PendingPrompt?>>()
 
@@ -46,7 +44,6 @@ class NewWorktreeDialogTest : BasePlatformTestCase() {
         val ws = FakeWorkspaceRpcApi().apply { models = workspace() }
         workspaces = KiloWorkspaceService(scope, ws)
         sessionRpc = FakeSessionRpcApi()
-        sessions = KiloSessionService(project, scope, sessionRpc)
     }
 
     override fun tearDown() {
@@ -138,11 +135,8 @@ class NewWorktreeDialogTest : BasePlatformTestCase() {
                 "main",
                 listOf("main"),
                 onCreate = { branch, base, prompt -> created.add(Triple(branch, base, prompt)) },
-                onImportPr = {},
-                onImportBranch = {},
                 app,
                 workspaces,
-                sessions,
             )
         }
     }