Merge pull request #9245 from Kilo-Org/fix/plan-followup-continue-button

fix(vscode): Fix plan continue-here button

Author: alex-alecu

.changeset/plan-followup-continue-click.md

@@ -0,0 +1,5 @@
+---
+"kilo-code": patch
+---
+
+Fix the "Continue here" button not submitting after a plan is finished. Picking an option on a single-question prompt now sends the reply immediately — matching the CLI behaviour — and the redundant "Type your own answer" row no longer appears on the plan follow-up question.

packages/kilo-vscode/tests/unit/question-dock-utils.test.ts

@@ -1,5 +1,6 @@
 import { describe, it, expect } from "bun:test"
 import {
+  pickOutcome,
   resolveOptimisticQuestionAgent,
   resolveQuestionMode,
   resolveSelectedQuestionMode,
@@ -128,3 +129,25 @@ describe("resolveOptimisticQuestionAgent", () => {
     expect(result).toEqual({ base: "ask", agent: "architect" })
   })
 })
+
+describe("pickOutcome", () => {
+  it("submits immediately on a single-question single-select option pick", () => {
+    expect(pickOutcome({ single: true, multi: false, custom: false })).toEqual({ kind: "submit" })
+  })
+
+  it("advances to the next tab on a multi-question single-select option pick", () => {
+    expect(pickOutcome({ single: false, multi: false, custom: false })).toEqual({ kind: "advance" })
+  })
+
+  it("stays on the current tab for a multi-select pick", () => {
+    expect(pickOutcome({ single: true, multi: true, custom: false })).toEqual({ kind: "stay" })
+  })
+
+  it("defers submission for a single-select custom-input pick (handleCustomSubmit owns the submit)", () => {
+    expect(pickOutcome({ single: true, multi: false, custom: true })).toEqual({ kind: "stay" })
+  })
+
+  it("stays on the current tab for a multi-select custom-input pick", () => {
+    expect(pickOutcome({ single: false, multi: true, custom: true })).toEqual({ kind: "stay" })
+  })
+})

packages/kilo-vscode/webview-ui/src/components/chat/QuestionDock.tsx

@@ -12,7 +12,12 @@ import { Icon } from "@kilocode/kilo-ui/icon"
 import { useSession } from "../../context/session"
 import { useLanguage } from "../../context/language"
 import type { QuestionRequest } from "../../types/messages"
