fix(jetbrains): start worktree session with picked mode, model, and effort

The New Worktree dialog only stashed the prompt text, so the freshly-opened worktree session resolved its own default mode/model and the first turn could fail to start. Carry the picked agent/provider/model/reasoning with the initial prompt and seed them as the session's preferred selection so the pickers and every later turn use the dialog's choice. - Extend PendingWorktreePrompt to carry the selection alongside the text - Thread the selection through WorktreeController.create and the dialog - Add SessionController.applySelection + PromptSelection, applied on submit - Drop the dialog's global default_agent write that changed the mode for every other workspace and raced the new session's model load
2026-09-24 16:02:55 +08:00 · 2026-08-11 17:12:54 -04:00
parent 579245d8f2
commit 9e09eafb8e
11 changed files with 224 additions and 44 deletions
@@ -156,7 +156,7 @@ class AgentManagerPanel(
            controller.defaultBranch,
            controller.branches,
            onCreate = { branch, base, prompt ->
-                controller.create(branch, base, prompt = prompt.takeIf { it.isNotEmpty() })
+                controller.create(branch, base, prompt = prompt)
            },
            onImportPr = { url -> controller.importPr(url) },
            onImportBranch = { branch -> controller.importBranch(branch) },
@@ -13,7 +13,6 @@ 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.rpc.dto.ConfigUpdateDto
 import ai.kilocode.rpc.dto.ModelsWorkspaceDto
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.application.ModalityState
@@ -62,7 +61,7 @@ internal class NewWorktreeDialog(
    private val suggestedName: String,
    private val defaultBase: String,
    private val branches: List<String>,
-    private val onCreate: (branch: String, base: String?, prompt: String) -> Unit,
+    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(),
@@ -203,8 +202,10 @@ internal class NewWorktreeDialog(
    }

    private fun selectAgent(id: String) {
+        // The picked agent travels with the initial prompt (see submitCreate), so the dialog no
+        // longer writes default_agent to the global config here — doing so changed the mode for
+        // every other workspace and raced the new session's own model load.
        agent = id
-        sessions?.let { svc -> app.scope.launch { svc.updateConfig(directory, ConfigUpdateDto(agent = id)) } }
        val saved = app.models.value.model[id]?.let { "${it.providerID}/${it.modelID}" }
        if (saved != null && items.any { it.key == saved }) {
            prompt.model.select(saved)
@@ -253,10 +254,24 @@ internal class NewWorktreeDialog(
        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() }, text.trim())
+        onCreate(resolved, base.editor.item?.toString()?.trim()?.takeIf { it.isNotEmpty() }, pending(text))
        close(OK_EXIT_CODE)
    }

+    /** Bundles the typed prompt with the picked mode / model / reasoning, or null when empty. */
+    private fun pending(text: String): PendingPrompt? {
+        val body = text.trim()
+        if (body.isEmpty()) return null
+        val item = items.firstOrNull { it.key == modelKey }
+        return PendingPrompt(
+            text = body,
+            agent = agent,
+            provider = item?.provider,
+            model = item?.id,
+            variant = modelKey?.let { app.models.value.variant[it] },
+        )
+    }
+
    private fun enhance(text: String, done: (Result<String>) -> Unit) {
        val svc = sessions ?: return done(Result.failure(IllegalStateException("Session service unavailable")))
        scope.launch {
@@ -3,24 +3,37 @@ package ai.kilocode.client.agentManager.worktree
 import com.intellij.openapi.components.Service
 import com.intellij.util.concurrency.annotations.RequiresEdt

+/**
+ * An initial worktree prompt plus the mode / model / reasoning the New Worktree dialog picked for
+ * it. Carrying the selection alongside the text lets the first turn run with exactly what the user
+ * chose, instead of relying on the freshly-opened session resolving its own defaults.
+ */
+data class PendingPrompt(
+    val text: String,
+    val agent: String? = null,
+    val provider: String? = null,
+    val model: String? = null,
+    val variant: String? = null,
+)
+
 /**
 * One-shot handoff of an initial prompt from the New Worktree dialog to the freshly-opened worktree
- * session editor. The dialog creates the worktree, stashes the typed prompt keyed by worktree path,
- * and the editor consumes it once when it creates that worktree's first session — mirroring the VS
- * Code flow of create worktree → create session → send the initial prompt.
+ * session editor. The dialog creates the worktree, stashes the typed prompt (with its picked
+ * mode/model) keyed by worktree path, and the editor consumes it once when it creates that
+ * worktree's first session — mirroring the VS Code flow of create worktree → create session → send
+ * the initial prompt.
 */
@Service(Service.Level.APP)
 class PendingWorktreePrompt {
-    private val prompts = HashMap<String, String>()
+    private val prompts = HashMap<String, PendingPrompt>()

    @RequiresEdt
-    fun put(path: String, text: String) {
-        val body = text.trim()
-        if (body.isEmpty()) return
-        prompts[normalizeWorktreePath(path)] = body
+    fun put(path: String, prompt: PendingPrompt) {
+        if (prompt.text.isBlank()) return
+        prompts[normalizeWorktreePath(path)] = prompt.copy(text = prompt.text.trim())
    }

    /** Returns and clears the pending prompt for [path], or null when none is queued. */
    @RequiresEdt
-    fun take(path: String): String? = prompts.remove(normalizeWorktreePath(path))
+    fun take(path: String): PendingPrompt? = prompts.remove(normalizeWorktreePath(path))
 }
@@ -103,9 +103,9 @@ class WorktreeController(

    /**
     * Creates a worktree. When [prompt] is set, it is stashed for the worktree's first session so the
-     * editor auto-sends it once it opens (see [PendingWorktreePrompt]).
+     * editor auto-sends it once it opens with its picked mode/model (see [PendingWorktreePrompt]).
     */
-    fun create(branch: String, base: String?, existingBranch: Boolean = false, prompt: String? = null) {
+    fun create(branch: String, base: String?, existingBranch: Boolean = false, prompt: PendingPrompt? = null) {
        val id = "pending:$branch:${System.nanoTime()}"
        val temp = WorktreeDto(id, branch, branch, id)
        edt {
@@ -136,7 +136,7 @@ class WorktreeController(
    private fun finishCreate(
        temp: WorktreeDto,
        branch: String,
-        prompt: String?,
+        prompt: PendingPrompt?,
        result: CreateWorktreeResultDto,
    ) {
        val created = result.worktree
@@ -11,6 +11,7 @@ import ai.kilocode.client.session.SessionManager
 import ai.kilocode.client.session.SessionRef
 import ai.kilocode.client.session.SessionUi
 import ai.kilocode.client.session.SessionUiFactory
+import ai.kilocode.client.session.controller.PromptSelection
 import ai.kilocode.client.session.history.HistoryTime
 import ai.kilocode.client.session.history.LocalHistoryItem
 import ai.kilocode.client.util.UiTimerSource
@@ -117,8 +118,11 @@ open class WorktreeSessionEditorManager(
    /** Sends the New Worktree dialog's queued prompt into this worktree's first session, once. */
    @RequiresEdt
    private fun consumePendingPrompt() {
-        val text = service<PendingWorktreePrompt>().take(worktree.directory) ?: return
-        currentUi()?.submitPrompt(text)
+        val prompt = service<PendingWorktreePrompt>().take(worktree.directory) ?: return
+        currentUi()?.submitPrompt(
+            prompt.text,
+            PromptSelection(prompt.agent, prompt.provider, prompt.model, prompt.variant),
+        )
    }

    @RequiresEdt
@@ -45,6 +45,7 @@ import ai.kilocode.client.session.ui.selection.SessionSelection
 import ai.kilocode.client.session.ui.style.SessionEditorStyle
 import ai.kilocode.client.session.ui.style.SessionEditorStyleTarget
 import ai.kilocode.client.session.controller.EVENT_FLUSH_MS
+import ai.kilocode.client.session.controller.PromptSelection
 import ai.kilocode.client.session.controller.SessionController
 import ai.kilocode.client.session.controller.SessionControllerEvent
 import ai.kilocode.client.session.context.EditorContextGatherer
@@ -288,11 +289,16 @@ class SessionUi(
    /**
     * Sends [text] as the session's first message. Used by the New Worktree flow to auto-start a
     * session with the prompt typed in the dialog, routing through the same path as a typed prompt.
+     * The optional [select] carries the mode / model / reasoning picked in the dialog so the first
+     * turn runs with them even before this session's own model state has loaded.
     */
    @RequiresEdt
-    internal fun submitPrompt(text: String) {
+    internal fun submitPrompt(text: String, select: PromptSelection? = null) {
        if (text.isBlank()) return
-        sendPrompt(text, emptyList())
+        // Seed the session's agent/model/reasoning so the pickers and later turns reflect the pick,
+        // then send the first turn carrying it too (so it applies before workspace-ready resolves).
+        select?.let { controller.applySelection(it) }
+        sendPrompt(text, emptyList(), select)
    }

    @RequiresEdt
@@ -691,7 +697,7 @@ class SessionUi(
        }
    }

-    private fun sendPrompt(text: String, files: List<PromptPartDto>) {
+    private fun sendPrompt(text: String, files: List<PromptPartDto>, select: PromptSelection? = null) {
        if (text.isBlank() && files.isEmpty()) return
        prompt.clear()
        val follow = scroll.following()
@@ -720,7 +726,7 @@ class SessionUi(
            val model = controller.model.model ?: "none"
            "${ChatLogSummary.prompt(PromptDto(parts = parts, editorContext = editor.context))} agent=$agent model=$model ready=${controller.ready}"
        }
-        controller.prompt(text, allFiles, editor.context)
+        controller.prompt(text, allFiles, editor.context, select)
        scroll.followBottom(follow)
    }

@@ -264,11 +264,16 @@ class SessionController(
        }
    }

-    fun prompt(text: String, files: List<PromptPartDto> = emptyList(), editorContext: EditorContextDto? = null) {
+    fun prompt(
+        text: String,
+        files: List<PromptPartDto> = emptyList(),
+        editorContext: EditorContextDto? = null,
+        select: PromptSelection? = null,
+    ) {
        assertEdt()
        val start = sid ?: ref?.key ?: "pending"
        val exists = sid != null
-        val dto = promptDto(text, files, editorContext)
+        val dto = promptDto(text, files, editorContext, select)
        val props = promptProps(files)
        LOG.debug { "${ChatLogSummary.sid(start)} ${ChatLogSummary.prompt(dto)} ${ChatLogSummary.dir(directory)}" }
        dispatch(Dispatch("prompt", "user", text, props, start, exists)) { id ->
@@ -683,6 +688,32 @@ class SessionController(
        capture("Reasoning Variant Selected", sessionProps() + mapOf("model" to key, "variant" to value))
    }

+    /**
+     * Seeds this session's agent / model / reasoning from an initial [select] (New Worktree flow),
+     * mirroring VS Code's setSessionAgent + setSessionModel + variant seeding. Attaching the pick to
+     * the first prompt alone only affects that one turn; setting it as the session's preferred
+     * selection makes the pickers and every later turn use it too, and survives the later
+     * workspace-ready model resolution because [prefAgent] / [prefModel] win in [syncModelSelection].
+     */
+    fun applySelection(select: PromptSelection) {
+        assertEdt()
+        val agent = select.agent ?: return
+        fire(SessionControllerEvent.WorkspaceReady) {
+            model.agent = agent
+            val provider = select.provider
+            val id = select.model
+            if (provider != null && id != null) {
+                val key = "$provider/$id"
+                app.selectModel(agent, provider, id)
+                select.variant?.let { app.selectVariant(key, it) }
+                prefAgent = agent
+                prefModel = key
+            }
+            syncModelSelection()
+            model.refreshHeader()
+        }
+    }
+
    // ------ permission / question resolution ------

    fun replyPermission(requestId: String, reply: PermissionReplyDto, rules: PermissionAlwaysRulesDto? = null) {
@@ -1827,19 +1858,24 @@ class SessionController(
        text: String,
        files: List<PromptPartDto> = emptyList(),
        editorContext: EditorContextDto? = null,
+        select: PromptSelection? = null,
    ): PromptDto {
-        val full = model.model
-        val sel = full?.let(::parseModel)
-        val variant = model.variant?.takeIf { it in model.variants }
+        val sel = model.model?.let(::parseModel)
+        val provider = select?.provider ?: sel?.first
+        val modelId = select?.model ?: sel?.second
+        val agent = select?.agent ?: model.agent
+        // An explicit variant comes from the dialog before the model catalog is loaded, so it can't
+        // be validated against model.variants yet; only the fallback is filtered.
+        val variant = select?.variant ?: model.variant?.takeIf { it in model.variants }
        val parts = buildList {
            text.takeIf { it.isNotBlank() }?.let { add(PromptPartDto(type = "text", text = it)) }
            addAll(files)
        }
        return PromptDto(
            parts = parts,
-            providerID = sel?.first,
-            modelID = sel?.second,
-            agent = model.agent,
+            providerID = provider,
+            modelID = modelId,
+            agent = agent,
            variant = variant,
            editorContext = editorContext,
        )
@@ -2529,6 +2565,18 @@ private fun selection(value: String): ModelSelectionDto? {
    return ModelSelectionDto(parsed.first, parsed.second)
 }

+/**
+ * An explicit agent / provider / model / reasoning selection to attach to a single prompt. Used by
+ * the New Worktree flow so the first turn runs with the mode and model picked in the dialog rather
+ * than whatever the freshly-opened session resolves as its default.
+ */
+data class PromptSelection(
+    val agent: String? = null,
+    val provider: String? = null,
+    val model: String? = null,
+    val variant: String? = null,
+)
+
 private fun parseModel(value: String): Pair<String, String>? {
    val slash = value.indexOf('/')
    if (slash <= 0 || slash >= value.length - 1) return null
@@ -3,6 +3,7 @@ package ai.kilocode.client.agentManager
 import ai.kilocode.client.agentManager.worktree.WorktreeIcons
 import ai.kilocode.client.agentManager.worktree.KiloWorktreeService
 import ai.kilocode.client.agentManager.worktree.WorktreeController
+import ai.kilocode.client.agentManager.worktree.PendingPrompt
 import ai.kilocode.client.agentManager.worktree.PendingWorktreePrompt
 import ai.kilocode.client.agentManager.worktree.WorktreeNameCache
 import ai.kilocode.client.agentManager.worktree.WorktreeNames
@@ -298,15 +299,22 @@ class WorktreeControllerTest : BasePlatformTestCase() {
        assertEquals("/wt/pr-7", controller.model.getElementAt(0).path)
    }

-    fun `test create stashes the prompt for the created worktree`() {
+    fun `test create stashes the prompt with its picked selection for the created worktree`() {
        val controller = controller()

-        ApplicationManager.getApplication().invokeAndWait { controller.create("feature/y", null, prompt = "fix the bug") }
+        ApplicationManager.getApplication().invokeAndWait {
+            controller.create(
+                "feature/y",
+                null,
+                prompt = PendingPrompt("fix the bug", agent = "plan", provider = "kilo", model = "gpt-5", variant = "high"),
+            )
+        }
        flush()

        val created = controller.model.getElementAt(0)
        ApplicationManager.getApplication().invokeAndWait {
-            assertEquals("fix the bug", service<PendingWorktreePrompt>().take(created.path))
+            val stashed = service<PendingWorktreePrompt>().take(created.path)
+            assertEquals(PendingPrompt("fix the bug", "plan", "kilo", "gpt-5", "high"), stashed)
            // A one-shot take clears it.
            assertNull(service<PendingWorktreePrompt>().take(created.path))
        }
@@ -37,7 +37,7 @@ class NewWorktreeDialogTest : BasePlatformTestCase() {
    private lateinit var sessionRpc: FakeSessionRpcApi
    private lateinit var sessions: KiloSessionService
    private var dialog: NewWorktreeDialog? = null
-    private val created = mutableListOf<Triple<String, String?, String>>()
+    private val created = mutableListOf<Triple<String, String?, PendingPrompt?>>()

    override fun setUp() {
        super.setUp()
@@ -71,14 +71,22 @@ class NewWorktreeDialogTest : BasePlatformTestCase() {
        }
    }

-    fun `test selecting a mode persists the agent for the workspace`() {
+    fun `test selecting a mode forwards it with the created prompt and writes no global config`() {
        open()
-        flushUntil { edt { mode().selectedForTest() != null } }
+        flushUntil { edt { model().selectionKeyForTest() != null } }

-        edt { mode().onSelect(ModePicker.Item("plan", "Plan")) }
-        flushUntil { sessionRpc.configs.any { it.second.agent == "plan" } }
+        edt {
+            mode().onSelect(ModePicker.Item("plan", "Plan"))
+            prompt().setText("do it")
+        }
+        flushUntil { edt { prompt().isSendEnabled } }
+        edt { prompt().send() }
+        flushUntil { created.isNotEmpty() }
+        dialog = null

-        assertTrue(sessionRpc.configs.any { it.first == "/test" && it.second.agent == "plan" })
+        assertEquals("plan", created.single().third?.agent)
+        // Picking a mode must no longer mutate the global default_agent config.
+        assertTrue(sessionRpc.configs.none { it.second.agent != null })
    }

    fun `test selecting a model persists it for the default agent`() {
@@ -99,8 +107,9 @@ class NewWorktreeDialogTest : BasePlatformTestCase() {
        assertEquals("high", app.models.value.variant["kilo/gpt-5"])
    }

-    fun `test creating forwards the prompt and resolved branch`() {
+    fun `test creating forwards the prompt, resolved branch, and default selection`() {
        open()
+        flushUntil { edt { model().selectionKeyForTest() != null } }
        edt {
            prompt().setText("build the thing")
        }
@@ -109,7 +118,14 @@ class NewWorktreeDialogTest : BasePlatformTestCase() {
        flushUntil { created.isNotEmpty() }
        dialog = null

-        assertEquals(listOf(Triple("agent/foo", "main", "build the thing")), created)
+        val entry = created.single()
+        assertEquals("agent/foo", entry.first)
+        assertEquals("main", entry.second)
+        val payload = requireNotNull(entry.third)
+        assertEquals("build the thing", payload.text)
+        assertEquals("build", payload.agent)
+        assertEquals("kilo", payload.provider)
+        assertEquals("gpt-5", payload.model)
    }

    private fun open() {
@@ -90,10 +90,15 @@ class WorktreeSessionEditorManagerTest : BasePlatformTestCase() {
        assertEquals(1, requested.size)
    }

-    fun `test new session sends the queued worktree prompt once`() {
+    fun `test new session sends the queued worktree prompt once with its picked selection`() {
        rpc.session = session("ses_new", updated = 4.0).copy(title = "New session")
        val manager = manager()
-        edt { service<PendingWorktreePrompt>().put(DIR, "fix the bug") }
+        edt {
+            service<PendingWorktreePrompt>().put(
+                DIR,
+                PendingPrompt("fix the bug", agent = "plan", provider = "kilo", model = "gpt-5", variant = "high"),
+            )
+        }

        edt { manager.newSession() }
        waitUntil { rpc.prompts.any { it.first == "ses_new" } }
@@ -101,6 +106,11 @@ class WorktreeSessionEditorManagerTest : BasePlatformTestCase() {
        val sent = rpc.prompts.single { it.first == "ses_new" }
        assertEquals(DIR, sent.second)
        assertTrue("prompt parts should carry the typed text", sent.third.parts.any { it.text == "fix the bug" })
+        // The dialog's mode / model / reasoning must ride along with the first turn.
+        assertEquals("plan", sent.third.agent)
+        assertEquals("kilo", sent.third.providerID)
+        assertEquals("gpt-5", sent.third.modelID)
+        assertEquals("high", sent.third.variant)
        // The queued prompt is consumed once, so it is cleared after the first new session.
        assertNull(edt { service<PendingWorktreePrompt>().take(DIR) })
    }
@@ -0,0 +1,60 @@
+package ai.kilocode.client.session.controller
+
+import ai.kilocode.rpc.dto.AgentDto
+import ai.kilocode.rpc.dto.ConfigDto
+import ai.kilocode.rpc.dto.KiloAppStateDto
+import ai.kilocode.rpc.dto.KiloAppStatusDto
+import ai.kilocode.rpc.dto.ModelDto
+import ai.kilocode.rpc.dto.ProviderDto
+
+/**
+ * Covers [SessionController.applySelection], the New Worktree entry point that seeds a fresh
+ * session's mode / model / reasoning so the pickers and every later turn use the dialog's pick —
+ * not just the single first prompt.
+ */
+class SessionControllerSelectionTest : SessionControllerTestBase() {
+
+    fun `test applySelection seeds agent, model and variant and rides the first prompt`() {
+        appRpc.state.value = KiloAppStateDto(KiloAppStatusDto.READY, config = ConfigDto(model = "kilo/gpt-5"))
+        projectRpc.state.value = workspaceReady(
+            agents = listOf(AgentDto("code", "Code", mode = "code"), AgentDto("plan", "Plan", mode = "plan")),
+            default = "code",
+            providers = listOf(
+                ProviderDto(
+                    id = "kilo",
+                    name = "Kilo",
+                    models = mapOf(
+                        "gpt-5" to ModelDto(id = "gpt-5", name = "GPT-5"),
+                        "opus" to ModelDto(id = "opus", name = "Opus", variants = listOf("low", "high")),
+                    ),
+                ),
+            ),
+        )
+        val m = controller()
+        flush()
+
+        edt { m.applySelection(PromptSelection("plan", "kilo", "opus", "high")) }
+        flush()
+
+        var agent: String? = null
+        var model: String? = null
+        var variant: String? = null
+        edt {
+            agent = m.model.agent
+            model = m.model.model
+            variant = m.model.variant
+        }
+        assertEquals("plan", agent)
+        assertEquals("kilo/opus", model)
+        assertEquals("high", variant)
+
+        edt { m.prompt("go") }
+        flush()
+
+        val sent = rpc.prompts.single()
+        assertEquals("plan", sent.third.agent)
+        assertEquals("kilo", sent.third.providerID)
+        assertEquals("opus", sent.third.modelID)
+        assertEquals("high", sent.third.variant)
+    }
+}