From 9e09eafb8e72dbf765941db8e3cc2b9242c9490a Mon Sep 17 00:00:00 2001 From: kirillk Date: Tue, 11 Aug 2026 17:12:54 -0400 Subject: [PATCH] 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 --- .../client/agentManager/AgentManagerPanel.kt | 2 +- .../worktree/NewWorktreeDialog.kt | 23 +++++-- .../worktree/PendingWorktreePrompt.kt | 31 ++++++--- .../worktree/WorktreeController.kt | 6 +- .../worktree/WorktreeSessionEditorManager.kt | 8 ++- .../ai/kilocode/client/session/SessionUi.kt | 14 ++-- .../session/controller/SessionController.kt | 64 ++++++++++++++++--- .../agentManager/WorktreeControllerTest.kt | 14 +++- .../worktree/NewWorktreeDialogTest.kt | 32 +++++++--- .../WorktreeSessionEditorManagerTest.kt | 14 +++- .../SessionControllerSelectionTest.kt | 60 +++++++++++++++++ 11 files changed, 224 insertions(+), 44 deletions(-) create mode 100644 packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/controller/SessionControllerSelectionTest.kt diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/agentManager/AgentManagerPanel.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/agentManager/AgentManagerPanel.kt index 5dba3ce42b7..2d17c52964a 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/agentManager/AgentManagerPanel.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/agentManager/AgentManagerPanel.kt @@ -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) }, diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/agentManager/worktree/NewWorktreeDialog.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/agentManager/worktree/NewWorktreeDialog.kt index a46498d4cb1..9a7652abc52 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/agentManager/worktree/NewWorktreeDialog.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/agentManager/worktree/NewWorktreeDialog.kt @@ -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, - 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) -> Unit) { val svc = sessions ?: return done(Result.failure(IllegalStateException("Session service unavailable"))) scope.launch { diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/agentManager/worktree/PendingWorktreePrompt.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/agentManager/worktree/PendingWorktreePrompt.kt index e0337c67cbb..1d7e680eec1 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/agentManager/worktree/PendingWorktreePrompt.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/agentManager/worktree/PendingWorktreePrompt.kt @@ -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() + private val prompts = HashMap() @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)) } diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/agentManager/worktree/WorktreeController.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/agentManager/worktree/WorktreeController.kt index 8da4f523408..f50e94d08c6 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/agentManager/worktree/WorktreeController.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/agentManager/worktree/WorktreeController.kt @@ -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 diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/agentManager/worktree/WorktreeSessionEditorManager.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/agentManager/worktree/WorktreeSessionEditorManager.kt index 2d98bb7315b..c64ca5c0f04 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/agentManager/worktree/WorktreeSessionEditorManager.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/agentManager/worktree/WorktreeSessionEditorManager.kt @@ -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().take(worktree.directory) ?: return - currentUi()?.submitPrompt(text) + val prompt = service().take(worktree.directory) ?: return + currentUi()?.submitPrompt( + prompt.text, + PromptSelection(prompt.agent, prompt.provider, prompt.model, prompt.variant), + ) } @RequiresEdt diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/SessionUi.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/SessionUi.kt index afc90a6dc12..5e5cf5db5b8 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/SessionUi.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/SessionUi.kt @@ -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) { + private fun sendPrompt(text: String, files: List, 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) } diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/controller/SessionController.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/controller/SessionController.kt index 27e3825512c..a468e102d45 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/controller/SessionController.kt +++ b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/controller/SessionController.kt @@ -264,11 +264,16 @@ class SessionController( } } - fun prompt(text: String, files: List = emptyList(), editorContext: EditorContextDto? = null) { + fun prompt( + text: String, + files: List = 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 = 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? { val slash = value.indexOf('/') if (slash <= 0 || slash >= value.length - 1) return null diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/agentManager/WorktreeControllerTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/agentManager/WorktreeControllerTest.kt index 14cded655ef..a5d2e55f1c9 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/agentManager/WorktreeControllerTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/agentManager/WorktreeControllerTest.kt @@ -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().take(created.path)) + val stashed = service().take(created.path) + assertEquals(PendingPrompt("fix the bug", "plan", "kilo", "gpt-5", "high"), stashed) // A one-shot take clears it. assertNull(service().take(created.path)) } diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/agentManager/worktree/NewWorktreeDialogTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/agentManager/worktree/NewWorktreeDialogTest.kt index 42d0f3735c9..2cc17b08191 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/agentManager/worktree/NewWorktreeDialogTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/agentManager/worktree/NewWorktreeDialogTest.kt @@ -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>() + private val created = mutableListOf>() 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() { diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/agentManager/worktree/WorktreeSessionEditorManagerTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/agentManager/worktree/WorktreeSessionEditorManagerTest.kt index ae2c639ae11..c88c3286944 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/agentManager/worktree/WorktreeSessionEditorManagerTest.kt +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/agentManager/worktree/WorktreeSessionEditorManagerTest.kt @@ -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().put(DIR, "fix the bug") } + edt { + service().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().take(DIR) }) } diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/controller/SessionControllerSelectionTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/controller/SessionControllerSelectionTest.kt new file mode 100644 index 00000000000..6b786ea1c2a --- /dev/null +++ b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/controller/SessionControllerSelectionTest.kt @@ -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) + } +}