-import { resolveOptimisticQuestionAgent, resolveSelectedQuestionMode, toggleAnswer } from "./question-dock-utils"
+import {
+  pickOutcome,
+  resolveOptimisticQuestionAgent,
+  resolveSelectedQuestionMode,
+  toggleAnswer,
+} from "./question-dock-utils"
 
 export const QuestionDock: Component<{ request: QuestionRequest }> = (props) => {
   const session = useSession()
@@ -119,7 +124,14 @@ export const QuestionDock: Component<{ request: QuestionRequest }> = (props) =>
 
     syncAgent(answers, kinds)
 
-    if (!single() && !multi()) {
+    const outcome = pickOutcome({ single: single(), multi: multi(), custom })
+    if (outcome.kind === "submit") {
+      // Mirror TUI behaviour: a single-question single-select option pick submits immediately.
+      // handleCustomSubmit covers the custom-input path via its own submit() call.
+      reply([[answer]])
+      return
+    }
+    if (outcome.kind === "advance") {
       setStore("tab", store.tab + 1)
     }
   }

packages/kilo-vscode/webview-ui/src/components/chat/question-dock-utils.ts

@@ -1,5 +1,22 @@
 import type { QuestionOption } from "../../types/messages"
 
+export type PickOutcome = { kind: "submit" } | { kind: "advance" } | { kind: "stay" }
+
+/**
+ * Decide what should happen after a user picks an option in the question dock.
+ *
+ * - Multi-select prompts: the pick only toggles local state; no tab change, no submit.
+ * - Single-question single-select, option pick: submit immediately (matches the TUI).
+ * - Multi-question single-select, option pick: advance to the next tab.
+ * - Custom-input path for a single-select is handled separately in handleCustomSubmit.
+ */
+export function pickOutcome(input: { single: boolean; multi: boolean; custom: boolean }): PickOutcome {
+  if (input.multi) return { kind: "stay" }
+  if (input.single && input.custom) return { kind: "stay" }
+  if (input.single) return { kind: "submit" }
+  return { kind: "advance" }
+}
+
 export function toggleAnswer(existing: string[], answer: string): string[] {
   const next = [...existing]
   const index = next.indexOf(answer)

packages/opencode/src/kilocode/plan-followup.ts

@@ -16,6 +16,7 @@ import { MessageV2 } from "@/session/message-v2"
 import { Todo } from "@/session/todo"
 import { makeRuntime } from "@/effect/run-service"
 import { Log } from "@/util/log"
+import { KiloSessionPromptQueue } from "@/kilocode/session/prompt-queue"
 import path from "path"
 import z from "zod"
 
@@ -255,6 +256,7 @@ export namespace PlanFollowup {
       text: input.text,
       synthetic: input.synthetic ?? true,
     } satisfies MessageV2.TextPart)
+    return msg
   }
 
   function prompt(input: { sessionID: SessionID; abort: AbortSignal }) {
@@ -264,7 +266,12 @@ export namespace PlanFollowup {
         {
           question: "Ready to implement?",
           header: "Implement",
-          custom: true,
+          // On CLI the main prompt input is hidden while a blocking question is active,
+          // so we need the custom-answer row to allow a free-text reply. On VS Code the
+          // main prompt input below the dock already routes typed text as a question
+          // reply, so "Type your own answer" would be redundant (originally hidden in
+          // 65566af7f8, flipped back during the v1.4.4 upstream merge).
+          custom: Flag.KILO_CLIENT === "cli",
           options: [
             {
               label: ANSWER_NEW_SESSION,
@@ -398,22 +405,24 @@ export namespace PlanFollowup {
       const code = await resolveCodeModel({
         model: user.model,
       })
-      await inject({
+      const msg = await inject({
         sessionID: input.sessionID,
         agent: "code",
         model: code.model,
         text: "Implement the plan above.",
       })
+      KiloSessionPromptQueue.retarget(input.sessionID, msg.id)
       return "continue"
     }
 
     Telemetry.trackPlanFollowup(input.sessionID, "custom")
-    await inject({
+    const msg = await inject({
       sessionID: input.sessionID,
       agent: "plan",
       model: user.model,
       text: answer,
     })
+    KiloSessionPromptQueue.retarget(input.sessionID, msg.id)
     return "continue"
   }
 }

packages/opencode/src/kilocode/session/prompt-queue.ts

@@ -9,10 +9,15 @@ type Slot = {
   readonly tail: Promise<void>
 }
 
+type Target = {
+  readonly base: MessageID
+  readonly extras: ReadonlySet<MessageID>
+}
+
 export namespace KiloSessionPromptQueue {
   const tails = new Map<SessionID, Promise<void>>()
   const versions = new Map<SessionID, number>()
-  const targets = new Map<SessionID, MessageID>()
+  const targets = new Map<SessionID, Target>()
 
   const version = (sessionID: SessionID) => versions.get(sessionID) ?? 0
   const settle = (promise: Promise<void>) =>
@@ -27,15 +32,30 @@ export namespace KiloSessionPromptQueue {
     })
   }
 
+  /**
+   * Exempt an injected user message from being hidden by scope().
+   * Called after PlanFollowup.inject() so the injected follow-up is visible
+   * without also unhiding unrelated prompts that were queued mid-turn.
+   */
+  export function retarget(sessionID: SessionID, id: MessageID) {
+    const current = targets.get(sessionID)
+    if (!current) return
+    const extras = new Set(current.extras)
+    extras.add(id)
+    targets.set(sessionID, { base: current.base, extras })
+  }
+
   export function scope(sessionID: SessionID, messages: MessageV2.WithParts[]) {
     const target = targets.get(sessionID)
     if (!target) return messages
 
     const hidden = new Set(
-      messages.filter((item) => item.info.role === "user" && item.info.id > target).map((item) => item.info.id),
+      messages
+        .filter((item) => item.info.role === "user" && item.info.id > target.base && !target.extras.has(item.info.id))
+        .map((item) => item.info.id),
     )
     const visible = messages.filter((item) => {
-      if (item.info.role === "user") return item.info.id <= target
+      if (item.info.role === "user") return !hidden.has(item.info.id)
       if (item.info.role === "assistant") return !hidden.has(item.info.parentID)
       return true
     })
@@ -45,12 +65,14 @@ export namespace KiloSessionPromptQueue {
     // was written after the queue event). Ordering by time_created alone puts
     // the queued prompt before the prior turn's final assistant reply, which
     // makes the next request end with an assistant message and trips Anthropic's
-    // prefill rejection. Move the target user message and any of its own turn's
-    // assistant messages to the end so the request always ends with the queued
-    // user prompt (or with its own turn's latest assistant step).
+    // prefill rejection. Move the target user message (and any injected
+    // follow-ups) plus their own turn's assistant messages to the end so the
+    // request always ends with the queued user prompt (or its latest assistant
+    // step).
+    const ownsID = (id: MessageID) => id === target.base || target.extras.has(id)
     const owns = (item: MessageV2.WithParts) => {
-      if (item.info.role === "user") return item.info.id === target
-      if (item.info.role === "assistant") return item.info.parentID === target
+      if (item.info.role === "user") return ownsID(item.info.id)
+      if (item.info.role === "assistant") return ownsID(item.info.parentID)
       return false
     }
     const before: MessageV2.WithParts[] = []
@@ -81,12 +103,12 @@ export namespace KiloSessionPromptQueue {
             if (slot.version !== version(sessionID)) return cancelled
             return Effect.acquireUseRelease(
               Effect.sync(() => {
-                targets.set(sessionID, target)
+                targets.set(sessionID, { base: target, extras: new Set() })
               }),
               () => work,
               () =>
                 Effect.sync(() => {
-                  if (targets.get(sessionID) === target) targets.delete(sessionID)
+                  if (targets.get(sessionID)?.base === target) targets.delete(sessionID)
                 }),
             )
           }),

packages/opencode/test/kilocode/plan-followup.test.ts

@@ -266,6 +266,65 @@ describe("plan follow-up", () => {
       await expect(pending).resolves.toBe("break")
     }))
 
+  test("ask - emits a single-select question with the canonical answers and custom enabled on CLI", () =>
+    withInstance(async () => {
+      const seeded = await seed({ text: "1. Build" })
+      const pending = PlanFollowup.ask({
+        sessionID: seeded.sessionID,
+        messages: seeded.messages,
+        abort: AbortSignal.any([]),
+      })
+
+      const item = await waitQuestion(seeded.sessionID)
+      expect(item).toBeDefined()
+      if (!item) return
+      const q = item.questions[0]
+      expect(q).toBeDefined()
+      if (!q) return
+
+      // On CLI the main prompt input is hidden while a blocking question is active, so
+      // "Type your own answer" must remain available — i.e. custom must not be false.
+      expect(q.custom).not.toBe(false)
+      expect(q.multiple).not.toBe(true)
+      expect(q.options.map((item) => item.label)).toEqual([
+        PlanFollowup.ANSWER_NEW_SESSION,
+        PlanFollowup.ANSWER_CONTINUE,
+      ])
+
+      await question.reject(item.id)
+      await expect(pending).resolves.toBe("break")
+    }))
+
+  test("ask - hides custom answer row on VS Code where the main prompt input handles typed replies", () =>
+    withInstance(async () => {
+      const prev = process.env.KILO_CLIENT
+      try {
+        process.env.KILO_CLIENT = "vscode"
+        const seeded = await seed({ text: "1. Build" })
+        const pending = PlanFollowup.ask({
+          sessionID: seeded.sessionID,
+          messages: seeded.messages,
+          abort: AbortSignal.any([]),
+        })
+
+        const item = await waitQuestion(seeded.sessionID)
+        expect(item).toBeDefined()
+        if (!item) return
+        const q = item.questions[0]
+        expect(q).toBeDefined()
+        if (!q) return
+
+        // On VS Code the dock's main prompt input already accepts free text as a reply,
+        // so the "Type your own answer" row is redundant and must be hidden.
+        expect(q.custom).toBe(false)
+
+        await question.reject(item.id)
+        await expect(pending).resolves.toBe("break")
+      } finally {
+        process.env.KILO_CLIENT = prev
+      }
+    }))
+
   test("ask - returns continue and creates code message on Continue here", () =>
     withInstance(async () => {
       const get = spyOn(PlanFollowupRuntime, "agent").mockImplementation(async (name: string) => {
@@ -349,6 +408,40 @@ describe("plan follow-up", () => {
       expect(part.synthetic).toBe(true)
     }))
 
+  test("ask - retargets prompt queue so injected message is visible in scope", () =>
+    withInstance(async () => {
+      const { KiloSessionPromptQueue } = await import("../../src/kilocode/session/prompt-queue")
+      const seeded = await seed({ text: "1. Refactor\n2. Ship" })
+
+      // Simulate the prompt queue having a target set (like during a running loop)
+      const original = seeded.messages.find((m) => m.info.role === "user")!.info.id
+      KiloSessionPromptQueue.retarget(seeded.sessionID, original)
+
+      const pending = PlanFollowup.ask({
+        sessionID: seeded.sessionID,
+        messages: seeded.messages,
+        abort: AbortSignal.any([]),
+      })
+
+      const item = await waitQuestion(seeded.sessionID)
+      expect(item).toBeDefined()
+      if (!item) return
+      await question.reply({
+        requestID: item.id,
+        answers: [[PlanFollowup.ANSWER_CONTINUE]],
+      })
+
+      await expect(pending).resolves.toBe("continue")
+
+      // The injected user message must be visible when scoped
+      const all = await Session.messages({ sessionID: seeded.sessionID })
+      const scoped = KiloSessionPromptQueue.scope(seeded.sessionID, all)
+      const injected = scoped.findLast((m) => m.info.role === "user")
+      expect(injected).toBeDefined()
+      const part = injected!.parts.find((p) => p.type === "text")
+      expect(part?.type === "text" && part.text).toBe("Implement the plan above.")
+    }))
+
   test("ask - creates a new session on Start new session with handover and todos", () =>
     withInstance(async () => {
       const get = spyOn(PlanFollowupRuntime, "agent").mockImplementation(async (name: string) => {

packages/opencode/test/kilocode/session-prompt-queue.test.ts

@@ -188,6 +188,39 @@ describe("session prompt queue", () => {
     expect(ids).toEqual([m1, a1, a1tail, m2, a2step1])
   })
 
+  test("retarget keeps older queued prompts hidden", async () => {
+    // Regression: retargeting used to move the visible-message boundary forward,
+    // which unhid any user prompts queued between the base and the injected
+    // follow-up. Exempt the follow-up without reopening the boundary.
+    const sessionID = SessionID.make("session_retarget_hide")
+    const base = MessageID.make("message_b1")
+    const ans = MessageID.make("message_b2")
+    const queued = MessageID.make("message_b3") // queued while base was running
+    const injected = MessageID.make("message_b4") // injected follow-up
+    const messages = [
+      user(sessionID, base),
+      assistant(sessionID, ans, base),
+      user(sessionID, queued),
+      user(sessionID, injected),
+    ]
+
+    const ids = await Effect.runPromise(
+      KiloSessionPromptQueue.enqueue(
+        sessionID,
+        base,
+        Effect.sync(() => {
+          KiloSessionPromptQueue.retarget(sessionID, injected)
+          return KiloSessionPromptQueue.scope(sessionID, messages).map((item) => item.info.id)
+        }),
+        Effect.succeed([]),
+      ),
+    )
+
+    expect(ids).not.toContain(queued)
+    expect(ids).toContain(injected)
+    expect(ids[ids.length - 1]).toBe(injected)
+  })
+
   test("continues a queued prompt after the active run finishes", async () => {
     const ready = Promise.withResolvers<void>()
     const release = Promise.withResolvers<void>()